fix: ensure casting to torch.uint8

Fixes comfyanonymous#3539
dnswd · May 24, 2024 · 8e26f42 · 8e26f42
1 parent b02bcce
commit 8e26f42
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/latent_preview.py b/latent_preview.py
@@ -44,7 +44,9 @@ def decode_latent_to_preview(self, x0):
         latents_ubyte = (((latent_image + 1) / 2)
                             .clamp(0, 1)  # change scale from -1..1 to 0..1
                             .mul(0xFF)  # to 0..255
-                            ).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device))
+                            )
+        latents_ubyte = latents_ubyte.to(dtype=torch.uint8)
+        latents_ubyte = latents_ubyte.to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device))
 
         return Image.fromarray(latents_ubyte.numpy())