diffsynth 2.1.2__py3-none-any.whl → 2.1.4__py3-none-any.whl

diffsynth/configs/model_configs.py

@@ -339,6 +339,14 @@ flux_series = [
         "model_class": "diffsynth.models.flux_dit.FluxDiT",
         "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverterFromDiffusers",
     },
+    {
+        # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="flux1-fill-dev.safetensors")
+        "model_hash": "f876d2a6fb247fb7da0f84fc7fb8a823",
+        "model_name": "flux_dit",
+        "model_class": "diffsynth.models.flux_dit.FluxDiT",
+        "extra_kwargs": {"input_dim": 384},
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_dit.FluxDiTStateDictConverter",
+    },
     {
         # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors")
         "model_hash": "94eefa3dac9cec93cb1ebaf1747d7b78",
@@ -481,6 +489,21 @@ flux_series = [
         "model_class": "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M",
         "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.SiglipStateDictConverter",
     },
+    {
+        # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_encoder/model.safetensors")
+        "model_hash": "6bd5343828144a2135d783519c665e2c",
+        "model_name": "flux_redux_image_encoder",
+        "model_class": "diffsynth.models.flux_ipadapter.SiglipVisionModelSO400M",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.flux_ipadapter.SiglipReduxStateDictConverter",
+    },
+    {
+        # Example: ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_embedder/diffusion_pytorch_model.safetensors")
+        # 2-layer MLP projecting SigLIP tokens (1152) to the 4096-dim FLUX text
+        # space; keys already match FluxReduxImageEncoder, so no converter needed.
+        "model_hash": "d8ea4618391460bc831fddaae5fb5288",
+        "model_name": "flux_redux_image_embedder",
+        "model_class": "diffsynth.models.flux_redux.FluxReduxImageEncoder",
+    },
     {
         # Example: ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors"),
         "model_hash": "d30fb9e02b1dbf4e509142f05cf7dd50",
@@ -1394,18 +1417,46 @@ minimax_h3_series = [
         "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_text_encoder.MiniMaxH3TextEncoderStateDictConverter",
         "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["qkv", "proj", "linear_fc1", "linear_fc2"]},
     },
+    {
+        # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors")
+        "model_hash": "d6473c909b72b73afc4fa00c273d1da2",
+        "model_name": "minimax_h3_text_encoder",
+        "model_class": "diffsynth.models.minimax_h3_text_encoder.MiniMaxH3TextEncoder",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_h3_text_encoder.MiniMaxH3TextEncoderComfyOrgStateDictConverter",
+        "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["qkv", "proj", "linear_fc1", "linear_fc2"]},
+    },
     {
         # Example: ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/transformer/model*.safetensors")
         "model_hash": "db0197b6919425a5c7102c54e73affc1",
         "model_name": "minimax_h3_dit",
         "model_class": "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT",
     },
+    {
+        # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors")
+        "model_hash": "fb6c399412920c6b817d5c1f43ec62cd",
+        "model_name": "minimax_h3_dit",
+        "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfy",
+        "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_patch_proj", "audio_patch_proj", "condition_proj", "time_embedder.proj_in", "time_embedder.proj_out", "final_layer.adaln_proj.linear", "final_layer.video_out", "final_layer.audio_out", "token_refiner.blocks.0.attn.qkv_proj", "token_refiner.blocks.0.attn.out_proj", "token_refiner.blocks.0.mlp.fc1", "token_refiner.blocks.0.mlp.fc2", "token_refiner.blocks.1.attn.qkv_proj", "token_refiner.blocks.1.attn.out_proj", "token_refiner.blocks.1.mlp.fc1", "token_refiner.blocks.1.mlp.fc2"]},
+    },
+    {
+        # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors")
+        "model_hash": "ac9fdcd56900c9a1a5ef21eaacf6cb76",
+        "model_name": "minimax_h3_dit",
+        "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+        "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["video_patch_proj", "audio_patch_proj", "condition_proj", "adaln_proj.linear", "final_layer.video_out", "final_layer.audio_out", "token_refiner.blocks.0.attn.qkv_proj", "token_refiner.blocks.0.attn.out_proj", "token_refiner.blocks.0.mlp.fc1", "token_refiner.blocks.0.mlp.fc2", "token_refiner.blocks.1.attn.qkv_proj", "token_refiner.blocks.1.attn.out_proj", "token_refiner.blocks.1.mlp.fc1", "token_refiner.blocks.1.mlp.fc2"]},
+    },
+    {
+        # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_pruned_fp8_scaled.safetensors")
+        "model_hash": "7e410be9acd325ef8db3b8bed3d814cd",
+        "model_name": "minimax_h3_dit",
+        "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
+        "quant_config": {"method": "comfy_kitchen_fp8_w8a8", "load_prequantized": True, "exclude_modules": ["video_patch_proj", "audio_patch_proj", "condition_proj", "adaln_proj.linear", "final_layer.video_out", "final_layer.audio_out", "token_refiner.blocks.0.attn.qkv_proj", "token_refiner.blocks.0.attn.out_proj", "token_refiner.blocks.0.mlp.fc1", "token_refiner.blocks.0.mlp.fc2", "token_refiner.blocks.1.attn.qkv_proj", "token_refiner.blocks.1.attn.out_proj", "token_refiner.blocks.1.mlp.fc1", "token_refiner.blocks.1.mlp.fc2"]},
+    },
     {
         # Example: ModelConfig(model_id="Comfy-Org/MiniMax-H3", origin_file_pattern="diffusion_models/minimax_h3_fl2va_pruned_bf16.safetensors")
         "model_hash": "b4e0df87da0c079979ae857edbcef49f",
         "model_name": "minimax_h3_dit",
-        "model_class": "diffsynth.models.minimax_h3_dit_comfy_pruned.MiniMaxH3DiTComfyPruned",
-        "extra_kwargs": {"adaln_curve_grid": 1025, "time_embed_dim": 8},
+        "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
     },
     {
         # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-fl2va-nf4.safetensors")
@@ -1425,16 +1476,14 @@ minimax_h3_series = [
         # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-fl2va-pruned-nf4.safetensors")
         "model_hash": "9b8f1cbc6fb5fcd2e2c02b6691480427",
         "model_name": "minimax_h3_dit",
-        "model_class": "diffsynth.models.minimax_h3_dit_comfy_pruned.MiniMaxH3DiTComfyPruned",
-        "extra_kwargs": {"adaln_curve_grid": 1025, "time_embed_dim": 8},
+        "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
         "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["time_embedder.proj_in", "time_embedder.proj_out", "video_patch_proj", "audio_patch_proj", "condition_proj", "final_layer.video_out", "final_layer.audio_out", "adaln_proj.linear"]},
     },
     {
         # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-ref2va-pruned-nf4.safetensors")
         "model_hash": "072bd13242d6d67df79a154e7c4274b3",
         "model_name": "minimax_h3_dit",
-        "model_class": "diffsynth.models.minimax_h3_dit_comfy_pruned.MiniMaxH3DiTComfyPruned",
-        "extra_kwargs": {"adaln_curve_grid": 1025, "time_embed_dim": 8},
+        "model_class": "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned",
         "quant_config": {"method": "bitsandbytes_nf4", "load_prequantized": True, "exclude_modules": ["time_embedder.proj_in", "time_embedder.proj_out", "video_patch_proj", "audio_patch_proj", "condition_proj", "final_layer.video_out", "final_layer.audio_out", "adaln_proj.linear"]},
     },
     {
@@ -1462,8 +1511,42 @@ minimax_h3_series = [
     },
 ]
 
+minimax_music3_series = [
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors")
+        "model_hash": "8d1f1aba6a5801220ccc3a531865f694",
+        "model_name": "minimax_music3_dit",
+        "model_class": "diffsynth.models.minimax_music3_dit.MiniMaxMusic3DiT",
+    },
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="condition_encoder/diffusion_pytorch_model.safetensors")
+        "model_hash": "1370ef6377206e5ceaa78b125b2d4e1f",
+        "model_name": "minimax_music3_condition_encoder",
+        "model_class": "diffsynth.models.minimax_music3_condition_encoder.MiniMaxMusic3ConditionEncoder",
+    },
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="vocoder/diffusion_pytorch_model.safetensors")
+        "model_hash": "b6eded627c26c85b39804628ba4e5d1d",
+        "model_name": "minimax_music3_vocoder",
+        "model_class": "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3Vocoder",
+    },
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="rvq_depth_decoder/diffusion_pytorch_model.safetensors")
+        "model_hash": "5e571ded532a131c02e9d8256cb3fded",
+        "model_name": "minimax_music3_rvq_depth_decoder",
+        "model_class": "diffsynth.models.minimax_music3_rvq_depth_decoder.MiniMaxMusic3RVQDepthDecoder",
+    },
+    {
+        # Example: ModelConfig(model_id="MiniMax/MiniMax-Music3", origin_file_pattern="language_model/model*.safetensors")
+        "model_hash": "474b2197ab43bd2d4b69746506a952ef",
+        "model_name": "minimax_music3_text_encoder",
+        "model_class": "diffsynth.models.minimax_music3_text_encoder.MiniMaxMusic3TextEncoder",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.minimax_music3_text_encoder.MiniMaxMusic3TextEncoderStateDictConverter",
+    },
+]
+
 MODEL_CONFIGS = (
     stable_diffusion_xl_series + stable_diffusion_series + qwen_image_series + wan_series + flux_series + flux2_series + ernie_image_series
     + z_image_series + ltx2_series + anima_series + mova_series + joyai_image_series + boogu_image_series + ace_step_series + hidream_o1_image_series
-    + image_metrics_series + ideogram4_series + krea2_series + lingbot_video_series + minimax_h3_series
+    + image_metrics_series + ideogram4_series + krea2_series + lingbot_video_series + minimax_h3_series + minimax_music3_series
 )

