RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix device casting for zimage in some instances

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
Show Stats 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))