From 237cab2aa1a8bc4a63616a6335d55f10bf3a1d88 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 14 Aug 2024 17:42:49 +0300 Subject: [PATCH] Add offload check to cascade's vqgan --- modules/model_stablecascade.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modules/model_stablecascade.py b/modules/model_stablecascade.py index d0c35f47a..7ff3569ec 100644 --- a/modules/model_stablecascade.py +++ b/modules/model_stablecascade.py @@ -322,6 +322,10 @@ class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline): ) if not output_type == "latent": + if shared.opts.diffusers_offload_mode == "balanced": + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + else: + self.maybe_free_model_hooks() # 10. Scale and decode the image latents with vq-vae latents = self.vqgan.config.scale_factor * latents images = self.vqgan.decode(latents).sample.clamp(0, 1) @@ -334,9 +338,10 @@ class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline): images = latents # Offload all models - self.maybe_free_model_hooks() if shared.opts.diffusers_offload_mode == "balanced": shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + else: + self.maybe_free_model_hooks() if not return_dict: return images