diff --git a/comfy/text_encoders/hidream.py b/comfy/text_encoders/hidream.py index 600b34480..cf4e79241 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,6 +92,8 @@ 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: @@ -104,8 +107,9 @@ class HiDreamTEModel(torch.nn.Module): 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: t5_output = self.t5xxl.encode_token_weights(token_weight_pairs_t5) t5_out, t5_pooled = t5_output[:2] @@ -120,13 +124,15 @@ class HiDreamTEModel(torch.nn.Module): ll_out = None if t5_out is None: - t5_out = torch.zeros((1, 128, 4096), device=comfy.model_management.intermediate_device()) + t5_path = folder_paths.get_full_path_or_raise("hidream_empty_latents", "t5_blank.pt") + t5_out = torch.load(t5_path, map_location=comfy.model_management.intermediate_device()) if ll_out is None: 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()) + pooled_path = folder_paths.get_full_path_or_raise("hidream_empty_latents", "pooled_blank.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 9c96540e3..3bd9f5763 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -47,6 +47,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) folder_names_and_paths["model_patches"] = ([os.path.join(models_dir, "model_patches")], supported_pt_extensions) diff --git a/models/hidream_empty_latents/pooled_blank.pt b/models/hidream_empty_latents/pooled_blank.pt new file mode 100644 index 000000000..594991e92 Binary files /dev/null and b/models/hidream_empty_latents/pooled_blank.pt differ diff --git a/models/hidream_empty_latents/t5_blank.pt b/models/hidream_empty_latents/t5_blank.pt new file mode 100644 index 000000000..3ebb5b7ac Binary files /dev/null and b/models/hidream_empty_latents/t5_blank.pt differ