diff --git a/tests/helpers/tiny_model.py b/tests/helpers/tiny_model.py new file mode 100644 index 000000000..b212a8e7f --- /dev/null +++ b/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 diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index 4120900fe..da599c2c2 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/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): diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 9bee71f9b..54e7a4016 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/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") diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 3f44afcae..c6c1cb1d9 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/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], + ], + ), } diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py index 6b9c7ff99..23ff60eeb 100644 --- a/tests/model_tests/diffusion/task_runners.py +++ b/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)) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 4bc740b01..9d0a4911a 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -30,7 +30,6 @@ EXCLUDED_MODELS = [ "OvisImagePipeline", "WanPipeline", "WanVACEPipeline", - "LTX2Pipeline", "LTX2ImageToVideoPipeline", "LTX2TwoStagesPipeline", "LTX2ImageToVideoTwoStagesPipeline", diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py index f24ed3780..d5cf3aba6 100644 --- a/tests/model_tests/diffusion/test_common_offline.py +++ b/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") diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index e4c6be0d4..31509c3b3 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/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") diff --git a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/model_index.json b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/model_index.json new file mode 100644 index 000000000..554a39bf1 --- /dev/null +++ b/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" + ] +} diff --git a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/text_encoder/config.json new file mode 100644 index 000000000..004914a7f --- /dev/null +++ b/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 +} diff --git a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/transformer/config.json b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/transformer/config.json new file mode 100644 index 000000000..6ff16e1e7 --- /dev/null +++ b/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 +} diff --git a/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/vae/config.json b/tests/model_tests/diffusion/tiny_configs/Flux2KleinPipeline/vae/config.json new file mode 100644 index 000000000..6410d25bd --- /dev/null +++ b/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 +} diff --git a/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/audio_vae/config.json b/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/audio_vae/config.json new file mode 100644 index 000000000..5ebd2f5e8 --- /dev/null +++ b/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 +} diff --git a/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/connectors/config.json b/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/connectors/config.json new file mode 100644 index 000000000..c63177c7c --- /dev/null +++ b/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 +} diff --git a/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/model_index.json b/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/model_index.json new file mode 100644 index 000000000..ef9de5d07 --- /dev/null +++ b/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" + ] +} diff --git a/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/text_encoder/config.json b/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/text_encoder/config.json new file mode 100644 index 000000000..0d722f0c1 --- /dev/null +++ b/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 + } +} diff --git a/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/transformer/config.json b/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/transformer/config.json new file mode 100644 index 000000000..577b0fb36 --- /dev/null +++ b/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 + ] +} diff --git a/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/vae/config.json b/tests/model_tests/diffusion/tiny_configs/LTX2Pipeline/vae/config.json new file mode 100644 index 000000000..a672971d4 --- /dev/null +++ b/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 + ] +}