From 18a117f832cd76df9d68486fa77dd1732ead826b Mon Sep 17 00:00:00 2001 From: Haoming Date: Wed, 25 Mar 2026 12:30:18 +0800 Subject: [PATCH] diffuser --- backend/loader.py | 38 +++++++++++++------ .../packages/huggingface_guess/__init__.py | 10 ++--- .../packages/huggingface_guess/detection.py | 30 +++++---------- 3 files changed, 39 insertions(+), 39 deletions(-) diff --git a/backend/loader.py b/backend/loader.py index b1dd86d1..262d048c 100644 --- a/backend/loader.py +++ b/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(): diff --git a/modules_forge/packages/huggingface_guess/__init__.py b/modules_forge/packages/huggingface_guess/__init__.py index 8facf059..9a315002 100644 --- a/modules_forge/packages/huggingface_guess/__init__.py +++ b/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 diff --git a/modules_forge/packages/huggingface_guess/detection.py b/modules_forge/packages/huggingface_guess/detection.py index 2d5437b7..79bc157d 100644 --- a/modules_forge/packages/huggingface_guess/detection.py +++ b/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):