[diffusion] feat: support nunchaku for Z-Image-Turbo and flux.1 (int4) (#18959)

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -31,7 +31,7 @@ def set_default_torch_dtype(dtype: torch.dtype):
 
 
 def get_param_names_mapping(
-    mapping_dict: dict[str, str],
+    mapping_dict: dict[str, str | tuple[str, int, int]],
 ) -> Callable[[str], tuple[str, Any, Any]]:
     """
     Creates a mapping function that transforms parameter names using regex patterns.
@@ -44,21 +44,50 @@ def get_param_names_mapping(
     """
 
     def mapping_fn(name: str) -> tuple[str, Any, Any]:
-        # Try to match and transform the name using the regex patterns in mapping_dict
-        for pattern, replacement in mapping_dict.items():
-            match = re.match(pattern, name)
-            if match:
-                merge_index = None
-                total_split_params = None
-                if isinstance(replacement, tuple):
-                    merge_index = replacement[1]
-                    total_split_params = replacement[2]
-                    replacement = replacement[0]
-                name = re.sub(pattern, replacement, name)
-                return name, merge_index, total_split_params
+        # support chained conversions, e.g.:
+        # transformer.xxx.lora_down -> xxx.lora_down -> xxx.proj_down
+        merge_index = None
+        total_split_params = None
+        max_steps = max(8, len(mapping_dict) * 2)
+        applied_patterns: set[str] = set()
+        visited_names: set[str] = {name}
 
-        # If no pattern matches, return the original name
-        return name, None, None
+        for _ in range(max_steps):
+            transformed = False
+            for pattern, replacement in mapping_dict.items():
+                # avoid re-applying the same rule on its own output
+                if pattern in applied_patterns:
+                    continue
+                if re.match(pattern, name) is None:
+                    continue
+
+                curr_merge_index = None
+                curr_total_split_params = None
+                if isinstance(replacement, tuple):
+                    curr_merge_index = replacement[1]
+                    curr_total_split_params = replacement[2]
+                    replacement = replacement[0]
+
+                new_name = re.sub(pattern, replacement, name)
+
+                if new_name != name:
+                    if curr_merge_index is not None:
+                        merge_index = curr_merge_index
+                        total_split_params = curr_total_split_params
+
+                    name = new_name
+                    applied_patterns.add(pattern)
+                    if name in visited_names:
+                        transformed = False
+                        break
+                    visited_names.add(name)
+                    transformed = True
+                    break
+
+            if not transformed:
+                break
+
+        return name, merge_index, total_split_params
 
     return mapping_fn
 
@@ -150,25 +179,5 @@ def _list_safetensors_files(model_path: str) -> list[str]:
 
 BYTES_PER_GB = 1024**3
 
-
-def get_memory_usage_of_component(module) -> float | None:
-    """
-    returned value is in GB, rounded to 2 decimal digits
-    """
-    if not isinstance(module, nn.Module):
-        return None
-    if hasattr(module, "get_memory_footprint"):
-        usage = module.get_memory_footprint() / BYTES_PER_GB
-    else:
-        # manually
-        param_size = sum(p.numel() * p.element_size() for p in module.parameters())
-        buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
-
-        total_size_bytes = param_size + buffer_size
-        usage = total_size_bytes / (1024**3)
-
-    return round(usage, 2)
-
-
 # component name ->  ComponentLoader class
 component_name_to_loader_cls: Dict[str, Type[Any]] = {}