From 644a6f9246fcfc24b6ee5e906cafb96b56553464 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 23 May 2026 10:37:33 -0600 Subject: [PATCH] Fix device casting for zimage in some instances --- extensions_built_in/diffusion_models/z_image/z_image.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/extensions_built_in/diffusion_models/z_image/z_image.py b/extensions_built_in/diffusion_models/z_image/z_image.py index 368ae9e7..b5a024f9 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -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))