From cda698df8dc458cec5d032cf921e83321a6d673e Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 21 Sep 2024 22:44:53 +0200 Subject: [PATCH] Move controlnet specific weight handling to separate class --- src/mflux/controlnet/flux_controlnet.py | 3 +- .../controlnet/weight_handler_controlnet.py | 45 +++++++++++++++++++ src/mflux/weights/weight_handler.py | 39 ---------------- 3 files changed, 47 insertions(+), 40 deletions(-) create mode 100644 src/mflux/controlnet/weight_handler_controlnet.py diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index c9e5b1c..0fba545 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/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( diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py new file mode 100644 index 0000000..00654c9 --- /dev/null +++ b/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 diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index 56d25cc..1db8f95 100644 --- a/src/mflux/weights/weight_handler.py +++ b/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)