From cd912963f17c9ae00ec12e1869293edb78720831 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Sat, 10 Jan 2026 14:31:31 -0800
Subject: [PATCH] Fix issue with t5 text encoder in fp4. (#11794)

---
 comfy/model_detection.py | 2 ++
 comfy/sd.py              | 4 ++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index 0853b3aec..aff5a50b9 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -237,6 +237,8 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         else:
             dit_config["vec_in_dim"] = None
 
+        dit_config["num_heads"] = dit_config["hidden_size"] // sum(dit_config["axes_dim"])
+
         dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
         dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')
         if '{}distilled_guidance_layer.0.norms.0.scale'.format(key_prefix) in state_dict_keys or '{}distilled_guidance_layer.norms.0.scale'.format(key_prefix) in state_dict_keys: #Chroma
diff --git a/comfy/sd.py b/comfy/sd.py
index 5a7221620..b689c0dfc 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -1059,9 +1059,9 @@ def detect_te_model(sd):
         return TEModel.JINA_CLIP_2
     if "encoder.block.23.layer.1.DenseReluDense.wi_1.weight" in sd:
         weight = sd["encoder.block.23.layer.1.DenseReluDense.wi_1.weight"]
-        if weight.shape[-1] == 4096:
+        if weight.shape[0] == 10240:
             return TEModel.T5_XXL
-        elif weight.shape[-1] == 2048:
+        elif weight.shape[0] == 5120:
             return TEModel.T5_XL
     if 'encoder.block.23.layer.1.DenseReluDense.wi.weight' in sd:
         return TEModel.T5_XXL_OLD