[CI][Diffusion] Add tiny model builder for StableDiffusion3Pipeline (#6423)

Signed-off-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
This commit is contained in:
Khairul Kabirandkhairulkabir1661 authored and GitHub committed 2026-09-24 19:23:28 +00:00
1 parent b539282a9a
commit 29d13e3311
3 files changed
+27 -1

No files matched your search

@@ -256,6 +256,28 @@ def tiny_flux_kontext_builder() -> str:
)
def _shrink_sd3_transformer(config: dict) -> dict:
config["num_layers"] = 2
config["attention_head_dim"] = 32
config["num_attention_heads"] = 4
config["caption_projection_dim"] = 128
config["dual_attention_layers"] = [0]
return config
def tiny_sd3_builder() -> str:
return build_tiny_from_configs(
"StableDiffusion3Pipeline",
"stabilityai/stable-diffusion-3.5-medium",
transform={
"text_encoder": _shrink_flux_clip_text_encoder,
"text_encoder_2": _shrink_flux_clip_text_encoder,
"text_encoder_3": _shrink_flux_t5_text_encoder,
"transformer": _shrink_sd3_transformer,
},
)
def tiny_flux2_builder() -> str:
def shrink_text_encoder(config: dict) -> dict:
config["tie_word_embeddings"] = False