From 3711b31dff3530ba584e9a30f1bb32feaf2b5886 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 19 Feb 2024 11:20:48 -0500
Subject: [PATCH] Support Stable Cascade in checkpoint format.

---
 comfy/supported_models.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index f29f7f3d9..5bb98d88a 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -319,6 +319,10 @@ class Stable_Cascade_C(supported_models_base.BASE):
         "shift": 2.0,
     }
 
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoder."]
+    clip_vision_prefix = "clip_l_vision."
+
     def process_unet_state_dict(self, state_dict):
         key_list = list(state_dict.keys())
         for y in ["weight", "bias"]:
@@ -355,6 +359,8 @@ class Stable_Cascade_B(Stable_Cascade_C):
         "shift": 1.0,
     }
 
+    clip_vision_prefix = None
+
     def get_model(self, state_dict, prefix="", device=None):
         out = model_base.StableCascade_B(self, device=device)
         return out