[Tests] Config driven tiny model builder (#5090)

Signed-off-by: Nick Cao <ncao@redhat.com>

tests/model_tests/diffusion/diff_model_builders.py

@@ -1,63 +1,17 @@
-import os
-import tempfile
+from pathlib import Path
 
-import torch
-from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
-from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
-from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
-from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
+from tests.helpers.tiny_model import build_tiny_from_configs
 
-TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
-
-
-def _get_tiny_model_path(name: str) -> str:
-    path = os.path.join(TINY_MODEL_DIR, name)
-    os.makedirs(path, exist_ok=True)
-    return path
+TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs"
 
 
 def tiny_flux2_klein_builder() -> str:
-    """Build a tiny Flux2Klein model."""
-    model_dir = _get_tiny_model_path("Flux2KleinPipeline")
-
-    pipe = Flux2KleinPipeline(
-        scheduler=FlowMatchEulerDiscreteScheduler(),
-        vae=AutoencoderKLFlux2(
-            down_block_types=("DownEncoderBlock2D",),
-            up_block_types=("UpDecoderBlock2D",),
-            block_out_channels=(32,),
-            layers_per_block=1,
-            latent_channels=16,
-            norm_num_groups=16,
-        ),
-        # NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
-        text_encoder=Qwen3ForCausalLM(
-            Qwen3Config(
-                hidden_size=32,
-                intermediate_size=64,
-                num_hidden_layers=28,
-                num_attention_heads=2,
-                num_key_value_heads=2,
-                head_dim=16,
-                vocab_size=151936,
-                max_position_embeddings=512,
-            )
-        ),
-        tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
-        # NOTE: For now we need at least 2 layers for the transformer
-        # due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
-        transformer=Flux2Transformer2DModel(
-            in_channels=64,
-            num_layers=2,
-            num_single_layers=2,
-            attention_head_dim=16,
-            num_attention_heads=2,
-            joint_attention_dim=96,
-            timestep_guidance_channels=32,
-            axes_dims_rope=(4, 4, 4, 4),
-        ),
+    """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"
     )
-    # Need dtypes to be consistent; for now we just put it on bfloat16
-    pipe.to(torch.bfloat16).save_pretrained(model_dir)
-    return model_dir
+
+
+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")