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): diff --git a/comfy/supported_models.py b/comfy/supported_models.py index a8c5b8477..696ddb2eb 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