Move controlnet specific weight handling to separate class

src/mflux/controlnet/flux_controlnet.py

@@ -10,6 +10,7 @@ from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.controlnet_util import ControlnetUtil
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
+from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
 from mflux.models.transformer.transformer import Transformer
@@ -79,7 +80,7 @@ class Flux1Controlnet:
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
 
-        weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer(
+        weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer(
             controlnet_id=CONTROLNET_ID
         )
         self.transformer_controlnet = TransformerControlnet(

src/mflux/controlnet/weight_handler_controlnet.py

@@ -0,0 +1,45 @@
+import json
+from pathlib import Path
+
+import mlx.core as mx
+from huggingface_hub import snapshot_download
+from mlx.utils import tree_unflatten
+
+from mflux.weights.weight_util import WeightUtil
+
+
+class WeightHandlerControlnet:
+    @staticmethod
+    def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
+        controlnet_path = Path(
+            snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])
+        )
+        file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
+        quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
+        weights = list(mx.load(str(file)).items())
+
+        if quantization_level is not None:
+            return tree_unflatten(weights), quantization_level
+
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
+        weights = WeightUtil.flatten(weights)
+        weights = tree_unflatten(weights)
+
+        # Quantized weights (i.e. ones exported from this project) don't need any post-processing.
+        if quantization_level is not None:
+            return weights, quantization_level
+
+        # Reshape and process the huggingface weights
+        if "transformer_blocks" in weights:
+            for block in weights["transformer_blocks"]:
+                block["ff"] = {
+                    "linear1": block["ff"]["net"][0]["proj"],
+                    "linear2": block["ff"]["net"][2]
+                }  # fmt: off
+                if block.get("ff_context") is not None:
+                    block["ff_context"] = {
+                        "linear1": block["ff_context"]["net"][0]["proj"],
+                        "linear2": block["ff_context"]["net"][2],
+                    }
+        config = json.load(open(controlnet_path / "config.json"))
+        return weights, quantization_level, config

src/mflux/weights/weight_handler.py

@@ -1,4 +1,3 @@
-import json
 from pathlib import Path
 
 import mlx.core as mx
@@ -88,44 +87,6 @@ class WeightHandler:
                     }
         return weights, quantization_level
 
-    @staticmethod
-    def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
-        controlnet_path = Path(
-            snapshot_download(
-                repo_id=controlnet_id,
-                allow_patterns=["*.safetensors", "config.json"],
-            )
-        )
-        file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
-        quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
-        weights = list(mx.load(str(file)).items())
-
-        if quantization_level is not None:
-            return tree_unflatten(weights), quantization_level
-
-        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
-        weights = WeightUtil.flatten(weights)
-        weights = tree_unflatten(weights)
-
-        # Quantized weights (i.e. ones exported from this project) don't need any post-processing.
-        if quantization_level is not None:
-            return weights, quantization_level
-
-        # Reshape and process the huggingface weights
-        if "transformer_blocks" in weights:
-            for block in weights["transformer_blocks"]:
-                block["ff"] = {
-                    "linear1": block["ff"]["net"][0]["proj"],
-                    "linear2": block["ff"]["net"][2],
-                }
-                if block.get("ff_context") is not None:
-                    block["ff_context"] = {
-                        "linear1": block["ff_context"]["net"][0]["proj"],
-                        "linear2": block["ff_context"]["net"][2],
-                    }
-        config = json.load(open(controlnet_path / "config.json"))
-        return weights, quantization_level, config
-
     @staticmethod
     def load_vae(root_path: Path) -> (dict, int):
         weights, quantization_level = WeightHandler._get_weights("vae", root_path)