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

Fix device casting for zimage in some instances

Browse Source
This commit is contained in:
Jaret Burkett
2026-05-23 10:37:33 -06:00
parent c6ecc03ccd
commit 644a6f9246
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
extensions_built_in/diffusion_models/z_image/z_image.py
View File

@@ -300,8 +300,8 @@ class ZImageModel(BaseModel):
generator: torch.Generator,
extra: dict,
):
self.model.to(self.device_torch, dtype=self.torch_dtype)
self.model.to(self.device_torch)
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
sc = self.get_bucket_divisibility()
gen_config.width = int(gen_config.width // sc * sc)
@@ -326,7 +326,8 @@ class ZImageModel(BaseModel):
text_embeddings: PromptEmbeds,
**kwargs,
):
self.model.to(self.device_torch)
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
latent_model_input = latent_model_input.unsqueeze(2)
latent_model_input_list = list(latent_model_input.unbind(dim=0))