191 lines
7.6 KiB
Python
191 lines
7.6 KiB
Python
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
from huggingface_hub import snapshot_download
|
|
from mlx.utils import tree_unflatten
|
|
|
|
from mflux.weights.lora_converter import LoRAConverter
|
|
from mflux.weights.weight_util import WeightUtil
|
|
|
|
|
|
@dataclass
|
|
class MetaData:
|
|
quantization_level: int | None = None
|
|
scale: float | None = None
|
|
is_lora: bool = False
|
|
is_mflux: bool = False
|
|
|
|
|
|
class WeightHandler:
|
|
def __init__(
|
|
self,
|
|
meta_data: MetaData,
|
|
clip_encoder: dict | None = None,
|
|
t5_encoder: dict | None = None,
|
|
vae: dict | None = None,
|
|
transformer: dict | None = None,
|
|
):
|
|
self.clip_encoder = clip_encoder
|
|
self.t5_encoder = t5_encoder
|
|
self.vae = vae
|
|
self.transformer = transformer
|
|
self.meta_data = meta_data
|
|
|
|
@staticmethod
|
|
def load_regular_weights(
|
|
repo_id: str | None = None,
|
|
local_path: str | None = None,
|
|
) -> "WeightHandler":
|
|
# Load the weights from disk, huggingface cache, or download from huggingface
|
|
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
|
|
|
|
clip_encoder, _ = WeightHandler._load_clip_encoder(root_path=root_path)
|
|
t5_encoder, _ = WeightHandler._load_t5_encoder(root_path=root_path)
|
|
vae, _ = WeightHandler._load_vae(root_path=root_path)
|
|
transformer, quantization_level, _ = WeightHandler.load_transformer(root_path=root_path)
|
|
return WeightHandler(
|
|
clip_encoder=clip_encoder,
|
|
t5_encoder=t5_encoder,
|
|
vae=vae,
|
|
transformer=transformer,
|
|
meta_data=MetaData(
|
|
quantization_level=quantization_level,
|
|
scale=None,
|
|
is_lora=False,
|
|
is_mflux=False,
|
|
),
|
|
)
|
|
|
|
def num_transformer_blocks(self) -> int:
|
|
return len(self.transformer["transformer_blocks"])
|
|
|
|
def num_single_transformer_blocks(self) -> int:
|
|
return len(self.transformer["single_transformer_blocks"])
|
|
|
|
@staticmethod
|
|
def _load_clip_encoder(root_path: Path) -> (dict, int):
|
|
weights, quantization_level, _ = WeightHandler._get_weights("text_encoder", root_path)
|
|
return weights, quantization_level
|
|
|
|
@staticmethod
|
|
def _load_t5_encoder(root_path: Path) -> (dict, int):
|
|
weights, quantization_level, _ = WeightHandler._get_weights("text_encoder_2", root_path)
|
|
|
|
# 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
|
|
weights["final_layer_norm"] = weights["encoder"]["final_layer_norm"]
|
|
for block in weights["encoder"]["block"]:
|
|
attention = block["layer"][0]
|
|
ff = block["layer"][1]
|
|
block.pop("layer")
|
|
block["attention"] = attention
|
|
block["ff"] = ff
|
|
|
|
weights["t5_blocks"] = weights["encoder"]["block"]
|
|
|
|
# Only the first layer has the weights for "relative_attention_bias", we duplicate them here to keep code simple
|
|
relative_attention_bias = weights["t5_blocks"][0]["attention"]["SelfAttention"]["relative_attention_bias"]
|
|
for block in weights["t5_blocks"][1:]:
|
|
block["attention"]["SelfAttention"]["relative_attention_bias"] = relative_attention_bias
|
|
|
|
weights.pop("encoder")
|
|
return weights, quantization_level
|
|
|
|
@staticmethod
|
|
def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int, str | None):
|
|
weights, quantization_level, mflux_version = WeightHandler._get_weights("transformer", root_path, lora_path)
|
|
|
|
if lora_path:
|
|
if "transformer" not in weights:
|
|
weights = LoRAConverter.load_weights(lora_path)
|
|
weights = weights["transformer"]
|
|
|
|
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
|
|
if quantization_level or mflux_version:
|
|
return weights, quantization_level, mflux_version
|
|
|
|
# Reshape and process the huggingface weights
|
|
if "transformer_blocks" in weights:
|
|
for block in weights["transformer_blocks"]:
|
|
if block.get("ff") is not None:
|
|
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 weights, quantization_level, mflux_version
|
|
|
|
@staticmethod
|
|
def _load_vae(root_path: Path) -> (dict, int):
|
|
weights, quantization_level, _ = WeightHandler._get_weights("vae", root_path)
|
|
|
|
# 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
|
|
weights["decoder"]["conv_in"] = {"conv2d": weights["decoder"]["conv_in"]}
|
|
weights["decoder"]["conv_out"] = {"conv2d": weights["decoder"]["conv_out"]}
|
|
weights["decoder"]["conv_norm_out"] = {"norm": weights["decoder"]["conv_norm_out"]}
|
|
weights["encoder"]["conv_in"] = {"conv2d": weights["encoder"]["conv_in"]}
|
|
weights["encoder"]["conv_out"] = {"conv2d": weights["encoder"]["conv_out"]}
|
|
weights["encoder"]["conv_norm_out"] = {"norm": weights["encoder"]["conv_norm_out"]}
|
|
return weights, quantization_level
|
|
|
|
@staticmethod
|
|
def _get_weights(
|
|
model_name: str,
|
|
root_path: Path | None = None,
|
|
lora_path: str | None = None,
|
|
) -> (dict, int, str | None):
|
|
weights = []
|
|
quantization_level = None
|
|
mflux_version = None
|
|
|
|
if root_path is not None:
|
|
for file in sorted(root_path.glob(model_name + "/*.safetensors")):
|
|
data = mx.load(str(file), return_metadata=True)
|
|
weight = list(data[0].items())
|
|
if len(data) > 1:
|
|
quantization_level = data[1].get("quantization_level")
|
|
weights.extend(weight)
|
|
|
|
if lora_path and root_path is None:
|
|
data = mx.load(lora_path, return_metadata=True)
|
|
weight = list(data[0].items())
|
|
if len(data) > 1:
|
|
mflux_version = data[1].get("mflux_version", None)
|
|
weights.extend(weight)
|
|
|
|
# Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.
|
|
if quantization_level is not None or mflux_version is not None:
|
|
return tree_unflatten(weights), quantization_level, mflux_version
|
|
|
|
# Huggingface weights needs to be reshaped
|
|
weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
|
|
weights = WeightUtil.flatten(weights)
|
|
unflatten = tree_unflatten(weights)
|
|
return unflatten, quantization_level, mflux_version
|
|
|
|
@staticmethod
|
|
def _download_or_get_cached_weights(repo_id: str) -> Path:
|
|
return Path(
|
|
snapshot_download(
|
|
repo_id=repo_id,
|
|
allow_patterns=[
|
|
"text_encoder/*.safetensors",
|
|
"text_encoder_2/*.safetensors",
|
|
"transformer/*.safetensors",
|
|
"vae/*.safetensors",
|
|
],
|
|
)
|
|
)
|