From 045b24c06014f37517c0616f8d6f26c8962001ab Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 2 Sep 2024 08:23:40 -0400 Subject: [PATCH] fix vae decode for backend original --- .vscode/settings.json | 3 ++- CHANGELOG.md | 3 +++ modules/processing_helpers.py | 3 ++- modules/sd_hijack.py | 1 - modules/sd_models.py | 18 ++++++++++++++++++ 5 files changed, 25 insertions(+), 3 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index 428fc2335..690477350 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -8,5 +8,6 @@ "./repositories/taming" ], "python.analysis.typeCheckingMode": "off", - "editor.formatOnSave": false + "editor.formatOnSave": false, + "python.REPL.enableREPLSmartSend": false } \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 9be12ecf4..e84b92245 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,7 +38,10 @@ Other improvements: improves quality of the flow-matching samplers - t5 support manually downloaded models applies to all models that use t5 transformer + +Fixes: - fix handling of model configs if offline config is not available +- fix vae decode in backend original Work-in-progress: - flux controlnet support: (*1) diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 47839ed43..6dfbdc535 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -196,7 +196,8 @@ def decode_first_stage(model, x, full_quality=True): try: if full_quality: if hasattr(model, 'decode_first_stage'): - x_sample = model.decode_first_stage(x) * 0.5 + 0.5 + # x_sample = model.decode_first_stage(x) * 0.5 + 0.5 + x_sample = model.decode_first_stage(x) elif hasattr(model, 'vae'): x_sample = processing_vae.vae_decode(latents=x, model=model, output_type='np', full_quality=full_quality) else: diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index dc07dea81..2f707d098 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -265,7 +265,6 @@ class EmbeddingsWithFixes(torch.nn.Module): def forward(self, input_ids): batch_fixes = self.embeddings.fixes self.embeddings.fixes = None - inputs_embeds = self.wrapped(input_ids) if batch_fixes is None or len(batch_fixes) == 0 or max([len(x) for x in batch_fixes]) == 0: diff --git a/modules/sd_models.py b/modules/sd_models.py index 5f19c409a..19fdee91f 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -737,6 +737,9 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): set_diffuser_offload(sd_model, op) def set_diffuser_offload(sd_model, op: str = 'model'): + if not shared.native: + shared.log.warning('Attempting to use offload with backend=original') + return if sd_model is None: shared.log.warning(f'{op} is not loaded') return @@ -837,6 +840,8 @@ def apply_balanced_offload(sd_model): shared.log.error(f'Balanced offload: module={module_name} {e}') devices.torch_gc(fast=True) + if not shared.native: + return apply_balanced_offload_to_module(sd_model) if hasattr(sd_model, "prior_pipe"): apply_balanced_offload_to_module(sd_model.prior_pipe) @@ -859,6 +864,19 @@ def normalize_device(device): def move_model(model, device=None, force=False): if model is None or device is None: return + + if not shared.native: + if type(model).__name__ == 'LatentDiffusion': + model = model.to(device) + if hasattr(model, 'model'): + model.model = model.model.to(device) + if hasattr(model, 'first_stage_model'): + model.first_stage_model = model.first_stage_model.to(device) + if hasattr(model, 'cond_stage_model'): + model.cond_stage_model = model.cond_stage_model.to(device) + devices.torch_gc() + return + if getattr(model, 'vae', None) is not None and get_diffusers_task(model) != DiffusersTaskType.TEXT_2_IMAGE: if device == devices.device and model.vae.device.type != "meta": # force vae back to gpu if not in txt2img mode model.vae.to(device)