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

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

tests/helpers/tiny_model.py

@@ -0,0 +1,87 @@
+"""Reusable utility for building tiny diffusion models with random weights.
+
+Mirrors the component loading loop in DiffusionPipeline.from_pretrained,
+but monkeypatches from_pretrained on ModelMixin and PreTrainedModel to
+initialize with random weights instead of loading checkpoint files.
+Components with vendored configs use those; others load configs from
+the upstream HF model.
+"""
+
+import os
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+import torch
+from diffusers import ModelMixin
+from diffusers.pipelines.pipeline_loading_utils import _get_pipeline_class, simple_get_class_obj
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from transformers import AutoConfig, PreTrainedModel
+
+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
+
+
+def _diffusers_from_config(cls, pretrained_model_name_or_path, **kwargs):
+    """Replacement for ModelMixin.from_pretrained that initializes random weights."""
+    subfolder = kwargs.get("subfolder")
+    load_kwargs = {"subfolder": subfolder} if subfolder else {}
+    config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
+    return cls.from_config(config)
+
+
+def _transformers_from_config(cls, pretrained_model_name_or_path, **kwargs):
+    """Replacement for PreTrainedModel.from_pretrained that initializes random weights."""
+    subfolder = kwargs.get("subfolder")
+    load_kwargs = {"subfolder": subfolder} if subfolder else {}
+    config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
+    return cls(config)
+
+
+def build_tiny_from_configs(pipeline_name: str, model_id: str, configs_dir: str | Path) -> str:
+    """Build a tiny diffusion model from vendored configs with random weights.
+
+    Args:
+        pipeline_name: Name of the pipeline (used as output directory name).
+        model_id: HuggingFace model ID for loading upstream components
+            (tokenizer, scheduler) that don't have vendored configs.
+        configs_dir: Path to the directory containing vendored config files
+            (model_index.json and per-component config.json).
+
+    Returns:
+        Path to the saved tiny model directory with safetensors weights.
+    """
+    model_dir = _get_tiny_model_path(pipeline_name)
+    config_dir = Path(configs_dir)
+
+    config_dict = DiffusionPipeline.load_config(config_dir)
+    pipeline_cls = _get_pipeline_class(DiffusionPipeline, config=config_dict)
+
+    init_dict, _, _ = pipeline_cls.extract_init_dict(config_dict)
+
+    # Pop non-component entries (optional pipeline kwargs like is_distilled),
+    # same as DiffusionPipeline.from_pretrained lines 345-350
+    _, optional_kwargs = DiffusionPipeline._get_signature_keys(pipeline_cls)
+    init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict}
+
+    with (
+        patch.object(ModelMixin, "from_pretrained", classmethod(_diffusers_from_config)),
+        patch.object(PreTrainedModel, "from_pretrained", classmethod(_transformers_from_config)),
+    ):
+        for name, (library_name, class_name) in init_dict.items():
+            cls = simple_get_class_obj(library_name, class_name)
+
+            # Use vendored config dir if available, otherwise upstream model
+            if (config_dir / name).exists():
+                init_kwargs[name] = cls.from_pretrained(config_dir / name)
+            else:
+                init_kwargs[name] = cls.from_pretrained(model_id, subfolder=name)
+
+    pipe = pipeline_cls(**init_kwargs)
+    pipe.to(torch.bfloat16).save_pretrained(model_dir)
+    return model_dir

tests/model_tests/diffusion/config_types.py

@@ -34,7 +34,8 @@ class DiffusionTasks(StrEnum):
 
     TEXT_TO_IMAGE = auto()
     IMAGE_TO_IMAGE = auto()
-    # Text to video, text to audio, etc should be added here as needed
+    TEXT_TO_VIDEO = auto()
+    # Text to audio, etc should be added here as needed
 
 
 class DiffusionModelTestOpts(NamedTuple):

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")

tests/model_tests/diffusion/model_settings.py

