From 0fe34c47cf9e58f733b497110745d7ff2cfc1b60 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 10 May 2026 10:42:33 +0200 Subject: [PATCH] fix ipadapters and improve offloading Co-authored-by: Copilot Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 +++ TODO.md | 6 ++---- modules/gr_hijack.py | 9 +++++++++ modules/ipadapter.py | 11 ++++++++--- modules/sd_offload.py | 23 +++++++++++++++-------- 5 files changed, 37 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index af3e2179e..610081f01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,6 +66,7 @@ For full details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/m - remove "processed preview" from ui preprocessor output can still be generated by clicking preview button in in control unit and it will render into normal output area - **Internal** + - `offload` auto-reapply hook on error - refactor `pip` installer, thanks @awsr - remove obsolete `lora` stepwise and functional code, thanks @awsr - interrupt model loading between components @@ -90,6 +91,8 @@ For full details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/m - `detailer` handle `lora` internally - vae preview flashes previous image - `torch.compile` improvements + - `gradio` preprocess exception handling + - `ipadapters` with offloading ## Update for 2026-04-28 diff --git a/TODO.md b/TODO.md index b35980fce..e607bb7cb 100644 --- a/TODO.md +++ b/TODO.md @@ -1,16 +1,14 @@ # TODO -- Outpaint -- IPAdapter -- https://github.com/PrunaAI/pruna - ## Features ### Assigned +- Check Outpaint, @vladmandic - Chat-based interface, @vladmandic - Control tab verify overrides handling, @vladmandic - Reimplement `llama` remover for Kanvas, @vladmandic +- Implement [pruna](https://github.com/PrunaAI/pruna), @vladmandic - Cloud providers, @CalamitousFelicitousness - Video processing add full API support, @CalamitousFelicitousness diff --git a/modules/gr_hijack.py b/modules/gr_hijack.py index 0fe7eb24e..082ace269 100644 --- a/modules/gr_hijack.py +++ b/modules/gr_hijack.py @@ -153,6 +153,7 @@ def patch_gradio(): orig_cancel_tasks = gradio.utils.cancel_tasks orig_restore_session_state = gradio.route_utils.restore_session_state orig_call_prediction = gradio.queueing.Queue.call_prediction + orig_blocks_preprocess_data = gradio.blocks.Blocks.preprocess_data async def wrap_cancel_tasks(task_ids: set[str]): log.error(f'Gradio cancel: task={task_ids}') @@ -200,9 +201,17 @@ def patch_gradio(): log.error(f"Gradio queue: events={len(events)} batch={batch} error: {e}") raise + def wrap_blocks_preprocess_data(self, fn_index: int, inputs: list, state: dict): + try: + return orig_blocks_preprocess_data(self, fn_index, inputs, state) + except Exception as e: + log.error(f"Gradio preprocess: {e}") + raise + gradio.queueing.Queue.call_prediction = wrap_call_prediction gradio.route_utils.restore_session_state = wrap_restore_session_state gradio.utils.cancel_tasks = wrap_cancel_tasks + gradio.blocks.Blocks.preprocess_data = wrap_blocks_preprocess_data def patch_gradio_future(): diff --git a/modules/ipadapter.py b/modules/ipadapter.py index 291ae896a..3b9d7147e 100644 --- a/modules/ipadapter.py +++ b/modules/ipadapter.py @@ -201,10 +201,10 @@ def load_image_encoder(pipe: DiffusionPipeline, adapter_names: list[str]): else: if clip_subfolder is None: image_encoder = transformers.CLIPVisionModelWithProjection.from_pretrained(clip_repo, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, use_safetensors=True, **offline_config) - log.debug(f'IP adapter load: encoder="{clip_repo}" cls={pipe.image_encoder.__class__.__name__}') + log.debug(f'IP adapter load: encoder="{clip_repo}" cls={image_encoder.__class__.__name__}') else: image_encoder = transformers.CLIPVisionModelWithProjection.from_pretrained(clip_repo, subfolder=clip_subfolder, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, use_safetensors=True, **offline_config) - log.debug(f'IP adapter load: encoder="{clip_repo}/{clip_subfolder}" cls={pipe.image_encoder.__class__.__name__}') + log.debug(f'IP adapter load: encoder="{clip_repo}/{clip_subfolder}" cls={image_encoder.__class__.__name__}') sd_models.clear_caches() image_encoder = model_quant.do_post_load_quant(image_encoder, allow=True) if hasattr(pipe, 'register_modules'): @@ -212,12 +212,16 @@ def load_image_encoder(pipe: DiffusionPipeline, adapter_names: list[str]): else: pipe.image_encoder = image_encoder clip_loaded = f'{clip_repo}/{clip_subfolder}' + pipe = sd_models.apply_balanced_offload(pipe, force=True) except Exception as e: log.error(f'IP adapter load: encoder="{clip_repo}/{clip_subfolder}" {e}') errors.display(e, 'IP adapter: type=encoder') return False shared.state.end(jobid) sd_models.move_model(pipe.image_encoder, devices.device) + if hasattr(pipe.unet, 'balanced_offload_device_map') and pipe.unet.balanced_offload_device_map.get('encoder_hid_proj', None) is None: + # image encoder patches unet, but diffusers creates full device map only during model load so if module is loaded later it will be missing + pipe.unet.balanced_offload_device_map['encoder_hid_proj'] = 0 return True @@ -235,13 +239,14 @@ def load_feature_extractor(pipe): pipe.register_modules(feature_extractor=feature_extractor) else: pipe.feature_extractor = feature_extractor - sd_models.apply_balanced_offload(pipe.feature_extractor) + pipe = sd_models.apply_balanced_offload(pipe, force=True, silent=True) log.debug(f'IP adapter load: extractor={pipe.feature_extractor.__class__.__name__}') except Exception as e: log.error(f'IP adapter load: extractor {e}') errors.display(e, 'IP adapter: type=extractor') return False shared.state.end(jobid) + sd_models.move_model(pipe.feature_extractor, devices.device) return True diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 90a0bd326..80ca7fe58 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -276,13 +276,20 @@ class OffloadHook(accelerate.hooks.ModelHook): log.trace(f'Offload: type=balanced op=dispatch map={device_map}') if device_map is not None: skip_keys = getattr(module, "_skip_keys", None) - module = accelerate.dispatch_model(module, - main_device=torch.device(devices.device), - device_map=device_map, - offload_dir=offload_dir, - skip_keys=skip_keys, - force_hooks=True, - ) + try: + module = accelerate.dispatch_model(module, + main_device=torch.device(devices.device), + device_map=device_map, + offload_dir=offload_dir, + skip_keys=skip_keys, + force_hooks=True, + ) + except Exception as e: # reapply hook + log.warning(f'Offload: type=balanced op=dispatch module={module.__class__.__name__} {e}') + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + module.balanced_offload_device_map = None + sd_models.move_model(module, devices.device, force=True) + module = accelerate.hooks.add_hook_to_module(module, self, append=True) module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access module.balanced_offload_device_map = device_map module.balanced_offload_max_memory = max_memory @@ -498,7 +505,7 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc continue module.module_name = module_name module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) - apply_balanced_offload_to_module(module, op='apply') + apply_balanced_offload_to_module(module, op='apply', force=force) if not silent: report_model_stats(module_name, module)