diff --git a/comfy/sd.py b/comfy/sd.py index d97873ba2..0c5a6dc6e 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -703,6 +703,7 @@ class CLIPType(Enum): COSMOS = 11 LUMINA2 = 12 WAN = 13 + HIDREAM = 14 def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}): @@ -827,6 +828,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data)) clip_target.tokenizer = comfy.text_encoders.lumina2.LuminaTokenizer tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None) + elif te_model == TEModel.LLAMA3_8: + print("Single LLAMA3_8 for HiDreams") + clip_target.clip = comfy.text_encoders.hidream.hidream_clip(False, **llama_detect(clip_data), clip_g=False, t5=False) + clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer else: if clip_type == CLIPType.SD3: clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=False, t5=False) diff --git a/comfy/text_encoders/hidream.py b/comfy/text_encoders/hidream.py index 6c34c5572..cce1e3cce 100644 --- a/comfy/text_encoders/hidream.py +++ b/comfy/text_encoders/hidream.py @@ -5,6 +5,7 @@ from comfy import sdxl_clip import comfy.model_management import torch import logging +import folder_paths class HiDreamTokenizer: @@ -91,38 +92,51 @@ class HiDreamTEModel(torch.nn.Module): token_weight_pairs_llama = token_weight_pairs["llama"] lg_out = None pooled = None + t5_out = None + ll_out = None extra = {} if len(token_weight_pairs_g) > 0 or len(token_weight_pairs_l) > 0: if self.clip_l is not None: + print("Encoding clip_l token weights") lg_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs_l) else: l_pooled = torch.zeros((1, 768), device=comfy.model_management.intermediate_device()) if self.clip_g is not None: + print("Encoding clip_g token weights") g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) else: g_pooled = torch.zeros((1, 1280), device=comfy.model_management.intermediate_device()) - - pooled = torch.cat((l_pooled, g_pooled), dim=-1) + + if self.clip_g is not None and self.clip_l is not None: + pooled = torch.cat((l_pooled, g_pooled), dim=-1) if self.t5xxl is not None: + print("Encoding t5 token weights") t5_output = self.t5xxl.encode_token_weights(token_weight_pairs_t5) t5_out, t5_pooled = t5_output[:2] if self.llama is not None: + print("Encoding llama token weights") ll_output = self.llama.encode_token_weights(token_weight_pairs_llama) ll_out, ll_pooled = ll_output[:2] ll_out = ll_out[:, 1:] if t5_out is None: - t5_out = torch.zeros((1, 1, 4096), device=comfy.model_management.intermediate_device()) + print("Loading t5_out from disk") + t5_path = folder_paths.get_full_path_or_raise("hidream_empty_latents", "t5_out.pt") + t5_out = torch.load(t5_path, map_location=comfy.model_management.intermediate_device()) if ll_out is None: + print("No llama encoder found, filling with zeroes") ll_out = torch.zeros((1, 32, 1, 4096), device=comfy.model_management.intermediate_device()) - if pooled is None: - pooled = torch.zeros((1, 768 + 1280), device=comfy.model_management.intermediate_device()) + if pooled is None: + print("Loading pooled from disk") + pooled_path = folder_paths.get_full_path_or_raise("hidream_empty_latents", "pooled.pt") + pooled = torch.load(pooled_path, map_location=comfy.model_management.intermediate_device()) + extra["conditioning_llama3"] = ll_out return t5_out, pooled, extra diff --git a/folder_paths.py b/folder_paths.py index 9a525e5a1..9e055c574 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -45,6 +45,7 @@ folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetwo folder_names_and_paths["photomaker"] = ([os.path.join(models_dir, "photomaker")], supported_pt_extensions) folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""}) +folder_names_and_paths["hidream_empty_latents"] = ([os.path.join(models_dir, "hidream_empty_latents")], supported_pt_extensions) output_directory = os.path.join(base_path, "output") temp_directory = os.path.join(base_path, "temp") diff --git a/models/hidream_empty_latents/pooled.pt b/models/hidream_empty_latents/pooled.pt new file mode 100644 index 000000000..2b1ee01bf Binary files /dev/null and b/models/hidream_empty_latents/pooled.pt differ diff --git a/models/hidream_empty_latents/t5_out.pt b/models/hidream_empty_latents/t5_out.pt new file mode 100644 index 000000000..d09b73f77 Binary files /dev/null and b/models/hidream_empty_latents/t5_out.pt differ diff --git a/nodes.py b/nodes.py index ae0a2e183..e4147909f 100644 --- a/nodes.py +++ b/nodes.py @@ -917,7 +917,7 @@ class CLIPLoader: @classmethod def INPUT_TYPES(s): return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ), - "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan"], ), + "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "llama"], ), }, "optional": { "device": (["default", "cpu"], {"advanced": True}),