0b884ebc6c
Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
99 lines
3.9 KiB
Python
99 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
|
|
import torch
|
|
from vllm.model_executor.models.utils import WeightsMapper
|
|
|
|
from .base import GGUFAdapter, gguf_quant_weights_iterator
|
|
|
|
FLUX2_TRANSFORMER_KEYS_RENAME_DICT = {
|
|
"single_blocks.": "single_transformer_blocks.",
|
|
# Image and text input projections
|
|
"img_in": "x_embedder",
|
|
"txt_in": "context_embedder",
|
|
# Timestep and guidance embeddings
|
|
"time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
|
|
"time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
|
|
"guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1",
|
|
"guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2",
|
|
# Modulation parameters
|
|
"double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
|
|
"double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
|
|
"single_stream_modulation.lin": "single_stream_modulation.linear",
|
|
# Final output layer
|
|
# "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params
|
|
"final_layer.linear": "proj_out",
|
|
}
|
|
|
|
FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = {
|
|
"final_layer.adaLN_modulation.1": "norm_out.linear",
|
|
}
|
|
|
|
FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = {
|
|
"double_blocks.": "transformer_blocks.",
|
|
# Handle fused QKV projections separately as we need to break into Q, K, V projections
|
|
"img_attn.norm.query_norm": "attn.norm_q",
|
|
"img_attn.norm.key_norm": "attn.norm_k",
|
|
"img_attn.proj": "attn.to_out.0",
|
|
"img_mlp.0": "ff.linear_in",
|
|
"img_mlp.2": "ff.linear_out",
|
|
"txt_attn.norm.query_norm": "attn.norm_added_q",
|
|
"txt_attn.norm.key_norm": "attn.norm_added_k",
|
|
"txt_attn.proj": "attn.to_add_out",
|
|
"txt_mlp.0": "ff_context.linear_in",
|
|
"txt_mlp.2": "ff_context.linear_out",
|
|
# Additional for fuse qkv
|
|
"img_attn.qkv": "attn.to_qkv",
|
|
"txt_attn.qkv": "attn.add_kv_proj",
|
|
}
|
|
|
|
FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = {
|
|
"linear1": "attn.to_qkv_mlp_proj",
|
|
"linear2": "attn.to_out",
|
|
"norm.query_norm": "attn.norm_q",
|
|
"norm.key_norm": "attn.norm_k",
|
|
}
|
|
|
|
|
|
class Flux2KleinGGUFAdapter(GGUFAdapter):
|
|
"""GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap."""
|
|
|
|
@staticmethod
|
|
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
|
|
model_class = od_config.model_class_name or ""
|
|
if model_class.startswith("Flux2"):
|
|
return True
|
|
cfg = od_config.tf_model_config
|
|
if cfg is not None:
|
|
model_type = str(cfg.get("model_type", "")).lower()
|
|
if model_type.startswith("flux"):
|
|
return True
|
|
return False
|
|
|
|
gguf_to_hf_mapper = WeightsMapper(
|
|
# double_stream_modulation
|
|
orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP,
|
|
orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP,
|
|
)
|
|
|
|
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
|
|
def custom_weights_adapter(
|
|
weights: Iterable[tuple[str, torch.Tensor]],
|
|
) -> Iterable[tuple[str, torch.Tensor]]:
|
|
for name, weight in weights:
|
|
# Handle the special case for adaLN modulation parameters that require swapping shift and scale
|
|
if name.endswith(".scale"):
|
|
name = name.replace(".scale", ".weight")
|
|
if name == "norm_out.linear.weight":
|
|
shift, scale = weight.chunk(2, dim=0)
|
|
weight = torch.cat([scale, shift], dim=0)
|
|
yield name, weight
|
|
else:
|
|
yield name, weight
|
|
|
|
weights = gguf_quant_weights_iterator(self.gguf_file)
|
|
weights = self.gguf_to_hf_mapper.apply(weights)
|
|
yield from custom_weights_adapter(weights)
|