github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added decode latent method to qwen image model

Browse Source
This commit is contained in:
Jaret Burkett
2026-05-25 10:02:09 -06:00
parent 3d836ac371
commit 7354def271
1 changed files with 43 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -125,11 +125,14 @@ class QwenImageModel(BaseModel):
quantize_model(self, transformer) quantize_model(self, transformer)
flush() 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( MemoryManager.attach(
transformer, transformer,
self.device_torch, 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: if self.model_config.low_vram:
@@ -151,11 +154,14 @@ class QwenImageModel(BaseModel):
if not self._qwen_image_keep_visual: if not self._qwen_image_keep_visual:
text_encoder.model.visual = None 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( MemoryManager.attach(
text_encoder, text_encoder,
self.device_torch, 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) text_encoder.to(self.device_torch, dtype=dtype)
@@ -354,7 +360,7 @@ class QwenImageModel(BaseModel):
def get_prompt_embeds(self, prompt: str) -> PromptEmbeds: def get_prompt_embeds(self, prompt: str) -> PromptEmbeds:
if self.pipeline.text_encoder.device != self.device_torch: if self.pipeline.text_encoder.device != self.device_torch:
self.pipeline.text_encoder.to(self.device_torch) self.pipeline.text_encoder.to(self.device_torch)
prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt( prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt(
prompt, prompt,
device=self.device_torch, device=self.device_torch,
@@ -446,3 +452,35 @@ class QwenImageModel(BaseModel):
latents = latents.squeeze(2) # remove the frame count dimension latents = latents.squeeze(2) # remove the frame count dimension
return latents 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)