qnloft
/
ComfyUI
mirror of https://github.com/comfyanonymous/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Small refactor of preview code.

This commit is contained in:
comfyanonymous 2024-05-31 11:18:37 -04:00
parent bf3e334d46
commit 04b308229e
1 changed files with 10 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
latent_preview.py
View File

@ -11,6 +11,13 @@ import logging
MAX_PREVIEW_RESOLUTION = 512
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))
return Image.fromarray(latents_ubyte.numpy())
class LatentPreviewer:
def decode_latent_to_preview(self, x0):
pass
@ -24,12 +31,8 @@ class TAESDPreviewerImpl(LatentPreviewer):
self.taesd = taesd
def decode_latent_to_preview(self, x0):
x_sample = self.taesd.decode(x0[:1])[0].detach()
x_sample = 255. * torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
x_sample = np.moveaxis(x_sample.to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(x_sample.device)).numpy(), 0, 2)
preview_image = Image.fromarray(x_sample)
return preview_image
x_sample = self.taesd.decode(x0[:1])[0].movedim(0, 2)
return preview_to_image(x_sample)
class Latent2RGBPreviewer(LatentPreviewer):
@ -39,13 +42,7 @@ class Latent2RGBPreviewer(LatentPreviewer):
def decode_latent_to_preview(self, x0):
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
).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device))
return Image.fromarray(latents_ubyte.numpy())
return preview_to_image(latent_image)
def get_previewer(device, latent_format):