From 96be1c553c01e01cc96483064424e2e6ad7e6f4d Mon Sep 17 00:00:00 2001 From: asagi4 <130366179+asagi4@users.noreply.github.com> Date: Mon, 23 Sep 2024 14:01:08 +0300 Subject: [PATCH 1/3] High quality Stable Cascade Stage C previews via previewer.safetensors --- README.md | 2 +- comfy/latent_formats.py | 1 + latent_preview.py | 17 ++++++++++++++++- 3 files changed, 18 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 745cf1a35d8..5cda4fe754d 100644 --- a/README.md +++ b/README.md @@ -232,7 +232,7 @@ To use a textual inversion concepts/embeddings in a text prompt put them in the Use ```--preview-method auto``` to enable previews. -The default installation includes a fast latent preview method that's low-resolution. To enable higher-quality previews with [TAESD](https://github.com/madebyollin/taesd), download the [taesd_decoder.pth, taesdxl_decoder.pth, taesd3_decoder.pth and taef1_decoder.pth](https://github.com/madebyollin/taesd/) and place them in the `models/vae_approx` folder. Once they're installed, restart ComfyUI and launch it with `--preview-method taesd` to enable high-quality previews. +The default installation includes a fast latent preview method that's low-resolution. To enable higher-quality previews with [TAESD](https://github.com/madebyollin/taesd) or the Stable Cascade previewer, download [taesd_decoder.pth, taesdxl_decoder.pth, taesd3_decoder.pth and taef1_decoder.pth](https://github.com/madebyollin/taesd/) and/or [previewer.safetensors](https://huggingface.co/stabilityai/stable-cascade/resolve/main/previewer.safetensors) and place them in the `models/vae_approx` folder. Once they're installed, restart ComfyUI and launch it with `--preview-method taesd` to enable high-quality previews. ## How to use TLS/SSL? Generate a self-signed certificate (not appropriate for shared/production use) and key by running the command: `openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -sha256 -days 3650 -nodes -subj "/C=XX/ST=StateName/L=CityName/O=CompanyName/OU=CompanySectionName/CN=CommonNameOrHostname"` diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index ee19faeae74..ecebdc3e37c 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -95,6 +95,7 @@ def __init__(self): [ 0.0542, 0.1545, 0.1325], [-0.0352, -0.1672, -0.2541] ] + taesd_decoder_name = "previewer.safetensors" class SC_B(LatentFormat): def __init__(self): diff --git a/latent_preview.py b/latent_preview.py index e14c72ce4d0..8474dd1e9a1 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -8,6 +8,7 @@ import folder_paths import comfy.utils import logging +from comfy.ldm.cascade.stage_c_coder import Previewer MAX_PREVIEW_RESOLUTION = args.preview_size @@ -35,6 +36,17 @@ def decode_latent_to_preview(self, x0): return preview_to_image(x_sample) +class StageCPreviewer(Previewer): + def __init__(self, path): + super().__init__() + sd = comfy.utils.load_torch_file(path, safe_load=True) + self.load_state_dict(sd, strict=True) + self.eval() + + def decode(self, latent): + return self(latent) + + class Latent2RGBPreviewer(LatentPreviewer): def __init__(self, latent_rgb_factors): self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu") @@ -64,7 +76,10 @@ def get_previewer(device, latent_format): if method == LatentPreviewMethod.TAESD: if taesd_decoder_path: - taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device) + if 'previewer' in taesd_decoder_path: + taesd = StageCPreviewer(taesd_decoder_path).to(device) + else: + taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device) previewer = TAESDPreviewerImpl(taesd) else: logging.warning("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name)) From 433a02b2e826d121920c185000e0d3228866b3de Mon Sep 17 00:00:00 2001 From: asagi4 <130366179+asagi4@users.noreply.github.com> Date: Tue, 24 Sep 2024 18:52:40 +0300 Subject: [PATCH 2/3] Refactor previewer model loading into the LatentFormat class also cleans up some unused imports in latent_preview.py --- README.md | 2 +- comfy/latent_formats.py | 34 +++++++++++++++++++++++++++++++- latent_preview.py | 43 ++++++----------------------------------- 3 files changed, 40 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 5cda4fe754d..1132674e293 100644 --- a/README.md +++ b/README.md @@ -232,7 +232,7 @@ To use a textual inversion concepts/embeddings in a text prompt put them in the Use ```--preview-method auto``` to enable previews. -The default installation includes a fast latent preview method that's low-resolution. To enable higher-quality previews with [TAESD](https://github.com/madebyollin/taesd) or the Stable Cascade previewer, download [taesd_decoder.pth, taesdxl_decoder.pth, taesd3_decoder.pth and taef1_decoder.pth](https://github.com/madebyollin/taesd/) and/or [previewer.safetensors](https://huggingface.co/stabilityai/stable-cascade/resolve/main/previewer.safetensors) and place them in the `models/vae_approx` folder. Once they're installed, restart ComfyUI and launch it with `--preview-method taesd` to enable high-quality previews. +The default installation includes a fast latent preview method that's low-resolution. To enable higher-quality previews with [TAESD](https://github.com/madebyollin/taesd) or the Stable Cascade previewer, download [taesd_decoder.pth, taesdxl_decoder.pth, taesd3_decoder.pth and taef1_decoder.pth](https://github.com/madebyollin/taesd/) and/or [previewer.safetensors](https://huggingface.co/stabilityai/stable-cascade/resolve/main/previewer.safetensors) and place them in the `models/vae_approx` folder (save `previewer.safetensors` as `cascade_previewer.safetensors`). Once they're installed, restart ComfyUI and launch it with `--preview-method taesd` to enable high-quality previews. ## How to use TLS/SSL? Generate a self-signed certificate (not appropriate for shared/production use) and key by running the command: `openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -sha256 -days 3650 -nodes -subj "/C=XX/ST=StateName/L=CityName/O=CompanyName/OU=CompanySectionName/CN=CommonNameOrHostname"` diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index ecebdc3e37c..dfb0c69aaac 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -1,10 +1,32 @@ import torch +import folder_paths +import logging + +from comfy.taesd.taesd import TAESD +from comfy.ldm.cascade.stage_c_coder import Previewer +import comfy.utils + class LatentFormat: scale_factor = 1.0 latent_channels = 4 latent_rgb_factors = None taesd_decoder_name = None + # Default if decoder name is defined + previewer_class = TAESD + + def load_previewer(self, device): + model = None + if not self.taesd_decoder_name: + return None + + filename = next((fn for fn in folder_paths.get_filename_list("vae_approx") if fn.startswith(self.taesd_decoder_name)), "") + model_path = folder_paths.get_full_path("vae_approx", filename) + if model_path: + model = self.previewer_class(decoder_path=model_path, latent_channels=self.latent_channels).to(device) + if not model: + logging.warning("Warning: Could not load previewer model: models/vae_approx/%s", self.taesd_decoder_name) + return model def process_in(self, latent): return latent * self.scale_factor @@ -73,8 +95,19 @@ def __init__(self): [ 0.2523, -0.0055, -0.1651] ] +class CascadePreviewWrapper(Previewer): + def __init__(self, decoder_path=None, **kwargs): + super().__init__() + self.load_state_dict(comfy.utils.load_torch_file(decoder_path, safe_load=True), strict=True) + self.eval() + + def decode(self, latent): + return self(latent) + class SC_Prior(LatentFormat): latent_channels = 16 + taesd_decoder_name = "cascade_previewer" + previewer_class = CascadePreviewWrapper def __init__(self): self.scale_factor = 1.0 self.latent_rgb_factors = [ @@ -95,7 +128,6 @@ def __init__(self): [ 0.0542, 0.1545, 0.1325], [-0.0352, -0.1672, -0.2541] ] - taesd_decoder_name = "previewer.safetensors" class SC_B(LatentFormat): def __init__(self): diff --git a/latent_preview.py b/latent_preview.py index 8474dd1e9a1..189699f11ca 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -1,14 +1,8 @@ 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 comfy.model_management -import folder_paths import comfy.utils -import logging -from comfy.ldm.cascade.stage_c_coder import Previewer MAX_PREVIEW_RESOLUTION = args.preview_size @@ -36,17 +30,6 @@ def decode_latent_to_preview(self, x0): return preview_to_image(x_sample) -class StageCPreviewer(Previewer): - def __init__(self, path): - super().__init__() - sd = comfy.utils.load_torch_file(path, safe_load=True) - self.load_state_dict(sd, strict=True) - self.eval() - - def decode(self, latent): - return self(latent) - - class Latent2RGBPreviewer(LatentPreviewer): def __init__(self, latent_rgb_factors): self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu") @@ -57,32 +40,18 @@ def decode_latent_to_preview(self, x0): return preview_to_image(latent_image) -def get_previewer(device, latent_format): +def get_previewer(device, latent_format, method=None): previewer = None - method = args.preview_method + if method is None: + method = args.preview_method if method != LatentPreviewMethod.NoPreviews: - # TODO previewer methods - taesd_decoder_path = None - if latent_format.taesd_decoder_name is not None: - taesd_decoder_path = next( - (fn for fn in folder_paths.get_filename_list("vae_approx") - if fn.startswith(latent_format.taesd_decoder_name)), - "" - ) - taesd_decoder_path = folder_paths.get_full_path("vae_approx", taesd_decoder_path) - if method == LatentPreviewMethod.Auto: method = LatentPreviewMethod.Latent2RGB if method == LatentPreviewMethod.TAESD: - if taesd_decoder_path: - if 'previewer' in taesd_decoder_path: - taesd = StageCPreviewer(taesd_decoder_path).to(device) - else: - taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device) - previewer = TAESDPreviewerImpl(taesd) - else: - logging.warning("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name)) + model = latent_format.load_previewer(device) + if model: + previewer = TAESDPreviewerImpl(model) if previewer is None: if latent_format.latent_rgb_factors is not None: From 18183ad22d03802e0c50573bbf9eb214bb0ea368 Mon Sep 17 00:00:00 2001 From: asagi4 <130366179+asagi4@users.noreply.github.com> Date: Tue, 24 Sep 2024 19:14:36 +0300 Subject: [PATCH 3/3] Avoid breaking popular extension --- latent_preview.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/latent_preview.py b/latent_preview.py index 189699f11ca..94ce29c9f93 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -4,6 +4,9 @@ import comfy.model_management import comfy.utils +# Import this here to avoid breaking ComfyUI-Impact-Pack +from comfy.taesd.taesd import TAESD + MAX_PREVIEW_RESOLUTION = args.preview_size def preview_to_image(latent_image):