diffsynth/configs/vram_management_module_maps.py

@@ -14,7 +14,11 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
         "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
-    "diffsynth.models.minimax_h3_dit_comfy_pruned.MiniMaxH3DiTComfyPruned": {
+    "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfy": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.minimax_h3_dit_comfy.MiniMaxH3DiTComfyPruned": {
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
         "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
@@ -470,6 +474,28 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "diffsynth.models.lingbot_video_dit.LingBotVideoRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
     },
+    "diffsynth.models.minimax_music3_dit.MiniMaxMusic3DiT": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_music3_dit.MiniMaxMusic3FourierEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.minimax_music3_rvq_depth_decoder.MiniMaxMusic3RVQDepthDecoder": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_music3_rvq_depth_decoder.MiniMaxMusic3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3Vocoder": {
+        "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3WeightNormConv1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3WeightNormConvTranspose1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.minimax_music3_vocoder.MiniMaxMusic3Snake1d": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.minimax_music3_text_encoder.MiniMaxMusic3TextEncoder": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
 }
 
 def QwenImageTextEncoder_Module_Map_Updater():

diffsynth/core/attention/attention.py

@@ -1,41 +1,52 @@
 import torch, os, inspect
 from einops import rearrange, repeat
 
+if os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID") is not None:
+    try:
+        from kernels import get_kernel
+        kernel_module = get_kernel(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID"), version=int(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_VERSION")))
+        customized_flash_attn_func = kernel_module.flash_attn_func
+        CUSTOMIZED_FA_KERNEL_AVAILABLE = True
+    except Exception:
+        CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+else:
+    CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+
 try:
     from flash_attn.cute import flash_attn_func as flash_attn_func_cute
     FLASH_ATTN_4_AVAILABLE = True
