273fc8eb1d
Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
from functools import partial
|
|
from pathlib import Path
|
|
|
|
from tests.helpers.tiny_model import build_tiny_from_configs
|
|
|
|
TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs"
|
|
|
|
|
|
def _shrink_dit_rope_config(
|
|
config: dict,
|
|
num_layers: int = 2,
|
|
num_single_layers: int | None = None,
|
|
target_head_dim: int = 32,
|
|
target_heads: int = 4,
|
|
default_axes_dims_rope: list[int] | None = None,
|
|
) -> dict:
|
|
config["num_layers"] = num_layers
|
|
if num_single_layers is not None:
|
|
config["num_single_layers"] = num_single_layers
|
|
axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
|
|
factor = config["attention_head_dim"] / target_head_dim
|
|
config["attention_head_dim"] = target_head_dim
|
|
config["num_attention_heads"] = target_heads
|
|
config["axes_dims_rope"] = [int(d / factor) for d in axes_dims_rope]
|
|
return config
|
|
|
|
|
|
def tiny_flux2_klein_builder() -> str:
|
|
"""Build a tiny Flux2Klein model from vendored configs."""
|
|
return build_tiny_from_configs(
|
|
"Flux2KleinPipeline", "black-forest-labs/FLUX.2-klein-4B", TINY_CONFIGS_DIR / "Flux2KleinPipeline"
|
|
)
|
|
|
|
|
|
def tiny_ltx2_builder() -> str:
|
|
"""Build a tiny LTX2 model from vendored configs."""
|
|
return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
|
|
|
|
|
|
def tiny_qwen_image_builder() -> str:
|
|
def shrink_text_encoder(config: dict) -> dict:
|
|
config["num_hidden_layers"] = 2
|
|
config["intermediate_size"] = 64
|
|
config["text_config"]["num_hidden_layers"] = 2
|
|
config["text_config"]["intermediate_size"] = 64
|
|
config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
|
|
config["vision_config"]["depth"] = 2
|
|
config["vision_config"]["intermediate_size"] = 64
|
|
config["vision_config"]["fullatt_block_indexes"] = [0, 1]
|
|
return config
|
|
|
|
return build_tiny_from_configs(
|
|
"QwenImagePipeline",
|
|
"Qwen/Qwen-Image",
|
|
transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
|
|
)
|
|
|
|
|
|
def tiny_flux_builder() -> str:
|
|
def shrink_clip_text_encoder(config: dict) -> dict:
|
|
config["num_hidden_layers"] = 2
|
|
return config
|
|
|
|
def shrink_t5_text_encoder(config: dict) -> dict:
|
|
config["num_layers"] = 2
|
|
config["num_decoder_layers"] = 2
|
|
config["d_ff"] = 64
|
|
config["num_heads"] = 4
|
|
return config
|
|
|
|
return build_tiny_from_configs(
|
|
"FluxPipeline",
|
|
"black-forest-labs/FLUX.1-schnell",
|
|
transform={
|
|
"text_encoder": shrink_clip_text_encoder,
|
|
"text_encoder_2": shrink_t5_text_encoder,
|
|
"transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
|
|
},
|
|
)
|
|
|
|
|
|
def tiny_flux2_builder() -> str:
|
|
def shrink_text_encoder(config: dict) -> dict:
|
|
# The real checkpoint ships language_model.lm_head.weight as its own
|
|
# tensor even though tie_word_embeddings defaults to True; keep it
|
|
# untied here too so transformers doesn't drop it on save.
|
|
config["tie_word_embeddings"] = False
|
|
config["text_config"]["num_hidden_layers"] = 31
|
|
config["text_config"]["intermediate_size"] = 64
|
|
config["text_config"]["num_attention_heads"] = 4
|
|
config["text_config"]["head_dim"] = 32
|
|
config["text_config"]["num_key_value_heads"] = 2
|
|
config["vision_config"]["num_hidden_layers"] = 2
|
|
config["vision_config"]["intermediate_size"] = 64
|
|
config["vision_config"]["num_attention_heads"] = 4
|
|
config["vision_config"]["head_dim"] = 32
|
|
return config
|
|
|
|
return build_tiny_from_configs(
|
|
"Flux2Pipeline",
|
|
"black-forest-labs/FLUX.2-dev",
|
|
transform={
|
|
"text_encoder": shrink_text_encoder,
|
|
"transformer": partial(_shrink_dit_rope_config, num_single_layers=2),
|
|
},
|
|
)
|