From 7b3f2ced4e13f7cc57bad3f5e62d791db5f7b0fb Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Thu, 12 Dec 2024 06:59:37 +0100
Subject: [PATCH] Add fast rgb preview to Sana

---
 Sana/loader.py | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/Sana/loader.py b/Sana/loader.py
index 129e898..413ffb0 100644
--- a/Sana/loader.py
+++ b/Sana/loader.py
@@ -15,8 +15,11 @@ from .diffusers_convert import convert_state_dict
 from ..utils.loader import load_state_dict_from_config
 
 class SanaLatent(comfy.latent_formats.LatentFormat):
-    scale_factor = 0.41407
     latent_channels = 32
+    def __init__(self):
+        self.scale_factor = 0.41407
+        self.latent_rgb_factors = latent_rgb_factors.copy()
+        self.latent_rgb_factors_bias = latent_rgb_factors_bias.copy()
 
 class SanaConfig(comfy.supported_models_base.BASE):
     unet_class = SanaMS
@@ -97,6 +100,43 @@ def model_config_from_unet(sd):
     logging.debug(f"Sana config:\n{config}")
     return model_config
 
+# for fast latent preview
+latent_rgb_factors = [
+    [-2.0022e-03, -6.0736e-03, -1.7096e-03],
+    [ 1.4221e-03,  3.6703e-03,  4.1083e-03],
+    [ 1.0081e-02,  2.6456e-04, -1.4333e-02],
+    [-2.4253e-03,  3.0967e-03, -1.0301e-03],
+    [ 2.2158e-03,  7.7274e-03, -1.3151e-02],
+    [ 1.1235e-02,  5.7630e-03,  3.6146e-03],
+    [-7.2899e-02,  1.1062e-02,  3.6103e-02],
+    [ 3.2346e-02,  2.8678e-02,  2.5014e-02],
+    [ 1.6469e-03, -1.1364e-03,  2.8366e-03],
+    [-3.5597e-02, -2.3447e-02, -3.1172e-03],
+    [-1.9985e-04, -2.0647e-03, -1.2702e-02],
+    [ 2.1318e-04,  1.2196e-03, -8.3461e-04],
+    [ 1.3766e-02,  2.7559e-03,  7.3567e-03],
+    [ 1.3027e-02,  2.6365e-03,  3.0405e-03],
+    [ 1.5335e-02,  9.4682e-03,  6.7312e-03],
+    [ 5.1827e-03, -9.4865e-03,  8.5080e-03],
+    [ 1.4365e-02, -3.2867e-03,  9.5108e-03],
+    [-4.1216e-03, -1.9177e-03, -3.3726e-03],
+    [-2.4757e-03,  5.1739e-04,  2.0280e-03],
+    [-3.5950e-03,  1.0720e-03,  5.3043e-03],
+    [-5.1758e-03,  8.1040e-03, -3.7564e-02],
+    [-3.8555e-03, -1.5529e-03, -3.5799e-03],
+    [-6.6175e-03, -6.8484e-03, -9.9609e-03],
+    [-2.1656e-03,  5.5770e-05,  1.4936e-03],
+    [-9.2857e-02, -1.1379e-01, -1.0919e-01],
+    [ 7.7044e-04,  5.5594e-03,  3.4755e-02],
+    [ 1.2714e-02,  2.9729e-02,  3.1989e-03],
+    [-1.1805e-03,  9.0548e-03, -4.1063e-04],
+    [ 8.3309e-04,  4.9694e-03,  2.3087e-03],
+    [ 7.8456e-03,  3.9750e-03,  3.5655e-03],
+    [-1.7552e-03,  4.9306e-03,  1.4210e-02],
+    [-1.4790e-03,  2.8837e-03, -4.5687e-03]
+]
+latent_rgb_factors_bias = [0.4358, 0.3814, 0.3388]
+
 # 512/1024/2K match, TODO: 4K is new, add on release
 from ..PixArt.loader import resolutions as pixart_res
 resolutions = {