From 1210d094c711cb798ee1ce0198eedc2925a247b5 Mon Sep 17 00:00:00 2001 From: Bradley Reynolds Date: Tue, 28 Jan 2025 07:22:54 -0600 Subject: [PATCH] Convert `latents_ubyte` to 8-bit unsigned int before converting to CPU (#6300) * Convert latents_ubyte to 8-bit unsigned int before converting to CPU * Only convert to unint8 if directml_enabled --- latent_preview.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/latent_preview.py b/latent_preview.py index 07f9cc68..95d3cb73 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -12,7 +12,10 @@ MAX_PREVIEW_RESOLUTION = args.preview_size def preview_to_image(latent_image): latents_ubyte = (((latent_image + 1.0) / 2.0).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)) + ) + if comfy.model_management.directml_enabled: + 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())