Slightly faster latent2rgb previews.

This commit is contained in:
comfyanonymous 2024-05-19 17:42:35 -04:00
parent f37a47110b
commit 4ae1515f14
1 changed files with 3 additions and 2 deletions

View File

@ -37,12 +37,13 @@ class Latent2RGBPreviewer(LatentPreviewer):
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())