Add optional inputs to force execution order between text encoding and model loading

This commit is contained in:
kijai 2024-10-24 19:01:28 +03:00
parent 2c67025577
commit 813d6aa92f

View File

@ -65,6 +65,9 @@ class DownloadAndLoadMochiModel:
"attention_mode": (["sdpa","flash_attn","sage_attn", "comfy"], "attention_mode": (["sdpa","flash_attn","sage_attn", "comfy"],
), ),
}, },
"optional": {
"trigger": ("CONDITIONING", {"tooltip": "Dummy input for forcing execution order",}),
},
} }
RETURN_TYPES = ("MOCHIMODEL", "MOCHIVAE",) RETURN_TYPES = ("MOCHIMODEL", "MOCHIVAE",)
@ -73,7 +76,7 @@ class DownloadAndLoadMochiModel:
CATEGORY = "MochiWrapper" CATEGORY = "MochiWrapper"
DESCRIPTION = "Downloads and loads the selected Mochi model from Huggingface" DESCRIPTION = "Downloads and loads the selected Mochi model from Huggingface"
def loadmodel(self, model, vae, precision, attention_mode): def loadmodel(self, model, vae, precision, attention_mode, trigger=None):
device = mm.get_torch_device() device = mm.get_torch_device()
offload_device = mm.unet_offload_device() offload_device = mm.unet_offload_device()
@ -159,8 +162,8 @@ class MochiTextEncode:
} }
} }
RETURN_TYPES = ("CONDITIONING",) RETURN_TYPES = ("CONDITIONING", "CLIP",)
RETURN_NAMES = ("conditioning",) RETURN_NAMES = ("conditioning", "clip", )
FUNCTION = "process" FUNCTION = "process"
CATEGORY = "MochiWrapper" CATEGORY = "MochiWrapper"
@ -192,7 +195,7 @@ class MochiTextEncode:
"embeds": embeds, "embeds": embeds,
"attention_mask": attention_mask["attention_mask"].bool(), "attention_mask": attention_mask["attention_mask"].bool(),
} }
return (t5_embeds, ) return (t5_embeds, clip,)
class MochiSampler: class MochiSampler: