83a7bf3519
* Use magic trailing commas instead of disabling formatting to keep args on separate lines * Scope `fmt: off`s better where that's not possible
71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
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_handler import MetaData
|
|
from mflux.weights.weight_util import WeightUtil
|
|
|
|
CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
|
|
|
|
|
|
class WeightHandlerControlnet:
|
|
def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None):
|
|
self.meta_data = meta_data
|
|
self.controlnet_transformer = controlnet_transformer
|
|
self.config = config
|
|
|
|
@staticmethod
|
|
def load_controlnet_transformer() -> "WeightHandlerControlnet":
|
|
controlnet_path = Path(snapshot_download(repo_id=CONTROLNET_ID, allow_patterns=["*.safetensors", "config.json"])) # fmt:off
|
|
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())
|
|
config = json.load(open(controlnet_path / "config.json"))
|
|
|
|
if quantization_level is not None:
|
|
return WeightHandlerControlnet(
|
|
config=config,
|
|
controlnet_transformer=tree_unflatten(weights),
|
|
meta_data=MetaData(quantization_level=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 WeightHandlerControlnet(
|
|
config=config,
|
|
controlnet_transformer=weights,
|
|
meta_data=MetaData(quantization_level=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],
|
|
}
|
|
|
|
return WeightHandlerControlnet(
|
|
config=config,
|
|
controlnet_transformer=weights,
|
|
meta_data=MetaData(quantization_level=quantization_level),
|
|
)
|
|
|
|
def num_transformer_blocks(self) -> int:
|
|
return self.config["num_layers"]
|
|
|
|
def num_single_transformer_blocks(self) -> int:
|
|
return self.config["num_single_layers"]
|