From a4e679765ec6a7adc6dda326c7bab915e2914bd6 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Tue, 6 May 2025 06:00:01 -0700
Subject: [PATCH 1/2] Change chroma to use Flux shift. (#7961)

---
 comfy/model_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index 3d33086d8..045df1317 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -1111,7 +1111,7 @@ class HiDream(BaseModel):
         return out
 
 class Chroma(Flux):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+    def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.chroma.model.Chroma)
 
     def extra_conds(self, **kwargs):

From 271c9c5b9eb027c682161ed7848bbf640253f973 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Tue, 6 May 2025 06:52:37 -0700
Subject: [PATCH 2/2] Better mem estimation for the LTXV 13B model. (#7963)

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

diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index d5210cfac..a1dea2343 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -785,6 +785,10 @@ class LTXV(supported_models_base.BASE):
     vae_key_prefix = ["vae."]
     text_encoder_key_prefix = ["text_encoders."]
 
+    def __init__(self, unet_config):
+        super().__init__(unet_config)
+        self.memory_usage_factor = (unet_config.get("cross_attention_dim", 2048) / 2048) * 5.5
+
     def get_model(self, state_dict, prefix="", device=None):
         out = model_base.LTXV(self, device=device)
         return out