monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move controlnet specific weight handling to separate class

Browse Source
This commit is contained in:
filipstrand 2024-09-21 22:44:53 +02:00
parent 8678102321
commit cda698df8d
3 changed files with 47 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/mflux/controlnet/flux_controlnet.py
View File

@ -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(

45
src/mflux/controlnet/weight_handler_controlnet.py Normal file
View File

@ -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

39
src/mflux/weights/weight_handler.py
View File

@ -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)