[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:
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
|
||||
|
||||
@@ -146,4 +146,9 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
|
||||
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
|
||||
),
|
||||
"StableDiffusion3Pipeline": DiffusionModelTestOpts(
|
||||
model="stabilityai/stable-diffusion-3.5-medium",
|
||||
builder=diff_model_builders.tiny_sd3_builder,
|
||||
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
||||
),
|
||||
}
|
||||
@@ -56,7 +56,6 @@ EXCLUDED_MODELS = [
|
||||
"MingImageDiffusionPipeline",
|
||||
"MingImageLayeredDiffusionPipeline",
|
||||
"InternVLAA1Pipeline",
|
||||
"StableDiffusion3Pipeline",
|
||||
"HunyuanImage3ForCausalMM",
|
||||
"ErnieImagePipeline",
|
||||
"NextStep11Pipeline",
|
||||
|
||||
Reference in new issue
Block a user