-except ModuleNotFoundError:
+except Exception:
     FLASH_ATTN_4_AVAILABLE = False
 
 try:
     import flash_attn_interface
     FLASH_ATTN_3_AVAILABLE = True
-except ModuleNotFoundError:
+except Exception:
     FLASH_ATTN_3_AVAILABLE = False
 
 try:
     import flash_attn
     FLASH_ATTN_2_AVAILABLE = True
-except ModuleNotFoundError:
+except Exception:
     FLASH_ATTN_2_AVAILABLE = False
 
 try:
     from sageattention import sageattn
     SAGE_ATTN_AVAILABLE = True
-except ModuleNotFoundError:
+except Exception:
     SAGE_ATTN_AVAILABLE = False
 
 try:
     import xformers.ops as xops
     XFORMERS_AVAILABLE = True
-except ModuleNotFoundError:
+except Exception:
     XFORMERS_AVAILABLE = False
 
 try:
     from torch.nn.attention.flex_attention import flex_attention as flex_attention_func
     flex_attention_func = torch.compile(flex_attention_func, dynamic=False, mode="max-autotune-no-cudagraphs", fullgraph=True, backend="inductor")
     FLEX_ATTN_AVAILABLE = True
-except (ModuleNotFoundError, ImportError):
+except Exception:
     FLEX_ATTN_AVAILABLE = False
 
 try:
@@ -43,13 +54,15 @@ try:
         TORCH_SUPPORT_GQA = True
     else:
         TORCH_SUPPORT_GQA = False
-except:
+except Exception:
     TORCH_SUPPORT_GQA = False
 
 
 def initialize_attention_priority():
     if os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION') is not None:
         return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').lower()
