Run hidream with single llama text encoder

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)

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

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")

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}),