up

Author: sayakpaul

src/diffusers/loaders/single_file_model.py

@@ -23,6 +23,7 @@
 
 from .. import __version__
 from ..quantizers import DiffusersAutoQuantizer
+from ..quantizers.quantization_config import NunchakuConfig
 from ..utils import deprecate, is_accelerate_available, is_torch_version, logging
 from ..utils.torch_utils import empty_device_cache
 from .single_file_utils import (
@@ -442,6 +443,7 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
             )
 
         # This step is better off here than above because `diffusers_format_checkpoint` holds the keys we expect.
+        # We can move it to a separate function as well.
         if quantization_config is not None:
             original_modules_to_not_convert = quantization_config.modules_to_not_convert or []
             determined_modules_to_not_convert = _maybe_determine_modules_to_not_convert(
@@ -450,12 +452,15 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =
             if determined_modules_to_not_convert:
                 determined_modules_to_not_convert.extend(original_modules_to_not_convert)
                 determined_modules_to_not_convert = list(set(determined_modules_to_not_convert))
-                logger.info(
+                logger.debug(
                     f"`modules_to_not_convert` in the quantization_config was updated from {quantization_config.modules_to_not_convert} to {determined_modules_to_not_convert}."
                 )
-                quantization_config.modules_to_not_convert = original_modules_to_not_convert
-                # Update the `quant_config`.
-                hf_quantizer.quantization_config = quantization_config
+                modified_quant_config = quantization_config.to_dict()
+                modified_quant_config["modules_to_not_convert"] = determined_modules_to_not_convert
+                # TODO: figure out a better way.
+                modified_quant_config = NunchakuConfig.from_dict(modified_quant_config)
+                setattr(hf_quantizer, "quantization_config", modified_quant_config)
+                logger.debug("TODO")
 
         # Check if `_keep_in_fp32_modules` is not None
         use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and (

src/diffusers/loaders/single_file_utils.py

@@ -2220,7 +2220,7 @@ def convert_nunchaku_flux_to_diffusers(checkpoint, **kwargs):
         if k.startswith("single_transformer_blocks."):
             # attention / qkv / norms
             new_k = new_k.replace(".qkv_proj.", ".attn.to_qkv.")
-            new_k = new_k.replace(".out_proj.", ".attn.to_out.")
+            new_k = new_k.replace(".out_proj.", ".proj_out.")
             new_k = new_k.replace(".norm_k.", ".attn.norm_k.")
             new_k = new_k.replace(".norm_q.", ".attn.norm_q.")
 
@@ -2279,7 +2279,11 @@ def convert_nunchaku_flux_to_diffusers(checkpoint, **kwargs):
     for k in model_state_dict:
         if k not in new_state_dict:
             # CPU device for now
-            new_state_dict[k] = torch.ones_like(k, device="cpu")
+            new_state_dict[k] = torch.ones_like(model_state_dict[k], device="cpu")
+
+    for k in new_state_dict:
+        if "single_transformer_blocks.0" in k and k.endswith(".weight"):
+            print(f"{k=}")
 
     return new_state_dict
 

src/diffusers/loaders/single_file_utils_nunchaku.py

@@ -6,8 +6,8 @@
 _QKV_ANCHORS_NUNCHAKU = ("to_qkv", "add_qkv_proj")
 _ALLOWED_SUFFIXES_NUNCHAKU = {
     "bias",
-    "lora_down",
-    "lora_up",
+    "proj_down",
+    "proj_up",
     "qweight",
     "smooth_factor",
     "smooth_factor_orig",
@@ -66,14 +66,16 @@ def _unpack_qkv_state_dict(
         '...to_q.bias', '...to_k.bias', '...to_v.bias' '...to_q.wscales', '...to_k.wscales', '...to_v.wscales'
     Returns a NEW dict; original is not modified.
 
-    Only keys with suffix in `allowed_suffixes` are processed. Keys with non-divisible-by-3 tensors raise a ValueError.
+    Only keys with suffix in `allowed_suffixes` are processed. Keys with non-divisible-by-3 tensors raise a ValueError.:
     """
     anchors = tuple(anchors)
     allowed_suffixes = set(allowed_suffixes)
 
     new_sd: dict = {}
-    for k, v in state_dict.items():
+    sd_keys = list(state_dict.keys())
+    for k in sd_keys:
         m = _QKV_NUNCHAKU_REGEX.match(k)
+        v = state_dict.pop(k)
         if m:
             suffix = m.group("suffix")
             if suffix not in allowed_suffixes:

src/diffusers/quantizers/auto.py

@@ -56,7 +56,7 @@
 
 class DiffusersAutoQuantizer:
     """
-     The auto diffusers quantizer class that takes care of automatically instantiating to the correct
+    The auto diffusers quantizer class that takes care of automatically instantiating to the correct
     `DiffusersQuantizer` given the `QuantizationConfig`.
     """
 

src/diffusers/quantizers/quantization_config.py

@@ -791,25 +791,6 @@ def post_init(self):
 
         # TODO: should there be a check for rank?
 
-    # Copied from diffusers.quantizers.quantization_config.BitsAndBytesConfig.to_diff_dict with BitsAndBytesConfig->NunchakuConfig
-    def to_diff_dict(self) -> Dict[str, Any]:
-        """
-        Removes all attributes from config which correspond to the default config attributes for better readability and
-        serializes to a Python dictionary.
-
-        Returns:
-            `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
-        """
+    def __repr__(self):
         config_dict = self.to_dict()
-
-        # get the default config dict
-        default_config_dict = NunchakuConfig().to_dict()
-
-        serializable_config_dict = {}
-
-        # only serialize values that differ from the default config
-        for key, value in config_dict.items():
-            if value != default_config_dict[key]:
-                serializable_config_dict[key] = value
-
-        return serializable_config_dict
+        return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"