| import torch |
| from PIL import Image |
| import struct |
| import numpy as np |
| from comfy.cli_args import args, LatentPreviewMethod |
| from comfy.taesd.taesd import TAESD |
| import folder_paths |
|
|
| MAX_PREVIEW_RESOLUTION = 512 |
|
|
| class LatentPreviewer: |
| def decode_latent_to_preview(self, x0): |
| pass |
|
|
| def decode_latent_to_preview_image(self, preview_format, x0): |
| preview_image = self.decode_latent_to_preview(x0) |
| return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION) |
|
|
| class TAESDPreviewerImpl(LatentPreviewer): |
| def __init__(self, taesd): |
| self.taesd = taesd |
|
|
| def decode_latent_to_preview(self, x0): |
| x_sample = self.taesd.decoder(x0)[0].detach() |
| |
| x_sample = x_sample.sub(0.5).mul(2) |
|
|
| x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0) |
| x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) |
| x_sample = x_sample.astype(np.uint8) |
|
|
| preview_image = Image.fromarray(x_sample) |
| return preview_image |
|
|
|
|
| class Latent2RGBPreviewer(LatentPreviewer): |
| 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 |
|
|
| latents_ubyte = (((latent_image + 1) / 2) |
| .clamp(0, 1) |
| .mul(0xFF) |
| .byte()).cpu() |
|
|
| return Image.fromarray(latents_ubyte.numpy()) |
|
|
|
|
| def get_previewer(device, latent_format): |
| previewer = None |
| method = args.preview_method |
| if method != LatentPreviewMethod.NoPreviews: |
| |
| taesd_decoder_path = None |
| if latent_format.taesd_decoder_name is not None: |
| taesd_decoder_path = folder_paths.get_full_path("vae_approx", latent_format.taesd_decoder_name) |
|
|
| if method == LatentPreviewMethod.Auto: |
| method = LatentPreviewMethod.Latent2RGB |
| if taesd_decoder_path: |
| method = LatentPreviewMethod.TAESD |
|
|
| if method == LatentPreviewMethod.TAESD: |
| if taesd_decoder_path: |
| taesd = TAESD(None, taesd_decoder_path).to(device) |
| previewer = TAESDPreviewerImpl(taesd) |
| else: |
| print("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name)) |
|
|
| if previewer is None: |
| if latent_format.latent_rgb_factors is not None: |
| previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors) |
| return previewer |
|
|
|
|
|
|
