From 7354def271d94ad23107d3a96bccee2e94e4277e Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 25 May 2026 10:02:09 -0600 Subject: [PATCH] Added decode latent method to qwen image model --- .../diffusion_models/qwen_image/qwen_image.py | 48 +++++++++++++++++-- 1 file changed, 43 insertions(+), 5 deletions(-) diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index 0e1e0ee..2726ab0 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -125,11 +125,14 @@ class QwenImageModel(BaseModel): quantize_model(self, transformer) flush() - if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0: + if ( + self.model_config.layer_offloading + and self.model_config.layer_offloading_transformer_percent > 0 + ): MemoryManager.attach( transformer, self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent + offload_percent=self.model_config.layer_offloading_transformer_percent, ) if self.model_config.low_vram: @@ -151,11 +154,14 @@ class QwenImageModel(BaseModel): if not self._qwen_image_keep_visual: text_encoder.model.visual = None - if self.model_config.layer_offloading and self.model_config.layer_offloading_text_encoder_percent > 0: + if ( + self.model_config.layer_offloading + and self.model_config.layer_offloading_text_encoder_percent > 0 + ): MemoryManager.attach( text_encoder, self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent + offload_percent=self.model_config.layer_offloading_text_encoder_percent, ) text_encoder.to(self.device_torch, dtype=dtype) @@ -354,7 +360,7 @@ class QwenImageModel(BaseModel): def get_prompt_embeds(self, prompt: str) -> PromptEmbeds: if self.pipeline.text_encoder.device != self.device_torch: self.pipeline.text_encoder.to(self.device_torch) - + prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt( prompt, device=self.device_torch, @@ -446,3 +452,35 @@ class QwenImageModel(BaseModel): latents = latents.squeeze(2) # remove the frame count dimension return latents + + def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + + latents = latents.to(device, dtype=dtype) + + # add frame count dim for wan vae + latents = latents.unsqueeze(2) + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + + images = self.vae.decode(latents).sample + + images = images.squeeze(2) # remove the frame count dimension + + return images.to(device, dtype=dtype)