Introduce weight util for shared methods

src/flux_1/weights/lora_util.py

@@ -6,7 +6,7 @@ from mlx.utils import tree_flatten
 from mlx.utils import tree_unflatten
 from safetensors import safe_open
 
-from flux_1.weights.weight_handler import WeightHandler
+from flux_1.weights.weight_util import WeightUtil
 
 log = logging.getLogger(__name__)
 
@@ -86,8 +86,8 @@ class LoraUtil:
     def _lora_transformer(lora_file: Path) -> (dict, int):
         quantization_level = safe_open(lora_file, framework="pt").metadata().get("quantization_level")
         weights = list(mx.load(str(lora_file)).items())
-        weights = [WeightHandler.reshape_weights(k, v) for k, v in weights]
-        weights = WeightHandler.flatten(weights)
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
+        weights = WeightUtil.flatten(weights)
         unflatten = tree_unflatten(weights)
         for block in unflatten["transformer"]["transformer_blocks"]:
             block["ff"] = {

src/flux_1/weights/weight_handler.py

@@ -5,8 +5,8 @@ from huggingface_hub import snapshot_download
 from mlx.utils import tree_unflatten
 from safetensors import safe_open
 
-from flux_1.config.config import Config
 from flux_1.weights.lora_util import LoraUtil
+from flux_1.weights.weight_util import WeightUtil
 
 
 class WeightHandler:
@@ -117,22 +117,11 @@ class WeightHandler:
             return tree_unflatten(weights), quantization_level
 
         # Huggingface weights needs to be reshaped
-        weights = [WeightHandler.reshape_weights(k, v) for k, v in weights]
-        weights = WeightHandler.flatten(weights)
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
+        weights = WeightUtil.flatten(weights)
         unflatten = tree_unflatten(weights)
         return unflatten, quantization_level
 
-    @staticmethod
-    def flatten(params):
-        return [(k, v) for p in params for (k, v) in p]
-
-    @staticmethod
-    def reshape_weights(key, value):
-        if len(value.shape) == 4:
-            value = value.transpose(0, 2, 3, 1)
-        value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
-        return [(key, value)]
-
     @staticmethod
     def _download_or_get_cached_weights(repo_id: str) -> Path:
         return Path(

src/flux_1/weights/weight_util.py

@@ -0,0 +1,15 @@
+from flux_1.config.config import Config
+
+
+class WeightUtil:
+
+    @staticmethod
+    def flatten(params):
+        return [(k, v) for p in params for (k, v) in p]
+
+    @staticmethod
+    def reshape_weights(key, value):
+        if len(value.shape) == 4:
+            value = value.transpose(0, 2, 3, 1)
+        value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
+        return [(key, value)]