[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
|
||||
|
||||
Reference in new issue
Block a user