Slightly faster latent2rgb previews.

Author: comfyanonymous

latent_preview.py

@@ -37,12 +37,13 @@ def __init__(self, latent_rgb_factors):
         self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu")
 
     def decode_latent_to_preview(self, x0):
-        latent_image = x0[0].permute(1, 2, 0).cpu() @ self.latent_rgb_factors
+        self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device)
+        latent_image = x0[0].permute(1, 2, 0) @ self.latent_rgb_factors
 
         latents_ubyte = (((latent_image + 1) / 2)
                             .clamp(0, 1)  # change scale from -1..1 to 0..1
                             .mul(0xFF)  # to 0..255
-                            .byte()).cpu()
+                            ).to(device="cpu", dtype=torch.uint8, non_blocking=True)
 
         return Image.fromarray(latents_ubyte.numpy())