@@ -38,5 +38,17 @@ DIFFUSION_TEST_SETTINGS = {
             [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
             [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
         ],
-    )
+    ),
+    "LTX2Pipeline": DiffusionModelTestOpts(
+        model="Lightricks/LTX-2",
+        builder=diff_model_builders.tiny_ltx2_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
+        check_determinism=False,
+        check_multi_output=False,
+        extra_test_groups=[
+            [DiffusionAccs.HSDP, DiffusionAccs.CACHE_DIT],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    ),
 }

tests/model_tests/diffusion/task_runners.py

@@ -28,6 +28,15 @@ IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
     seed=42,
 )
 
+VIDEO_NUM_FRAMES = 9
+VIDEO_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
+    num_inference_steps=4,
+    height=HEIGHT,
+    width=WIDTH,
+    num_frames=VIDEO_NUM_FRAMES,
+    seed=42,
+)
+
 # Online extra_body for diffusion requests
 IMAGE_GEN_EXTRA_BODY = {
     "height": HEIGHT,
@@ -36,6 +45,15 @@ IMAGE_GEN_EXTRA_BODY = {
     "seed": 42,
 }
 
+# Online form_data for video generation requests (multipart /v1/videos API)
+VIDEO_GEN_FORM_DATA = {
+    "height": HEIGHT,
+    "width": WIDTH,
+    "num_inference_steps": 4,
+    "num_frames": VIDEO_NUM_FRAMES,
+    "seed": 42,
+}
+
 
 ### Shared validation
 def _validate_images(images: list[Image.Image], expected_n: int = 1):
@@ -70,11 +88,31 @@ def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
     return images
 
 
+def _validate_video(outputs: list[OmniRequestOutput], expected_n: int = 1):
+    """Given a set of outputs, ensure we got video frames with the expected shape."""
+    assert len(outputs) == expected_n
+    for output in outputs:
+        # Video models return numpy arrays via output.images
+        images = output.images
+        assert len(images) > 0
+        for frame_data in images:
+            assert isinstance(frame_data, np.ndarray)
+            # (num_outputs, num_frames, H, W, C) or (num_frames, H, W, C)
+            assert frame_data.ndim in (4, 5), f"Expected 4D or 5D video array, got shape {frame_data.shape}"
+            assert frame_data.shape[-3] == HEIGHT, f"Expected height {HEIGHT}, got {frame_data.shape[-3]}"
+            assert frame_data.shape[-2] == WIDTH, f"Expected width {WIDTH}, got {frame_data.shape[-2]}"
+            assert frame_data.shape[-1] == 3, f"Expected 3 channels (RGB), got {frame_data.shape[-1]}"
+
+
 ### Offline helpers
 def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
     return omni.generate({"prompt": PROMPT}, params)
 
 
+def _run_offline_t2v(omni: Omni, params: OmniDiffusionSamplingParams = VIDEO_GEN_SAMPLING_PARAMS):
+    return omni.generate({"prompt": PROMPT}, params)
+
+
 def _run_offline_i2i(omni: Omni):
     return omni.generate(
         {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
@@ -130,6 +168,11 @@ def run_and_validate_image_to_image_request(omni: Omni):
     _validate_images(_get_offline_images(_run_offline_i2i(omni)))
 
 
+def run_and_validate_text_to_video_request(omni: Omni):
+    """Run and validate a text to video request."""
+    _validate_video(_run_offline_t2v(omni))
+
+
 def run_and_validate_text_to_image_determinism(omni: Omni):
     """Checks for determinism; for now we just keep this for TTI."""
     _validate_image_gen_determinism(
@@ -144,6 +187,25 @@ def run_and_validate_text_to_image_multi_output(omni: Omni):
     _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
 
 
+def _run_online_t2v(
+    server: OmniServer, client: OpenAIClientHandler, form_data: dict | None = None
+) -> list[DiffusionResponse]:
+    """Run a text to video request through the server's /v1/videos API."""
+    data = dict(form_data or VIDEO_GEN_FORM_DATA)
+    data.setdefault("prompt", PROMPT)
+    data.setdefault("model", server.model)
+    return client.send_video_diffusion_request({"form_data": data})
+
+
+def _get_online_videos(responses: list[DiffusionResponse]) -> list:
+    """Extract the videos from a server response."""
+    assert len(responses) == 1
+    videos = responses[0].videos
+    assert videos is not None
+    assert len(videos) > 0
+    return videos
+
+
 ### Online task runners
 def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
     """Run and validate a text to image request through the server."""
@@ -170,3 +232,8 @@ def run_and_validate_online_text_to_image_multi_output(server: OmniServer, clien
         _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
         expected_n=2,
     )
+
+
+def run_and_validate_online_text_to_video_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate a text to video request through the server."""
+    _get_online_videos(_run_online_t2v(server, client))

tests/model_tests/diffusion/test_alignment.py

@@ -30,7 +30,6 @@ EXCLUDED_MODELS = [
     "OvisImagePipeline",
     "WanPipeline",
     "WanVACEPipeline",
-    "LTX2Pipeline",
     "LTX2ImageToVideoPipeline",
     "LTX2TwoStagesPipeline",
     "LTX2ImageToVideoTwoStagesPipeline",

tests/model_tests/diffusion/test_common_offline.py

@@ -12,6 +12,7 @@ from tests.model_tests.diffusion.task_runners import (
     run_and_validate_text_to_image_determinism,
     run_and_validate_text_to_image_multi_output,
     run_and_validate_text_to_image_request,
+    run_and_validate_text_to_video_request,
 )
 
 # NOTE : Hardware marks are added dynamically based on test requirements
@@ -53,6 +54,8 @@ def test_pipeline_on_supported_tasks(
                     run_and_validate_text_to_image_request(omni)
                 elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
                     run_and_validate_image_to_image_request(omni)
+                elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
+                    run_and_validate_text_to_video_request(omni)
                 else:
                     raise ValueError(f"Task type {task_type} is not yet supported")
 

tests/model_tests/diffusion/test_common_online.py

@@ -19,6 +19,7 @@ from tests.model_tests.diffusion.task_runners import (
     run_and_validate_online_text_to_image_determinism,
     run_and_validate_online_text_to_image_multi_output,
     run_and_validate_online_text_to_image_request,
+    run_and_validate_online_text_to_video_request,
 )
 
 # NOTE : Hardware marks are added dynamically based on test requirements
@@ -61,6 +62,8 @@ def test_online_on_supported_tasks(
                     run_and_validate_online_text_to_image_request(server, client)
                 elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
                     run_and_validate_online_image_to_image_request(server, client)
+                elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
+                    run_and_validate_online_text_to_video_request(server, client)
                 else:
                     raise ValueError(f"Task type {task_type} is not yet supported")
 

tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/model_index.json

@@ -0,0 +1,25 @@
+{
+  "_class_name": "Flux2KleinPipeline",
+  "_diffusers_version": "0.37.0.dev0",
+  "is_distilled": true,
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3ForCausalLM"
+  ],
+  "tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "Flux2Transformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ]
+}

tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json

@@ -0,0 +1,60 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 16,
+  "hidden_act": "silu",
+  "hidden_size": 32,
+  "initializer_range": 0.02,
+  "intermediate_size": 64,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 512,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "num_attention_heads": 2,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.56.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/transformer/config.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "Flux2Transformer2DModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "_name_or_path": "/raid/yiyi/klein-4b-distilled-diffusers/transformer",
+  "attention_head_dim": 16,
+  "axes_dims_rope": [
+    4,
+    4,
+    4,
+    4
+  ],
+  "eps": 1e-06,
+  "guidance_embeds": false,
+  "in_channels": 64,
+  "joint_attention_dim": 96,
+  "mlp_ratio": 3.0,
+  "num_attention_heads": 2,
+  "num_layers": 2,
+  "num_single_layers": 2,
+  "out_channels": null,
+  "patch_size": 1,
+  "rope_theta": 2000,
+  "timestep_guidance_channels": 32
+}

tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/vae/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.37.0.dev0",
+  "_name_or_path": "black-forest-labs/FLUX.2-dev",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    32
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 16,
+  "layers_per_block": 1,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 16,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/audio_vae/config.json

@@ -0,0 +1,25 @@
+{
+  "_class_name": "AutoencoderKLLTX2Audio",
+  "_diffusers_version": "0.37.0.dev0",
+  "attn_resolutions": null,
+  "base_channels": 128,
+  "causality_axis": "height",
+  "ch_mult": [
+    1,
+    2,
+    4
+  ],
+  "double_z": true,
+  "dropout": 0.0,
+  "in_channels": 2,
+  "is_causal": true,
+  "latent_channels": 8,
+  "mel_bins": 64,
+  "mel_hop_length": 160,
+  "mid_block_add_attention": false,
+  "norm_type": "pixel",
+  "num_res_blocks": 1,
+  "output_channels": 2,
+  "resolution": 256,
+  "sample_rate": 16000
+}

tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/connectors/config.json

@@ -0,0 +1,21 @@
+{
+  "_class_name": "LTX2TextConnectors",
+  "_diffusers_version": "0.37.0.dev0",
+  "audio_connector_attention_head_dim": 16,
+  "audio_connector_num_attention_heads": 2,
+  "audio_connector_num_layers": 1,
+  "audio_connector_num_learnable_registers": 4,
+  "caption_channels": 32,
+  "causal_temporal_positioning": false,
+  "connector_rope_base_seq_len": 32,
+  "rope_double_precision": true,
+  "rope_theta": 10000.0,
+  "rope_type": "split",
+  "text_proj_in_factor": 3,
+  "video_connector_attention_head_dim": 16,
+  "video_connector_num_attention_heads": 2,
+  "video_connector_num_layers": 1,
+  "video_connector_num_learnable_registers": 4,
+  "video_hidden_dim": 32,
+  "audio_hidden_dim": 32
+}

tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/model_index.json

@@ -0,0 +1,36 @@
+{
+  "_class_name": "LTX2Pipeline",
+  "_diffusers_version": "0.37.0.dev0",
+  "audio_vae": [
+    "diffusers",
+    "AutoencoderKLLTX2Audio"
+  ],
+  "connectors": [
+    "ltx2",
+    "LTX2TextConnectors"
+  ],
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "Gemma3ForConditionalGeneration"
+  ],
+  "tokenizer": [
+    "transformers",
+    "GemmaTokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "LTX2VideoTransformer3DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLLTX2Video"
+  ],
+  "vocoder": [
+    "ltx2",
+    "LTX2Vocoder"
+  ]
+}

tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/text_encoder/config.json

@@ -0,0 +1,68 @@
+{
+  "architectures": [
+    "Gemma3ForConditionalGeneration"
+  ],
+  "boi_token_index": 255999,
+  "dtype": "float32",
+  "eoi_token_index": 256000,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_index": 262144,
+  "initializer_range": 0.02,
+  "mm_tokens_per_image": 256,
+  "model_type": "gemma3",
+  "text_config": {
+    "_sliding_window_pattern": 6,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_logit_softcapping": null,
+    "cache_implementation": "hybrid",
+    "dtype": "float32",
+    "final_logit_softcapping": null,
+    "head_dim": 16,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 32,
+    "initializer_range": 0.02,
+    "intermediate_size": 64,
+    "layer_types": [
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 512,
+    "model_type": "gemma3_text",
+    "num_attention_heads": 2,
+    "num_hidden_layers": 2,
+    "num_key_value_heads": 2,
+    "query_pre_attn_scalar": 16,
+    "rms_norm_eps": 1e-06,
+    "rope_local_base_freq": 10000,
+    "rope_scaling": {
+      "factor": 8.0,
+      "rope_type": "linear"
+    },
+    "rope_theta": 1000000,
+    "sliding_window": 1024,
+    "sliding_window_pattern": 6,
+    "use_bidirectional_attention": false,
+    "use_cache": true,
+    "vocab_size": 262208
+  },
+  "transformers_version": "4.57.3",
+  "vision_config": {
+    "attention_dropout": 0.0,
+    "dtype": "float32",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 32,
+    "image_size": 896,
+    "intermediate_size": 64,
+    "layer_norm_eps": 1e-06,
+    "model_type": "siglip_vision_model",
+    "num_attention_heads": 2,
+    "num_channels": 3,
+    "num_hidden_layers": 2,
+    "patch_size": 14,
+    "vision_use_head": false
+  }
+}

tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/transformer/config.json

@@ -0,0 +1,44 @@
+{
+  "_class_name": "LTX2VideoTransformer3DModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "activation_fn": "gelu-approximate",
+  "attention_bias": true,
+  "attention_head_dim": 16,
+  "attention_out_bias": true,
+  "audio_attention_head_dim": 16,
+  "audio_cross_attention_dim": 32,
+  "audio_hop_length": 160,
+  "audio_in_channels": 128,
+  "audio_num_attention_heads": 2,
+  "audio_out_channels": 128,
+  "audio_patch_size": 1,
+  "audio_patch_size_t": 1,
+  "audio_pos_embed_max_pos": 20,
+  "audio_sampling_rate": 16000,
+  "audio_scale_factor": 4,
+  "base_height": 2048,
+  "base_width": 2048,
+  "caption_channels": 32,
+  "causal_offset": 1,
+  "cross_attention_dim": 32,
+  "cross_attn_timestep_scale_multiplier": 1000,
+  "in_channels": 128,
+  "norm_elementwise_affine": false,
+  "norm_eps": 1e-06,
+  "num_attention_heads": 2,
+  "num_layers": 2,
+  "out_channels": 128,
+  "patch_size": 1,
+  "patch_size_t": 1,
+  "pos_embed_max_pos": 20,
+  "qk_norm": "rms_norm_across_heads",
+  "rope_double_precision": true,
+  "rope_theta": 10000.0,
+  "rope_type": "split",
+  "timestep_scale_multiplier": 1000,
+  "vae_scale_factors": [
+    8,
+    32,
+    32
+  ]
+}

tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/vae/config.json

@@ -0,0 +1,81 @@
+{
+  "_class_name": "AutoencoderKLLTX2Video",
+  "_diffusers_version": "0.37.0.dev0",
+  "block_out_channels": [
+    256,
+    512,
+    1024,
+    2048
+  ],
+  "decoder_block_out_channels": [
+    256,
+    512,
+    1024
+  ],
+  "decoder_causal": false,
+  "decoder_inject_noise": [
+    false,
+    false,
+    false,
+    false
+  ],
+  "decoder_layers_per_block": [
+    1,
+    1,
+    1,
+    1
+  ],
+  "decoder_spatial_padding_mode": "reflect",
+  "decoder_spatio_temporal_scaling": [
+    true,
+    true,
+    true
+  ],
+  "down_block_types": [
+    "LTX2VideoDownBlock3D",
+    "LTX2VideoDownBlock3D",
+    "LTX2VideoDownBlock3D",
+    "LTX2VideoDownBlock3D"
+  ],
+  "downsample_type": [
+    "spatial",
+    "temporal",
+    "spatiotemporal",
+    "spatiotemporal"
+  ],
+  "encoder_causal": true,
+  "encoder_spatial_padding_mode": "zeros",
+  "in_channels": 3,
+  "latent_channels": 128,
+  "layers_per_block": [
+    1,
+    1,
+    1,
+    1,
+    1
+  ],
+  "out_channels": 3,
+  "patch_size": 4,
+  "patch_size_t": 1,
+  "resnet_norm_eps": 1e-06,
+  "scaling_factor": 1.0,
+  "spatial_compression_ratio": 32,
+  "spatio_temporal_scaling": [
+    true,
+    true,
+    true,
+    true
+  ],
+  "temporal_compression_ratio": 8,
+  "timestep_conditioning": false,
+  "upsample_factor": [
+    2,
+    2,
+    2
+  ],
+  "upsample_residual": [
+    true,
+    true,
+    true
+  ]
+}