High quality Stable Cascade Stage C previews via previewer.safetensors

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"`

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):

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))