+    elif CUSTOMIZED_FA_KERNEL_AVAILABLE:
+        return "customized_fa_kernel"
     elif FLASH_ATTN_4_AVAILABLE:
         return "flash_attention_4"
     elif FLASH_ATTN_3_AVAILABLE:
@@ -178,6 +191,17 @@ def flash_attention_2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_patte
     return out
 
 
+def flash_attention_customized_kernel(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = customized_flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
 def sage_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
     required_in_pattern, required_out_pattern= "b n s d", "b n s d"
     q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
@@ -213,6 +237,8 @@ def attention_forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_patte
             # Sliding Window Attention is not compatible with `is_causal` and `attn_mask`.
             assert is_causal == False and attn_mask is None
             return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+    elif ATTENTION_IMPLEMENTATION == "customized_fa_kernel":
+        return flash_attention_customized_kernel(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
     elif ATTENTION_IMPLEMENTATION == "flash_attention_4":
         return flash_attention_4(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
     elif ATTENTION_IMPLEMENTATION == "flash_attention_3":

diffsynth/core/data/operators.py

@@ -307,8 +307,8 @@ class LoadAudioWithTorchaudio(DataProcessingOperator, FrameSamplerByRateMixin):
                 padding = target_samples - current_samples
                 waveform = torch.nn.functional.pad(waveform, (0, padding))
             return waveform, sample_rate
-        except:
-            warnings.warn(f"Cannot load audio in {data}. The audio will be `None`.")
+        except Exception as e:
+            warnings.warn(f"Cannot load audio in {data} due to {e}. The audio will be `None`.")
             return None
 
 

diffsynth/core/device/npu_compatible_device.py

@@ -83,17 +83,17 @@ def enable_high_precision_for_bf16():
 
 
 def parse_device_type(device):
-    if isinstance(device, str):
-        if device.startswith("cuda"):
-            return "cuda"
-        elif device.startswith("npu"):
+    if isinstance(device, torch.device):
+        return device.type
+    elif isinstance(device, str):
+        # `torch.device("npu")` raises when torch_npu is not installed, so keep npu ahead of the generic parse.
+        if device.startswith("npu"):
             return "npu"
-        elif device.startswith("mps"):
-            return "mps"
-        else:
+        try:
+            return torch.device(device).type
+        except RuntimeError:
+            # Strings that are not device names at all
             return "cpu"
-    elif isinstance(device, torch.device):
-        return device.type
 
 
 def parse_nccl_backend(device_type):

diffsynth/core/loader/file.py

@@ -1,5 +1,5 @@
 from safetensors import safe_open
-import torch, hashlib
+import torch, hashlib, json
 
 
 def load_state_dict(file_path, torch_dtype=None, device="cpu", pin_memory=False, verbose=0):
@@ -34,10 +34,26 @@ def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"):
 
 
 def load_metadata_from_safetensors(file_path):
+    if isinstance(file_path, list):
+        metadata = {}
+        for file_path_ in file_path:
+            for key, value in load_metadata_from_safetensors(file_path_).items():
+                metadata[key] = merge_metadata_value(metadata[key], value) if key in metadata else value
+        return metadata
     with safe_open(file_path, framework="pt", device="cpu") as f:
         return f.metadata() or {}
 
 
+def merge_metadata_value(old_value, new_value):
+    try:
+        old_items, new_items = json.loads(old_value), json.loads(new_value)
+    except (json.JSONDecodeError, TypeError):
+        return new_value
+    if isinstance(old_items, list) and isinstance(new_items, list):
+        return json.dumps(old_items + [i for i in new_items if i not in old_items])
+    return new_value
+
+
 def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"):
     state_dict = torch.load(file_path, map_location=device, weights_only=True)
     if len(state_dict) == 1:

diffsynth/core/loader/model.py

@@ -37,14 +37,13 @@ def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, devic
             devices = [vram_config[k] for k in ("offload_device", "onload_device", "preparing_device", "computation_device")]
             load_device = [d for d in devices if d != "disk"][0]
             disk_map = DiskMap(path, load_device, torch_dtype=None, state_dict_converter=state_dict_converter)
-            metadata = load_metadata_from_safetensors(path[0] if isinstance(path, list) else path)
+            metadata = load_metadata_from_safetensors(path)
             model = quantize.prepare_for_prequantized_load(model, compute_dtype=vram_config["computation_dtype"])
             model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=vram_limit, quantize=quantize, metadata=metadata)
         else:
             offload_device = vram_config["offload_device"]
             computation_device = vram_config["computation_device"]
             computation_dtype = vram_config["computation_dtype"]
-            offload_dtype = vram_config["offload_dtype"]
             load_dtype = None if quantize.load_prequantized else computation_dtype
             if state_dict is None: state_dict = DiskMap(path, offload_device, torch_dtype=load_dtype)
             if state_dict_converter is not None:
@@ -61,7 +60,7 @@ def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, devic
 
             model = quantize.quantize_model(model, compute_device=computation_device, model_device=offload_device)
             model = quantize.dequantize_model(model, compute_dtype=computation_dtype, compute_device=computation_device, model_device=offload_device)
-            model = model.to(dtype=offload_dtype, device=offload_device)
+            model = model.to(dtype=computation_dtype, device=offload_device)
             model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=None, vram_limit=vram_limit, quantize=quantize)
     elif quantize is not None:
         # Weight-only quantization (see `diffsynth.core.quant`), isolated from the normal path below.

diffsynth/core/quant/__init__.py

@@ -1,3 +1,3 @@
-from .base import QuantBackend, QUANT_BACKENDS, register_quant_backend, check_differentiable, check_backend_contract, resolve_checkpoint_keys
+from .base import QuantBackend, BackendConfig, QUANT_BACKENDS, register_quant_backend, check_differentiable, check_backend_contract
 from .config import QuantizeConfig, MixedQuantizeConfig, QuantMethodSpec, QUANT_METHODS, register_quant_method, describe_quant_method
 from . import backends

diffsynth/core/quant/backends/__init__.py

@@ -1 +1,28 @@
-from . import bitsandbytes, torchao
+import importlib
+
+_LAZY_BACKENDS = {
+    "bitsandbytes": ".bitsandbytes",
+    "torchao": ".torchao",
+    "comfy_kitchen": ".comfy_kitchen",
+}
+_loaded = set()
+
+
+def load_backend(name):
+    module = _LAZY_BACKENDS.get(name)
+    if module is not None and name not in _loaded:
+        importlib.import_module(module, __name__)
+        _loaded.add(name)
+
+
+def load_all_backends():
+    for name in _LAZY_BACKENDS:
+        load_backend(name)
+
+
+def load_backend_for_method(method):
+    from ..config import QUANT_METHODS
+    for name in _LAZY_BACKENDS:
+        load_backend(name)
+        if method in QUANT_METHODS:
+            return

diffsynth/core/quant/backends/bitsandbytes.py

@@ -1,7 +1,15 @@
+from dataclasses import dataclass, field
+
 import torch
-from ..base import QuantBackend, register_quant_backend
+from ..base import QuantBackend, BackendConfig, register_quant_backend
 from ..config import register_quant_method
 
+try:
+    import bitsandbytes as bnb
+    BITSANDBYTES_AVAILABLE = True
+except ImportError:
+    BITSANDBYTES_AVAILABLE = False
+
 
 def _assign_params_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
     """`_load_from_state_dict` replacement for shells: assigns tensors as-is, since packed weights do not pass the standard shape check."""
@@ -29,10 +37,10 @@ def _assign_params_from_state_dict(self, state_dict, prefix, local_metadata, str
 class BitsAndBytesQuantBackend(QuantBackend):
     """Adapter over `bitsandbytes.nn.Linear4bit` (NF4 / FP4, weight-only)."""
 
+    project_url = "https://github.com/bitsandbytes-foundation/bitsandbytes"
+
     def validate_environment(self):
-        try:
-            import bitsandbytes  # noqa: F401
-        except ImportError:
+        if not BITSANDBYTES_AVAILABLE:
             raise ImportError(
                 "bitsandbytes is required for bitsandbytes quantization methods. "
                 'Please install it via `pip install bitsandbytes` or `pip install "diffsynth[quant]"`.'
@@ -47,23 +55,14 @@ class BitsAndBytesQuantBackend(QuantBackend):
         }
 
     def quantized_linear_classes(self):
-        try:
-            import bitsandbytes as bnb
-        except ImportError:
-            return ()
         return (bnb.nn.Linear4bit,)
 
     def create_quantized_linear(self, linear, compute_device=None, model_device=None):
         """The `meta` shell avoids allocating an fp weight; bnb quantizes while `Params4bit` moves to `compute_device`."""
-        import bitsandbytes as bnb
-
+        cfg = self.config
         compute_dtype = linear.weight.dtype
         if compute_device is None:
             compute_device = linear.weight.device
-        quant_type = self.config.get("quant_type", "nf4")
-        compress_statistics = self.config.get("compress_statistics", True)
-        blocksize = self.config.get("blocksize")
-        quant_storage = self.config.get("quant_storage", torch.uint8)
 
         with torch.device("meta"):
             quant_linear = bnb.nn.Linear4bit(
@@ -71,17 +70,17 @@ class BitsAndBytesQuantBackend(QuantBackend):
                 linear.out_features,
                 bias=linear.bias is not None,
                 compute_dtype=compute_dtype,
-                compress_statistics=compress_statistics,
-                quant_type=quant_type,
-                quant_storage=quant_storage,
+                compress_statistics=cfg.compress_statistics,
+                quant_type=cfg.quant_type,
+                quant_storage=cfg.quant_storage,
             )
         weight = bnb.nn.Params4bit(
             linear.weight.data.contiguous(),
             requires_grad=False,
-            compress_statistics=compress_statistics,
-            blocksize=blocksize,
-            quant_type=quant_type,
-            quant_storage=quant_storage,
+            compress_statistics=cfg.compress_statistics,
+            blocksize=cfg.blocksize,
+            quant_type=cfg.quant_type,
+            quant_storage=cfg.quant_storage,
         )
         quant_linear.weight = weight.to(compute_device)
         if linear.bias is not None:
@@ -89,22 +88,20 @@ class BitsAndBytesQuantBackend(QuantBackend):
         return quant_linear.to(device=model_device)
 
     def create_quantized_linear_shell(self, linear, compute_dtype):
-        import bitsandbytes as bnb
         with torch.device("meta"):
             shell = bnb.nn.Linear4bit(
                 linear.in_features,
                 linear.out_features,
                 bias=linear.bias is not None,
                 compute_dtype=compute_dtype,
-                compress_statistics=self.config.get("compress_statistics", True),
-                quant_type=self.config.get("quant_type", "nf4"),
-                quant_storage=self.config.get("quant_storage", torch.uint8),
+                compress_statistics=self.config.compress_statistics,
+                quant_type=self.config.quant_type,
+                quant_storage=self.config.quant_storage,
             )
         shell._load_from_state_dict = _assign_params_from_state_dict.__get__(shell)
         return shell
 
     def unflatten_state_dict(self, state_dict, metadata):
-        import bitsandbytes as bnb
         rebuilt = {}
         quant_state = {}
         for key, value in state_dict.items():
@@ -122,7 +119,6 @@ class BitsAndBytesQuantBackend(QuantBackend):
         return rebuilt
 
     def dequantize_to_linear(self, module, compute_dtype, compute_device=None, model_device=None):
-        import bitsandbytes as bnb
         if compute_device is not None:
             module = module.to(device=compute_device)
         weight = module.weight
@@ -136,17 +132,24 @@ class BitsAndBytesQuantBackend(QuantBackend):
         return linear
 
 
-def _linear4bit_config(quant_type):
-    def factory(backend_config_kwargs):
-        return {
-            "quant_type": quant_type,
-            "compress_statistics": True,
-            "blocksize": None,
-            "quant_storage": torch.uint8,
-            **backend_config_kwargs,
-        }
-    return factory
+@dataclass
+class BitsAndBytes4bitConfig(BackendConfig):
+    """Shared 4bit knobs; `quant_type` is pinned by the per-method subclass."""
+
+    compress_statistics: bool = True
+    blocksize: int = None                                # None -> bnb default (64)
+    quant_storage: torch.dtype = torch.uint8
+
+
+@dataclass
+class BitsAndBytesNF4Config(BitsAndBytes4bitConfig):
+    quant_type: str = field(init=False, default="nf4")
+
+
+@dataclass
+class BitsAndBytesFP4Config(BitsAndBytes4bitConfig):
+    quant_type: str = field(init=False, default="fp4")
 
 
-register_quant_method("bitsandbytes_nf4", "bitsandbytes", _linear4bit_config("nf4"), label="4bit, nf4, weight-only")
-register_quant_method("bitsandbytes_fp4", "bitsandbytes", _linear4bit_config("fp4"), label="4bit, fp4, weight-only")
+register_quant_method("bitsandbytes_nf4", "bitsandbytes", BitsAndBytesNF4Config.from_kwargs, label="4bit, nf4, weight-only")
+register_quant_method("bitsandbytes_fp4", "bitsandbytes", BitsAndBytesFP4Config.from_kwargs, label="4bit, fp4, weight-only")