diffuser

backend/loader.py

@@ -319,10 +319,6 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
             load_device = memory_management.get_torch_device()
             offload_device = memory_management.unet_offload_device()
 
-            if (new_dict := convert_diffusers_mmdit(state_dict, "")) is not None:
-                del state_dict
-                state_dict = new_dict
-
             unet_config = guess.unet_config.copy()
             state_dict_parameters = utils.calculate_parameters(state_dict)
             state_dict_dtype = utils.weight_dtype(state_dict)
@@ -665,14 +661,37 @@ def process_anima(dit: dict[str, torch.Tensor], enc: dict[str, torch.Tensor]):
             enc[k] = dit.pop(k)
 
 
-def split_state_dict(sd, additional_state_dicts: list = None):
+def _load_unet(path: os.PathLike):
     import huggingface_guess
 
-    sd, metadata = load_torch_file(sd, return_metadata=True)
+    sd, metadata = load_torch_file(path, return_metadata=True)
     sd, metadata = convert_quantization(sd, metadata)
     sd = preprocess_state_dict(sd)
     guess = huggingface_guess.guess(sd)
 
+    return sd, metadata, guess
+
+
+def _load_diffuser(path: os.PathLike):
+    import huggingface_guess
+
+    sd, metadata = load_torch_file(path, return_metadata=True)
+    sd, metadata = convert_quantization(sd, metadata)
+    sd = convert_diffusers_mmdit(sd, "")
+    sd = preprocess_state_dict(sd)
+    guess = huggingface_guess.guess(sd)
+
+    return sd, metadata, guess
+
+
+def split_state_dict(path: os.PathLike, additional_state_dicts: list[os.PathLike] = None):
+    try:
+        sd, metadata, guess = _load_unet(path)
+    except Exception:
+        sd, metadata, guess = _load_diffuser(path)
+    finally:
+        memory_management.soft_empty_cache()
+
     if getattr(guess, "nunchaku", False) and ("Z-Image" in guess.huggingface_repo or "Qwen" in guess.huggingface_repo):
         import json
 
@@ -723,11 +742,8 @@ def split_state_dict(sd, additional_state_dicts: list = None):
 def forge_loader(sd: os.PathLike, additional_state_dicts: list[os.PathLike] = None):
     try:
         state_dicts, estimated_config = split_state_dict(sd, additional_state_dicts=additional_state_dicts)
-    except Exception as e:
-        from modules.errors import display
-
-        display(e, "forge_loader")
-        raise ValueError("Failed to recognize model type!")
+    except Exception:
+        raise ValueError("Failed to recognize model type!") from None
 
     repo_name = estimated_config.huggingface_repo
     if "xl" in repo_name and "rectified" in str(sd).lower():

modules_forge/packages/huggingface_guess/__init__.py

@@ -20,14 +20,10 @@ from .detection import model_config_from_unet, unet_prefix_from_state_dict
 
 def guess(state_dict):
     unet_key_prefix = unet_prefix_from_state_dict(state_dict)
-    result = model_config_from_unet(
-        state_dict, unet_key_prefix, use_base_if_no_match=False
-    )
+    result = model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False)
     result.unet_key_prefix = [unet_key_prefix]
-    if "image_model" in result.unet_config:
-        del result.unet_config["image_model"]
-    if "audio_model" in result.unet_config:
-        del result.unet_config["audio_model"]
+    result.unet_config.pop("image_model", None)
+    result.unet_config.pop("audio_model", None)
     return result
 
 

modules_forge/packages/huggingface_guess/detection.py

@@ -37,10 +37,10 @@ def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
     return None
 
 
-def detect_unet_config(state_dict: dict, key_prefix: str):
+def detect_unet_config(state_dict: dict, key_prefix: str) -> dict:
     state_dict_keys = list(state_dict.keys())
 
-    if "{}cap_embedder.1.weight".format(key_prefix) in state_dict_keys:  # Lumina 2
+    if "{}cap_embedder.1.weight".format(key_prefix) in state_dict_keys and "{}noise_refiner.0.attention.k_norm.weight".format(key_prefix) in state_dict_keys:  # Lumina 2
         dit_config = {}
         dit_config["image_model"] = "lumina2"
         dit_config["patch_size"] = 2
@@ -58,9 +58,6 @@ def detect_unet_config(state_dict: dict, key_prefix: str):
             dit_config["axes_lens"] = [300, 512, 512]
             dit_config["rope_theta"] = 10000.0
             dit_config["ffn_dim_multiplier"] = 4.0
-            ctd_weight = state_dict.get("{}clip_text_pooled_proj.0.weight".format(key_prefix), None)
-            if ctd_weight is not None:  # NewBie
-                dit_config["clip_text_dim"] = int(ctd_weight.shape[0])
         elif dit_config["dim"] == 3840:  # Z-image
             dit_config["nunchaku"] = "{}layers.0.attention.to_out.0.qweight".format(key_prefix) in state_dict_keys
             dit_config["n_heads"] = 30
@@ -395,24 +392,15 @@ def top_candidate(state_dict, candidates):
     return top, counts[top]
 
 
-def unet_prefix_from_state_dict(state_dict):
-    candidates = [
+def unet_prefix_from_state_dict(state_dict: dict) -> str:
+    candidates = (
         "model.diffusion_model.",  # ldm/sgm models
-        "model.model.",  # audio models
         "net.",  # cosmos
-    ]
-    counts = {k: 0 for k in candidates}
-    for k in state_dict:
-        for c in candidates:
-            if k.startswith(c):
-                counts[c] += 1
-                break
-
-    top = max(counts, key=counts.get)
-    if counts[top] > 5:
-        return top
-    else:
-        return "model."  # etc.
+    )
+    for prefix in candidates:
+        if sum(1 for k in state_dict if k.startswith(prefix)) > 5:
+            return prefix
+    return "model."  # etc.
 
 
 def convert_config(unet_config):