[Tests] Add tiny model test for StableDiffusionXLPipeline
Covers base text-to-image plus CPU offload; UNet stays full-size since its architecture is hardcoded rather than config-driven. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Nick Cao <ncao@redhat.com>
This commit is contained in:
1 parent
9146284c16
commit
51b0184969
3 files changed
+20
-1
No files matched your search
@@ -278,6 +278,20 @@ def tiny_sd3_builder() -> str:
|
||||
)
|
||||
|
||||
|
||||
def tiny_sdxl_builder() -> str:
|
||||
# The UNet's architecture (block_out_channels, transformer_layers_per_block, etc.)
|
||||
# is hardcoded in SDXLUNet2DConditionModel rather than read from config, so it
|
||||
# can't be shrunk here; only the two CLIP text encoders are reduced.
|
||||
return build_tiny_from_configs(
|
||||
"StableDiffusionXLPipeline",
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
transform={
|
||||
"text_encoder": _shrink_flux_clip_text_encoder,
|
||||
"text_encoder_2": _shrink_flux_clip_text_encoder,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def tiny_flux2_builder() -> str:
|
||||
def shrink_text_encoder(config: dict) -> dict:
|
||||
config["tie_word_embeddings"] = False
|
||||
|
||||
@@ -181,4 +181,10 @@ DIFFUSION_TEST_SETTINGS = {
|
||||
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
||||
model_type_marker=ModelTypeMarker.DIFFUSION,
|
||||
),
|
||||
"StableDiffusionXLPipeline": DiffusionModelTestOpts(
|
||||
model="stabilityai/stable-diffusion-xl-base-1.0",
|
||||
builder=diff_model_builders.tiny_sdxl_builder,
|
||||
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
|
||||
extra_test_groups=[[DiffusionAccs.CPU_OFFLOAD]],
|
||||
),
|
||||
}
|
||||
@@ -84,7 +84,6 @@ EXCLUDED_MODELS = [
|
||||
"HiDreamImagePipeline",
|
||||
"HiDreamO1ImagePipeline",
|
||||
"DreamZeroPipeline",
|
||||
"StableDiffusionXLPipeline",
|
||||
"Gr00tN1d7Pipeline",
|
||||
"Pi0Pipeline",
|
||||
"Pi05Pipeline",
|
||||
|
||||
Reference in new issue
Block a user