diffsynth 2.1.1__tar.gz → 2.1.2__tar.gz

{diffsynth-2.1.1 → diffsynth-2.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.1.1
+Version: 2.1.2
 Summary: Enjoy the magic of Diffusion models!
 Author: ModelScope Team
 License: Apache-2.0

{diffsynth-2.1.1 → diffsynth-2.1.2}/README.md

@@ -1527,6 +1527,8 @@ Example code for LingBot-Video is available at: [/examples/lingbot_video/](/exam
 |[Robbyant/lingbot-video-moe-30b-a3b: T2V](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b)|[code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2v.py)|[code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2v.py)|[code](/examples/lingbot_video/model_training/full/lingbot-video-moe-30b-a3b_t2v.sh)|[code](/examples/lingbot_video/model_training/validate_full/lingbot-video-moe-30b-a3b_t2v.py)|[code](/examples/lingbot_video/model_training/lora/lingbot-video-moe-30b-a3b_t2v.sh)|[code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-moe-30b-a3b_t2v.py)|
 |[Robbyant/lingbot-video-moe-30b-a3b: TI2V](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b)|[code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_ti2v.py)|[code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_ti2v.py)|[code](/examples/lingbot_video/model_training/full/lingbot-video-moe-30b-a3b_ti2v.sh)|[code](/examples/lingbot_video/model_training/validate_full/lingbot-video-moe-30b-a3b_ti2v.py)|[code](/examples/lingbot_video/model_training/lora/lingbot-video-moe-30b-a3b_ti2v.sh)|[code](/examples/lingbot_video/model_training/validate_lora/lingbot-video-moe-30b-a3b_ti2v.py)|
 |[Robbyant/lingbot-video-moe-30b-a3b: T2I](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b)|[code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2i.py)|[code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2i.py)|-|-|-|-|
+|[Robbyant/lingbot-video-moe-30b-a3b: T2V + Refinement](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b)|[code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2v_refiner.py)|[code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_t2v_refiner.py)|-|-|-|-|
+|[Robbyant/lingbot-video-moe-30b-a3b: TI2V + Refinement](https://modelscope.cn/models/Robbyant/lingbot-video-moe-30b-a3b)|[code](/examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_ti2v_refiner.py)|[code](/examples/lingbot_video/model_inference_low_vram/lingbot-video-moe-30b-a3b_ti2v_refiner.py)|-|-|-|-|
 
 </details>
 
@@ -1593,6 +1595,12 @@ Example code for MiniMax-H3 is available at: [/examples/minimax_h3/](/examples/m
 |[MiniMax/MiniMax-H3: Retake](https://www.modelscope.cn/models/MiniMax/MiniMax-H3)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-Retake.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Retake.py)|-|-|-|-|
 |[DiffSynth-Studio/MiniMax-H3-NF4: FL2VA](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-FL2VA.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-FL2VA.py)|-|-|[code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-FL2VA.sh)|[code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-FL2VA.py)|
 |[DiffSynth-Studio/MiniMax-H3-NF4: Ref2VA](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-NF4-Ref2VA.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-NF4-Ref2VA.py)|-|-|[code](/examples/minimax_h3/model_training/lora/MiniMax-H3-NF4-Ref2VA.sh)|[code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-NF4-Ref2VA.py)|
+|[Comfy-Org/MiniMax-H3: FL2VA pruned](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-FL2VA.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-FL2VA.py)|[code](/examples/minimax_h3/model_training/full/MiniMax-H3-Pruned-FL2VA.sh)|[code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Pruned-FL2VA.py)|[code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-FL2VA.sh)|[code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-FL2VA.py)|
+|[Comfy-Org/MiniMax-H3: Ref2VA pruned](https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-Ref2VA.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-Ref2VA.py)|[code](/examples/minimax_h3/model_training/full/MiniMax-H3-Pruned-Ref2VA.sh)|[code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Pruned-Ref2VA.py)|[code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-Ref2VA.sh)|[code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-Ref2VA.py)|
+|[DiffSynth-Studio/MiniMax-H3-NF4: FL2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-NF4-FL2VA.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-NF4-FL2VA.py)|-|-|[code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-NF4-FL2VA.sh)|[code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-NF4-FL2VA.py)|
+|[DiffSynth-Studio/MiniMax-H3-NF4: Ref2VA pruned](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-NF4)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-Pruned-NF4-Ref2VA.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Pruned-NF4-Ref2VA.py)|-|-|[code](/examples/minimax_h3/model_training/lora/MiniMax-H3-Pruned-NF4-Ref2VA.sh)|[code](/examples/minimax_h3/model_training/validate_lora/MiniMax-H3-Pruned-NF4-Ref2VA.py)|
+|[lightx2v/Minimax-h3-Turbo: FL2VA 4steps](https://www.modelscope.cn/models/lightx2v/Minimax-h3-Turbo)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/minimax_h3/model_inference/MiniMax-H3-FL2VA-Turbo.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-FL2VA-Turbo.py)|-|-|-|-|
+|[DiffSynth-Studio/MiniMax-H3-Text-Embeddings](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-Text-Embeddings)|[code](/examples/minimax_h3/model_inference/MiniMax-H3-Text-Embeddings.py)|[code](/examples/minimax_h3/model_inference_low_vram/MiniMax-H3-Text-Embeddings.py)|[code](/examples/minimax_h3/model_training/full/MiniMax-H3-Text-Embeddings.sh)|[code](/examples/minimax_h3/model_training/validate_full/MiniMax-H3-Text-Embeddings.py)|-|-|
 
 </details>
 

{diffsynth-2.1.1 → diffsynth-2.1.2}/diffsynth/configs/model_configs.py

@@ -1400,6 +1400,13 @@ minimax_h3_series = [
         "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_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},
+    },
     {
         # Example: ModelConfig(model_id="DiffSynth-Studio/MiniMax-H3-NF4", origin_file_pattern="minimax-h3-fl2va-nf4.safetensors")
         "model_hash": "4b27efacefbc4d8670e0d7b876699648",
@@ -1414,6 +1421,22 @@ minimax_h3_series = [
         "model_class": "diffsynth.models.minimax_h3_dit.MiniMaxH3DiT",
         "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"]},
     },
+    {
+        # 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},
+        "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},
+        "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="MiniMax/MiniMax-H3-FL2VA", origin_file_pattern="video_vae/source/model.safetensors")
         "model_hash": "24b80900992e2024fab17c991c57da23",

{diffsynth-2.1.1 → diffsynth-2.1.2}/diffsynth/configs/vram_management_module_maps.py

@@ -14,6 +14,10 @@ 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": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
     "diffsynth.models.minimax_h3_text_encoder.MiniMaxH3TextEncoder": {
         "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
         "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",

{diffsynth-2.1.1 → diffsynth-2.1.2}/diffsynth/core/quant/__init__.py

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

{diffsynth-2.1.1 → diffsynth-2.1.2}/diffsynth/core/quant/backends/torchao.py

@@ -1,10 +1,15 @@
 import importlib.util
+import json
 
 import torch
 from ..base import QuantBackend, register_quant_backend
 from ..config import register_quant_method
 
 
+class TorchaoLinear(torch.nn.Linear):
+    """Marker class for torchao-quantized Linears."""
+
+
 @register_quant_backend("torchao")
 class TorchaoQuantBackend(QuantBackend):
     """Adapter over torchao `quantize_` (weight-only configs); the quantization lives in the weight tensor subclass, not the module class."""
@@ -21,27 +26,32 @@ class TorchaoQuantBackend(QuantBackend):
     def capabilities(self):
         return {
             "is_serializable": True,
-            "is_differentiable": False,
+            "is_differentiable": True,
             "is_compileable": True,
             "requires_calibration": False,
         }
 
-    def is_quantized_linear(self, module) -> bool:
-        weight = getattr(module, "weight", None)
-        return isinstance(module, torch.nn.Linear) and weight is not None and "torchao" in type(weight).__module__
+    def quantized_linear_classes(self):
+        return (TorchaoLinear,)
+
+    def checkpoint_key_patterns(self):
+        return ("weight", "_weight_qdata", "_weight_scale", "_weight_zero_point", "bias")
 
     def create_quantized_linear(self, linear, compute_device=None, model_device=None):
         from torchao.quantization import quantize_
         linear.requires_grad_(False)
         if compute_device is not None:
             linear = linear.to(device=compute_device)
-        quantize_(linear, self.config)
+        quant_linear = TorchaoLinear(linear.in_features, linear.out_features, bias=linear.bias is not None, device="meta")
+        quant_linear.weight = linear.weight
+        quant_linear.bias = linear.bias
+        quantize_(quant_linear, self.config)
         if model_device is not None:
-            linear = linear.to(device=model_device)
-        return linear
+            quant_linear = quant_linear.to(device=model_device)
+        return quant_linear
 
     def create_quantized_linear_shell(self, linear, compute_dtype):
-        return torch.nn.Linear(linear.in_features, linear.out_features, bias=linear.bias is not None, device="meta")
+        return TorchaoLinear(linear.in_features, linear.out_features, bias=linear.bias is not None, device="meta")
 
     def flatten_state_dict(self, state_dict):
         self._require_safetensors_support()
@@ -61,9 +71,16 @@ class TorchaoQuantBackend(QuantBackend):
                 "safetensors header), so its tensor subclasses cannot be rebuilt. It was most "
                 "likely not saved by torchao."
             )
+        tensor_names = json.loads(metadata["tensor_names"])
+        root_names = [name for name in tensor_names if "." not in name]
+        if root_names:
+            metadata = {**metadata, "tensor_names": json.dumps([name for name in tensor_names if "." in name])}
         rebuilt = unflatten_tensor_state_dict(state_dict, metadata)
         if isinstance(rebuilt, tuple):
             rebuilt = rebuilt[0]
+        for name in root_names:
+            if name in state_dict:
+                rebuilt[name] = state_dict[name]
         return rebuilt
 
     def _require_safetensors_support(self):

diffsynth-2.1.2/diffsynth/core/quant/base.py

@@ -0,0 +1,241 @@
+from abc import ABC
+import torch
+
+
+QUANT_BACKENDS = {}
+
+
+class QuantBackend(ABC):
+    """
+    Adapter between the framework and a quantization library (bnb / torchao / custom).
+    Subclasses are registered in `QUANT_BACKENDS` ({name -> class}) and instantiated by
+    `QuantizeConfig` with the method's backend config. A backend operates on single
+    Linears; model-level traversal and replacement is done by the `QuantizeConfig` methods.
+
+    The quantized Linear produced by a backend must satisfy:
+    (a) It is an `nn.Linear` drop-in: `forward(x)` internally performs dequant + matmul.
+    (b) `.to(...)` moves devices but never re-types the packed weight / quant state:
+        a dtype cast (`.to(dtype)` / `.half()` / `.float()`) must leave their storage
+        format and values intact.
+    (c) `state_dict()` / `load_state_dict(assign=True)` round-trips (optionally via
+        `flatten_state_dict` / `unflatten_state_dict`).
+    (d) (Training branch only) `forward` is differentiable w.r.t. its input, so
+        gradients can pass through frozen quantized layers to reach LoRA branches.
+        Declared statically by `capabilities()["is_differentiable"]` and verifiable
+        at runtime by calling `check_differentiable` on a quantized Linear the
+        backend produced.
+    """
+
+    name: str = ""
+
+    def __init__(self, config=None):
+        self.config = config
+
+    def capabilities(self) -> dict:
+        return {
+            "is_serializable": False,
+            "is_differentiable": False,
+            "is_compileable": False,
+            "requires_calibration": False,
+        }
+
+    def validate_environment(self):
+        return
+
+    def create_quantized_linear(self, linear: torch.nn.Linear, compute_device=None, model_device=None) -> torch.nn.Module:
+        raise NotImplementedError(
+            f"Backend `{self.name}` cannot quantize an fp model online. Use a method whose "
+            "backend supports it, or load an already quantized checkpoint."
+        )
+
+    def create_quantized_linear_shell(self, linear: torch.nn.Linear, compute_dtype: torch.dtype) -> torch.nn.Module:
+        raise NotImplementedError(
+            f"Backend `{self.name}` cannot load pre-quantized checkpoints. Use "
+            "`load_prequantized=False` to quantize an fp model online instead."
+        )
+
+    def dequantize_to_linear(self, module: torch.nn.Module, compute_dtype: torch.dtype, compute_device=None, model_device=None) -> torch.nn.Linear:
+        raise NotImplementedError(
+            f"Backend `{self.name}` cannot dequantize back to `nn.Linear`, so "
+            '`mode="dequant_once"` is unavailable.'
+        )
+
+    def quantized_linear_classes(self) -> tuple:
+        """The Linear classes this backend produces. MUST be subclass of `torch.nn.Linear`."""
+        raise NotImplementedError(
+            f"Backend `{self.name}` must declare the Linear classes it produces."
+        )
+
+    def checkpoint_key_patterns(self) -> tuple:
+        """Checkpoint entries one quantized Linear needs, relative to its dotted name."""
+        return ("weight", "weight.", "bias")
+
+    def is_quantized_linear(self, module) -> bool:
+        return isinstance(module, self.quantized_linear_classes())
+
+    def flatten_state_dict(self, state_dict: dict):
+        self._require_serializable()
+        return state_dict, {}
+
+    def unflatten_state_dict(self, state_dict: dict, metadata: dict):
+        self._require_serializable()
+        return state_dict
+
+    def _require_serializable(self):
+        if not self.capabilities().get("is_serializable", False):
+            raise NotImplementedError(
+                f"Backend `{self.name}` declares `is_serializable=False`, so its quantized state "
+                "dict cannot be flattened or rebuilt. Override these methods if it actually can."
+            )
+
+def register_quant_backend(name):
+    def decorator(cls):
+        cls.name = name
+        QUANT_BACKENDS[name] = cls
+        return cls
+    return decorator
+
+
+def resolve_checkpoint_keys(patterns, layer_name: str, available_keys) -> list:
+    """Expand a backend's `checkpoint_key_patterns` into the absolute keys of one layer.
+
+    `available_keys` is anything supporting `in` and iteration (a dict or a `DiskMap`).
+    Exact patterns are probed with `in`, so a whole-file key index costs O(1) per pattern;
+    only patterns ending in "." require scanning the index.
+    """
+    prefix = f"{layer_name}." if layer_name else ""
+    nested_prefixes = tuple(prefix + pattern for pattern in patterns if pattern.endswith("."))
+    keys = [prefix + pattern for pattern in patterns
+            if not pattern.endswith(".") and prefix + pattern in available_keys]
+    if nested_prefixes:
+        keys += [key for key in available_keys if key.startswith(nested_prefixes)]
+    return list(dict.fromkeys(keys))
+
+
+def check_differentiable(module: torch.nn.Module, example_input: torch.Tensor = None, verbose: bool = True) -> bool:
+    """
+    Check whether gradients pass through `module` w.r.t. its input: run a real
+    backward pass from the output (`torch.autograd.grad`) and verify a finite
+    gradient arrives at the input. This is what LoRA training requires from
+    frozen (e.g. quantized) layers. The module is cast to bfloat16 in place and
+    probed with a bfloat16 input; if `example_input` is None, a random one is
+    built for modules exposing `in_features`.
+
+    Example (probing a torchao-quantized Linear):
+
+        import torch
+        from diffsynth.core.quant import check_differentiable
+        from torchao.quantization import quantize_, Int8WeightOnlyConfig
+
+        linear = torch.nn.Linear(1024, 1024, dtype=torch.bfloat16, device="cuda")
+        quantize_(linear, Int8WeightOnlyConfig(version=2))
+        check_differentiable(linear)
+    """
+    def report(result, detail):
+        if verbose:
+            print(f"check_differentiable ({type(module).__name__}): {'OK' if result else 'FAIL'} -- {detail}")
+        return result
+
+    try:
+        module = module.to(torch.bfloat16)
+        if example_input is None:
+            if not hasattr(module, "in_features"):
+                raise ValueError("`example_input` is required for modules without `in_features`.")
+            device = next((t.device for t in list(module.parameters()) + list(module.buffers())), torch.device("cpu"))
+            example_input = torch.randn(4, module.in_features, device=device)
+        x = example_input.detach().to(torch.bfloat16).requires_grad_(True)
+        y = module(x)
+        if not y.requires_grad:
+            return report(False, "the output does not require grad, so no autograd graph was recorded")
+        input_grad = torch.autograd.grad(y, x, grad_outputs=torch.randn_like(y), allow_unused=True)[0]
+    except Exception as error:
+        return report(False, f"{type(error).__name__}: {error}")
+    if input_grad is None:
+        return report(False, "backward finished but no gradient reached the input")
+    if not torch.isfinite(input_grad.float()).all():
+        return report(False, "the input gradient contains non-finite values")
+    return report(True, "gradients pass through the module to its input")
+
+
+def check_backend_contract(backend, in_features: int = 512, out_features: int = 512,
+                           compute_dtype: torch.dtype = torch.bfloat16,
+                           compute_device: str = "cuda", verbose: bool = True) -> bool:
+    """
+    Verify a backend satisfies the quantized-Linear contract: it declares its classes,
+    both factory methods return instances of them, and every declared class subclasses
+    `torch.nn.Linear` so LoRA target detection and VRAM management can see it. The
+    checkpoint key patterns are checked against the keys the backend actually writes,
+    since a pattern list that misses a scale makes disk offload load corrupt layers
+    without raising.
+
+    Example:
+
+        from diffsynth.core.quant import QUANT_BACKENDS, QUANT_METHODS, check_backend_contract
+        spec = QUANT_METHODS["bitsandbytes_nf4"]
+        check_backend_contract(QUANT_BACKENDS[spec.backend](spec.config_factory({})))
+    """
+    failures = []
+
+    def check(condition, detail):
+        if not condition:
+            failures.append(detail)
+        if verbose:
+            print(f"  [{'PASS' if condition else 'FAIL'}] {detail}")
+        return condition
+
+    if verbose:
+        print(f"check_backend_contract ({backend.name}):")
+    try:
+        classes = backend.quantized_linear_classes()
+    except NotImplementedError as error:
+        if verbose:
+            print(f"  [FAIL] quantized_linear_classes() is not implemented: {error}")
+        return False
+    if not check(len(classes) > 0, f"quantized_linear_classes() is non-empty: {[cls.__name__ for cls in classes]}"):
+        return False
+    for cls in classes:
+        check(issubclass(cls, torch.nn.Linear), f"{cls.__name__} subclasses torch.nn.Linear")
+
+    plain = torch.nn.Linear(in_features, out_features, bias=True, dtype=compute_dtype, device=compute_device)
+    plain.requires_grad_(False)
+    check(not backend.is_quantized_linear(plain), "a plain nn.Linear is not reported as quantized")
+
+    try:
+        shell = backend.create_quantized_linear_shell(plain, compute_dtype)
+    except NotImplementedError:
+        shell = None
+        if verbose:
+            print("  [SKIP] create_quantized_linear_shell() is unsupported by this backend")
+    if shell is not None:
+        check(isinstance(shell, classes), f"create_quantized_linear_shell() returns a declared class, got {type(shell).__name__}")
+        check(backend.is_quantized_linear(shell), "the shell is recognized before load_state_dict (disk offload routing)")
+
+    try:
+        quantized = backend.create_quantized_linear(plain, compute_device=compute_device)
+    except NotImplementedError:
+        quantized = None
+        if verbose:
+            print("  [SKIP] create_quantized_linear() is unsupported by this backend")
+    if quantized is not None:
+        check(isinstance(quantized, classes), f"create_quantized_linear() returns a declared class, got {type(quantized).__name__}")
+
+    patterns = backend.checkpoint_key_patterns()
+    check(len(patterns) > 0, f"checkpoint_key_patterns() is non-empty: {list(patterns)}")
+    saved = quantized if quantized is not None else shell
+    if saved is None:
+        if verbose:
+            print("  [SKIP] neither factory method is supported, so the stored keys cannot be checked")
+    else:
+        state_dict = {f"proj.{key}": value for key, value in saved.state_dict().items()}
+        if backend.capabilities().get("is_serializable", False):
+            try:
+                state_dict = backend.flatten_state_dict(state_dict)[0]
+            except Exception as error:
+                if verbose:
+                    print(f"  [SKIP] flatten_state_dict() failed, falling back to the raw state dict keys: {type(error).__name__}: {error}")
+        uncovered = sorted(set(state_dict) - set(resolve_checkpoint_keys(patterns, "proj", state_dict)))
+        check(len(uncovered) == 0, f"checkpoint_key_patterns() covers every stored key; uncovered: {uncovered}")
+
+    if verbose:
+        print(f"  => {'OK' if not failures else str(len(failures)) + ' FAILED'}")
+    return not failures

{diffsynth-2.1.1 → diffsynth-2.1.2}/diffsynth/core/quant/config.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field, fields, is_dataclass
 from typing import Any, Callable, Optional
 import torch
-from .base import QUANT_BACKENDS
+from .base import QUANT_BACKENDS, resolve_checkpoint_keys
 
 
 QUANT_METHODS = {}
@@ -166,6 +166,40 @@ class QuantizeConfig:
         """Whether `module` is one of this config's backend-native quantized Linears."""
         return self.backend.is_quantized_linear(module)
 
+    def checkpoint_keys(self, module, layer_name: str, available_keys) -> list:
+        """
+        Resolve the backend's `checkpoint_key_patterns` for `layer_name` against
+        `available_keys` (anything supporting `in` and iteration, including a `DiskMap`).
+        Raises when the layer contributes no key at all, since silently loading a
+        quantized layer without its packed weight or scale corrupts it without any error.
+
+        Parameters:
+            module: the quantized layer the keys are fetched for.
+            layer_name: its dotted name inside the checkpoint.
+            available_keys: the key index of the whole checkpoint.
+        """
+        keys = resolve_checkpoint_keys(self.backend.checkpoint_key_patterns(), layer_name, available_keys)
+        if not keys:
+            raise ValueError(
+                f"Found no checkpoint entry for the quantized layer `{layer_name}` "
+                f"(backend `{self.backend.name}`, patterns {list(self.backend.checkpoint_key_patterns())}). "
+                "Check that the checkpoint really holds this layer quantized and that "
+                "`target_modules` matches the layers it quantized."
+            )
+        return keys
+
+    def build_quantized_shell(self, module, compute_dtype: torch.dtype, **kwargs):
+        """
+        Build an empty quantized Linear matching `module`, used to release a layer's
+        weights while keeping it routable, and to stage a transient copy on the
+        computation device.
+
+        Parameters:
+            module: the quantized layer whose shape and bias presence are mirrored.
+            compute_dtype: dtype the shell dequantizes to at forward time.
+        """
+        return self.backend.create_quantized_linear_shell(module, compute_dtype)
+
     def prepare_for_prequantized_load(self, model: torch.nn.Module, compute_dtype: torch.dtype = torch.bfloat16):
         """
         Replace the targeted `nn.Linear` layers with empty quantized Linears matching a
@@ -226,6 +260,8 @@ class QuantizeConfig:
     def _should_quantize(self, full_name, module):
         if not isinstance(module, torch.nn.Linear):
             return False
+        if self.is_quantized_linear(module):
+            return False
         if self.target_modules is not None and not self._name_matches(full_name, self.target_modules):
             return False
         if self._name_matches(full_name, self.exclude_modules):
@@ -302,6 +338,7 @@ class MixedQuantizeConfig:
         modes = {config.mode for config in self.configs}
         if len(modes) > 1:
             raise ValueError(f"All configs in `MixedQuantizeConfig` should share the same `mode`, but got {sorted(modes)}.")
+        self._ownership = {}
 
     @property
     def method(self):
@@ -324,7 +361,7 @@ class MixedQuantizeConfig:
         """
         if self.load_prequantized:
             return model
-        self._validate_disjoint(model)
+        self._build_ownership(model)
         for config in self.configs:
             config.quantize_model(model, compute_device=compute_device, model_device=model_device)
         return model
@@ -351,6 +388,43 @@ class MixedQuantizeConfig:
         """Whether `module` is a quantized Linear of any config's backend."""
         return any(config.is_quantized_linear(module) for config in self.configs)
 
+    def checkpoint_keys(self, module, layer_name: str, available_keys) -> list:
+        """
+        Resolve the checkpoint keys of `layer_name` using the config whose backend owns
+        that layer, so each layer set is read with its own key shape.
+
+        Parameters:
+            module: the quantized layer the keys are fetched for.
+            layer_name: its dotted name inside the checkpoint.
+            available_keys: the key index of the whole checkpoint.
+        """
+        return self._owning_config(module, layer_name).checkpoint_keys(module, layer_name, available_keys)
+
+    def build_quantized_shell(self, module, compute_dtype: torch.dtype, layer_name: str = None):
+        """
+        Build an empty quantized Linear for `module` using the config that owns it.
+
+        Parameters:
+            module: the quantized layer whose shape and bias presence are mirrored.
+            compute_dtype: dtype the shell dequantizes to at forward time.
+            layer_name: its dotted name inside the model; required to pick the right config
+                when several configs share one backend, since their quantized Linears are
+                then the same class.
+        """
+        return self._owning_config(module, layer_name).build_quantized_shell(module, compute_dtype)
+
+    def _owning_config(self, module, layer_name=None):
+        config = self._ownership.get(layer_name)
+        if config is not None:
+            return config
+        for config in self.configs:
+            if config.is_quantized_linear(module):
+                return config
+        raise ValueError(
+            f"`{type(module).__name__}` is not a quantized Linear of any config in this "
+            f"`MixedQuantizeConfig` (methods: {self.method})."
+        )
+
     def prepare_for_prequantized_load(self, model: torch.nn.Module, compute_dtype: torch.dtype = torch.bfloat16):
         """
         Replace each config's targeted Linears by that backend's shells, after verifying
@@ -360,26 +434,29 @@ class MixedQuantizeConfig:
             model: the freshly constructed model whose targeted layers become shells.
             compute_dtype: dtype the quantized layers dequantize to at forward time.
         """
-        self._validate_disjoint(model)
+        self._build_ownership(model)
         for config in self.configs:
             config.prepare_for_prequantized_load(model, compute_dtype=compute_dtype)
         return model
 
     def unflatten_state_dict(self, state_dict: dict, metadata: dict):
         """
-        Rebuild composite quantized tensors from a flat (safetensors) state dict,
-        running each distinct backend's `unflatten_state_dict` once.
+        Rebuild composite quantized tensors from a flat (safetensors) state dict.
         """
-        for config in self._distinct_backend_configs():
-            state_dict = config.unflatten_state_dict(state_dict, metadata)
-        return state_dict
+        if not self._ownership:
+            for config in self._distinct_backend_configs():
+                state_dict = config.unflatten_state_dict(state_dict, metadata)
+            return state_dict
+        rebuilt = {}
+        for config, keys in self._grouped_by_backend(state_dict):
+            rebuilt.update(config.unflatten_state_dict(keys, metadata) if config else keys)
+        return rebuilt
 
     def flatten_state_dict(self, state_dict: dict):
         """
-        Flatten the mixed model's state dict into plain tensors and string-only
-        metadata, running each distinct backend's `flatten_state_dict` once. Returns
-        (state_dict, metadata), ready for `safetensors.torch.save_file(tensors, path,
-        metadata=metadata)`.
+        Flatten the mixed model's state dict into plain tensors and string-only metadata.
+        Returns (state_dict, metadata), ready for `safetensors.torch.save_file(tensors,
+        path, metadata=metadata)`.
         """
         for config in self.configs:
             if not config.backend.capabilities().get("is_serializable", False):
@@ -389,10 +466,21 @@ class MixedQuantizeConfig:
                     "state dict cannot be flattened for saving."
                 )
         merged_metadata = {}
-        for config in self._distinct_backend_configs():
-            state_dict, metadata = config.backend.flatten_state_dict(state_dict)
-            merged_metadata.update(metadata)
-        tensors = {key: value.contiguous() for key, value in state_dict.items()}
+        if not self._ownership:
+            flattened = state_dict
+            for config in self._distinct_backend_configs():
+                flattened, metadata = config.backend.flatten_state_dict(flattened)
+                merged_metadata.update(metadata)
+        else:
+            flattened = {}
+            for config, keys in self._grouped_by_backend(state_dict):
+                if config is None:
+                    flattened.update(keys)
+                    continue
+                tensors, metadata = config.backend.flatten_state_dict(keys)
+                flattened.update(tensors)
+                merged_metadata.update(metadata)
+        tensors = {key: value.contiguous() for key, value in flattened.items()}
         metadata = {"format": "pt", **{key: value if isinstance(value, str) else str(value) for key, value in merged_metadata.items()}}
         return tensors, metadata
 
@@ -405,21 +493,25 @@ class MixedQuantizeConfig:
                 distinct.append(config)
         return distinct
 
-    def _validate_disjoint(self, model):
-        matched = [
-            {name for name, module in model.named_modules() if name and config._should_quantize(name, module)}
-            for config in self.configs
-        ]
-        for i in range(len(matched)):
-            for j in range(i + 1, len(matched)):
-                overlap = matched[i] & matched[j]
-                if overlap:
-                    samples = ", ".join(sorted(overlap)[:5])
-                    raise ValueError(
-                        f"Configs {i} (`{self.configs[i].method}`) and {j} (`{self.configs[j].method}`) "
-                        f"in `MixedQuantizeConfig` both match {len(overlap)} layers (e.g. {samples}). "
-                        "Layer sets must be pairwise disjoint, because re-quantizing an already "
-                        "quantized layer corrupts its packed weight silently. "
-                        "Adjust `target_modules` / `exclude_modules` so the sets are complementary."
-                    )
-
+    def _grouped_by_backend(self, state_dict):
+        groups = {}
+        for key, value in state_dict.items():
+            parts = key.split(".")
+            config = next((owner for end in range(len(parts) - 1, 0, -1)
+                           if (owner := self._ownership.get(".".join(parts[:end]))) is not None), None)
+            groups.setdefault(config.backend.name if config else None, (config, {}))[1][key] = value
+        return groups.values()
+
+    def _build_ownership(self, model):
+        ownership = {}
+        for index, config in enumerate(self.configs):
+            names = [name for name, module in model.named_modules() if name and config._should_quantize(name, module)]
+            overlap = sorted(set(names) & ownership.keys())
+            if overlap:
+                owner = self.configs.index(ownership[overlap[0]])
+                raise ValueError(
+                    f"Configs {owner} (`{self.configs[owner].method}`) and {index} (`{config.method}`) "
+                    f"in `MixedQuantizeConfig` both match {len(overlap)} layers (e.g. {', '.join(overlap[:5])}). "
+                )
+            ownership.update({name: config for name in names})
+        self._ownership.update(ownership)