[Model] Add native SANA-Video 2B T2V and I2V support (#5508)

Signed-off-by: leo <92263211+napleon-liu@users.noreply.github.com> Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com> Signed-off-by: liuyao0322 <yaoliu548926@gmail.com> Co-authored-by: leo <92263211+napleon-liu@users.noreply.github.com>
2026-08-25 09:46:02 +08:00
1 parent 4af0ddc1c2
commit 50b8a150f9
@@ -0,0 +1,992 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _DummyComponent:
+    pass
+
+
+def _make_request_batch(prompt, request_id="sana-video-test", **sampling_overrides):
+    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
+    request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=request_id)
+    return DiffusionRequestBatch([request])
+
+
+def test_sana_video_pipeline_import_and_registry():
+    from vllm_omni.diffusion.models.sana_video import (
+        SanaImageToVideoPipeline,
+        SanaVideoPipeline,
+        SanaVideoTransformer3DModel,
+        get_sana_video_i2v_post_process_func,
+        get_sana_video_i2v_pre_process_func,
+        get_sana_video_post_process_func,
+    )
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+        _DIFFUSION_PRE_PROCESS_FUNCS,
+    )
+
+    assert SanaImageToVideoPipeline is not None
+    assert SanaVideoPipeline is not None
+    assert SanaVideoTransformer3DModel is not None
+    assert get_sana_video_i2v_post_process_func is not None
+    assert get_sana_video_i2v_pre_process_func is not None
+    assert get_sana_video_post_process_func is not None
+    assert _DIFFUSION_MODELS["SanaVideoPipeline"] == (
+        "sana_video",
+        "pipeline_sana_video",
+        "SanaVideoPipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["SanaVideoPipeline"] == "get_sana_video_post_process_func"
+    assert _DIFFUSION_MODELS["SanaImageToVideoPipeline"] == (
+        "sana_video",
+        "pipeline_sana_video_i2v",
+        "SanaImageToVideoPipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_post_process_func"
+    assert _DIFFUSION_PRE_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_pre_process_func"
+
+
+def test_component_discovery_declarations():
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    assert SanaVideoPipeline._dit_modules == ["transformer"]
+    assert SanaVideoPipeline._encoder_modules == ["text_encoder"]
+    assert SanaVideoPipeline._vae_modules == ["vae"]
+    assert SanaVideoPipeline.supports_step_execution is False
+
+
+def test_sana_video_loads_concrete_gemma_tokenizer(monkeypatch):
+    from vllm_omni.diffusion.models.sana_video import pipeline_sana_video
+
+    captured: dict[str, object] = {}
+    expected = object()
+
+    def fake_from_pretrained_with_prefetch(loader, model, **kwargs):
+        captured.update(loader=loader, model=model, kwargs=kwargs)
+        return expected
+
+    monkeypatch.setattr(
+        pipeline_sana_video,
+        "from_pretrained_with_prefetch",
+        fake_from_pretrained_with_prefetch,
+    )
+
+    result = pipeline_sana_video._load_sana_tokenizer(
+        "Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+        ["tokenizer", "text_encoder"],
+        local_files_only=False,
+    )
+
+    assert result is expected
+    assert captured["loader"] == pipeline_sana_video.GemmaTokenizer.from_pretrained
+    assert captured["model"] == "Efficient-Large-Model/SANA-Video_2B_480p_diffusers"
+    assert captured["kwargs"] == {
+        "subfolder": "tokenizer",
+        "prefetch_list": ["tokenizer", "text_encoder"],
+        "local_files_only": False,
+    }
+
+
+def test_sana_video_declares_extra_body_params():
+    from vllm_omni.model_extras import get_extra_body_params
+
+    assert get_extra_body_params("SanaVideoPipeline") == {
+        "clean_caption",
+        "motion_score",
+        "use_resolution_binning",
+    }
+    assert get_extra_body_params("SanaImageToVideoPipeline") == {
+        "clean_caption",
+        "motion_score",
+        "use_resolution_binning",
+    }
+
+
+def test_sana_video_i2v_preprocesses_image_and_preserves_aspect_ratio():
+    from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func
+
+    request = OmniDiffusionRequest(
+        prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (640, 360))}},
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_id="sana-video-i2v-preprocess",
+    )
+    processed = get_sana_video_i2v_pre_process_func(SimpleNamespace())(request)
+
+    assert processed.sampling_params.height == 448
+    assert processed.sampling_params.width == 832
+    assert processed.prompt["multi_modal_data"]["image"].size == (832, 448)
+
+
+def test_sana_video_i2v_remote_720p_uses_loaded_transformer_config():
+    from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func
+
+    od_config = SimpleNamespace(
+        model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+        tf_model_config={"sample_size": 22},
+    )
+    request = OmniDiffusionRequest(
+        prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (1280, 704))}},
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_id="sana-video-i2v-remote-720p",
+    )
+
+    processed = get_sana_video_i2v_pre_process_func(od_config)(request)
+
+    assert processed.sampling_params.height == 704
+    assert processed.sampling_params.width == 1280
+    assert processed.prompt["multi_modal_data"]["image"].size == (1280, 704)
+
+
+@pytest.mark.parametrize(
+    "frames",
+    [
+        [[Image.new("RGB", (16, 16)) for _ in range(2)]],
+        np.zeros((1, 2, 16, 16, 3), dtype=np.float32),
+        [[np.zeros((16, 16, 3), dtype=np.float32) for _ in range(2)]],
+    ],
+)
+@pytest.mark.parametrize(
+    "factory_name",
+    ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"],
+)
+def test_sana_video_diffusers_preserves_postprocessed_frames(frames, factory_name):
+    from vllm_omni.diffusion.models import sana_video
+
+    post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers"))
+
+    result = post_process(frames)
+
+    assert result["payload"]["video"] is frames
+
+
+@pytest.mark.parametrize(
+    "factory_name",
+    ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"],
+)
+def test_sana_video_diffusers_still_postprocesses_native_tensor(factory_name):
+    from vllm_omni.diffusion.models import sana_video
+
+    decoded_video = torch.zeros(1, 3, 2, 16, 16)
+    post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers"))
+
+    result = post_process(decoded_video)
+
+    assert isinstance(result["payload"]["video"], np.ndarray)
+    assert result["payload"]["video"].shape == (1, 2, 16, 16, 3)
+
+
+@pytest.mark.parametrize(
+    "factory_name",
+    ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"],
+)
+def test_sana_video_postprocess_preserves_requested_latents(factory_name):
+    from vllm_omni.diffusion.models import sana_video
+
+    latents = torch.randn(1, 4, 2, 2, 2)
+    post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="native"))
+
+    result = post_process(latents, sampling_params=SimpleNamespace(output_type="latent"))
+
+    assert result is latents
+
+
+def test_sana_video_720p_model_id_fallback(monkeypatch):
+    from vllm_omni.diffusion.models.sana_video import pipeline_sana_video
+
+    monkeypatch.setattr(pipeline_sana_video, "get_hf_file_to_dict", lambda *_args, **_kwargs: None)
+    od_config = SimpleNamespace(
+        model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+        tf_model_config={},
+    )
+
+    assert pipeline_sana_video.resolve_sana_video_sample_size(od_config) == 22
+
+
+@pytest.mark.parametrize(
+    ("requested_output_type", "expected_internal_output_type"),
+    [
+        (None, "raw"),
+        ("np", "raw"),
+        ("latent", "latent"),
+    ],
+)
+def test_sana_video_i2v_forward_maps_image_request(requested_output_type, expected_internal_output_type):
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+    pipeline = object.__new__(SanaImageToVideoPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    calls = []
+
+    def fake_generate_i2v(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(kwargs["num_videos_per_prompt"], 3, 9, 192, 320)
+
+    pipeline._generate_i2v = fake_generate_i2v
+    image = Image.new("RGB", (320, 192))
+    custom_timesteps = torch.tensor([900, 100])
+    req = _make_request_batch(
+        {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}},
+        height=192,
+        width=320,
+        num_frames=9,
+        num_inference_steps=2,
+        timesteps=custom_timesteps,
+        guidance_scale=4.5,
+        num_outputs_per_prompt=2,
+        eta=0.25,
+        seed=42,
+        output_type=requested_output_type,
+        extra_args={"clean_caption": True, "motion_score": 30, "use_resolution_binning": False},
+    )
+
+    output = pipeline.forward(req)
+
+    assert output.output.shape == (2, 3, 9, 192, 320)
+    assert calls[0]["image"] is image
+    assert calls[0]["prompt"] == "a robot walks motion score: 30."
+    assert calls[0]["negative_prompt"] == "blurry"
+    assert calls[0]["frames"] == 9
+    assert calls[0]["timesteps"] is custom_timesteps
+    assert calls[0]["sigmas"] is None
+    assert calls[0]["num_videos_per_prompt"] == 2
+    assert calls[0]["eta"] == 0.25
+    assert calls[0]["clean_caption"] is True
+    assert calls[0]["generator"].initial_seed() == 42
+    assert calls[0]["output_type"] == expected_internal_output_type
+
+
+def test_sana_video_i2v_forward_maps_custom_sigmas():
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+    pipeline = object.__new__(SanaImageToVideoPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    calls = []
+
+    def fake_generate_i2v(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(1, 3, 9, 192, 320)
+
+    pipeline._generate_i2v = fake_generate_i2v
+    sigmas = [1.0, 0.5]
+    output = pipeline.forward(
+        _make_request_batch(
+            {
+                "prompt": "a robot walks",
+                "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
+            },
+            height=192,
+            width=320,
+            num_frames=9,
+            sigmas=sigmas,
+            generator_device="cpu",
+            extra_args={"use_resolution_binning": False},
+        )
+    )
+
+    assert output.output.shape == (1, 3, 9, 192, 320)
+    assert calls[0]["timesteps"] is None
+    assert calls[0]["sigmas"] is sigmas
+
+
+def test_sana_video_i2v_omitted_num_frames_uses_model_default():
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+    pipeline = object.__new__(SanaImageToVideoPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    calls = []
+
+    def fake_generate_i2v(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(1, 3, 81, 2, 2)
+
+    pipeline._generate_i2v = fake_generate_i2v
+    req = _make_request_batch(
+        {
+            "prompt": "a robot walks",
+            "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
+        },
+        height=192,
+        width=320,
+        num_inference_steps=1,
+        generator_device="cpu",
+        extra_args={"use_resolution_binning": False},
+    )
+
+    assert req.sampling_params.num_frames == 1
+    output = pipeline.forward(req)
+
+    assert output.output.shape == (1, 3, 81, 2, 2)
+    assert calls[0]["frames"] == 81
+
+
+def test_sana_video_i2v_clean_caption_request_is_explicitly_rejected():
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+    pipeline = object.__new__(SanaImageToVideoPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    pipeline.text_encoder = SimpleNamespace(dtype=torch.float32)
+    pipeline.tokenizer = None
+    pipeline.check_inputs = lambda **_kwargs: None
+
+    req = _make_request_batch(
+        {
+            "prompt": "a robot walks",
+            "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
+        },
+        height=192,
+        width=320,
+        num_frames=9,
+        generator_device="cpu",
+        extra_args={"clean_caption": True, "use_resolution_binning": False},
+    )
+
+    with pytest.raises(ValueError, match="does not support `clean_caption=True`"):
+        pipeline.forward(req)
+
+
+@pytest.mark.parametrize(
+    ("custom_timesteps", "custom_sigmas"),
+    [
+        (torch.tensor([4.0, 2.0]), None),
+        (None, [1.0, 0.5]),
+    ],
+)
+def test_sana_video_i2v_consumes_batch_schedule_and_step_kwargs(custom_timesteps, custom_sigmas):
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+    calls = SimpleNamespace(
+        encode_prompt=[],
+        prepare_latents=[],
+        set_timesteps=[],
+        step=[],
+        transformer=[],
+    )
+
+    class StubTransformer:
+        dtype = torch.float32
+        config = SimpleNamespace(
+            sample_size=30,
+            in_channels=4,
+            out_channels=4,
+            patch_size=(1, 1, 1),
+        )
+
+        def __call__(
+            self,
+            hidden_states,
+            *,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            timestep,
+            return_dict,
+        ):
+            calls.transformer.append(
+                (
+                    hidden_states.shape,
+                    encoder_hidden_states.shape,
+                    encoder_attention_mask.shape,
+                    timestep.shape,
+                    return_dict,
+                )
+            )
+            return (torch.zeros_like(hidden_states),)
+
+    class StubScheduler:
+        order = 1
+
+        def set_timesteps(
+            self,
+            num_inference_steps=None,
+            *,
+            device=None,
+            timesteps=None,
+            sigmas=None,
+        ):
+            calls.set_timesteps.append(
+                {
+                    "num_inference_steps": num_inference_steps,
+                    "timesteps": timesteps,
+                    "sigmas": sigmas,
+                }
+            )
+            schedule_length = len(timesteps if timesteps is not None else sigmas)
+            self.timesteps = torch.arange(schedule_length, 0, -1, dtype=torch.float32, device=device)
+
+        def step(
+            self,
+            _noise_pred,
+            _timestep,
+            current_latents,
+            *,
+            eta,
+            generator,
+            return_dict,
+        ):
+            calls.step.append({"eta": eta, "generator": generator, "return_dict": return_dict})
+            return (current_latents + 1,)
+
+    class VaeMustNotBeUsed:
+        @property
+        def dtype(self):
+            pytest.fail("latent output must return before accessing the VAE")
+
+        def decode(self, *_args, **_kwargs):
+            pytest.fail("latent output must not call VAE decode")
+
+    pipeline = object.__new__(SanaImageToVideoPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = StubTransformer()
+    pipeline.scheduler = StubScheduler()
+    pipeline.vae = VaeMustNotBeUsed()
+    pipeline.check_inputs = lambda **_kwargs: None
+    pipeline.video_processor = SimpleNamespace(
+        preprocess=lambda _image, height, width: torch.zeros(1, 3, height, width),
+    )
+    pipeline.progress_bar = lambda **_kwargs: nullcontext(SimpleNamespace(update=lambda: None))
+
+    def capture_encode_prompt(*_args, **kwargs):
+        calls.encode_prompt.append(kwargs)
+        batch_size = kwargs["num_videos_per_prompt"]
+        return (
+            torch.zeros(batch_size, 1, 1),
+            torch.ones(batch_size, 1, dtype=torch.bool),
+            torch.zeros(batch_size, 1, 1),
+            torch.ones(batch_size, 1, dtype=torch.bool),
+        )
+
+    pipeline.encode_prompt = capture_encode_prompt
+    initial_latents = torch.zeros(2, 4, 2, 2, 2)
+
+    def capture_prepare_latents(*_args, **kwargs):
+        calls.prepare_latents.append(kwargs)
+        return initial_latents
+
+    pipeline._prepare_i2v_latents = capture_prepare_latents
+    generator = torch.Generator(device="cpu").manual_seed(42)
+
+    output = pipeline._generate_i2v(
+        image=Image.new("RGB", (2, 2)),
+        prompt="test",
+        negative_prompt="",
+        height=2,
+        width=2,
+        frames=2,
+        num_inference_steps=50,
+        timesteps=custom_timesteps,
+        sigmas=custom_sigmas,
+        guidance_scale=4.5,
+        num_videos_per_prompt=2,
+        eta=0.25,
+        generator=generator,
+        latents=None,
+        clean_caption=True,
+        use_resolution_binning=False,
+        max_sequence_length=1,
+        output_type="latent",
+    )
+
+    assert calls.encode_prompt[0]["num_videos_per_prompt"] == 2
+    assert calls.encode_prompt[0]["clean_caption"] is True
+    assert calls.prepare_latents[0]["batch_size"] == 2
+    assert len(calls.set_timesteps) == 1
+    assert calls.set_timesteps[0]["num_inference_steps"] is None
+    assert calls.set_timesteps[0]["timesteps"] is custom_timesteps
+    assert calls.set_timesteps[0]["sigmas"] is custom_sigmas
+    assert len(calls.step) == 2
+    assert all(call == {"eta": 0.25, "generator": generator, "return_dict": False} for call in calls.step)
+    assert all(
+        call
+        == (
+            torch.Size([4, 4, 2, 2, 2]),
+            torch.Size([4, 1, 1]),
+            torch.Size([4, 1]),
+            torch.Size([4, 1, 2, 2, 2]),
+            False,
+        )
+        for call in calls.transformer
+    )
+    expected = initial_latents.clone()
+    expected[:, :, 1:] += 2
+    torch.testing.assert_close(output, expected)
+
+
+def test_sana_video_i2v_uses_diffusers_complex_instruction_default():
+    import inspect
+
+    from diffusers import SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline
+
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+    reference_default = (
+        inspect.signature(DiffusersSanaImageToVideoPipeline.__call__).parameters["complex_human_instruction"].default
+    )
+    native_default = (
+        inspect.signature(SanaImageToVideoPipeline._generate_i2v).parameters["complex_human_instruction"].default
+    )
+    assert native_default == reference_default
+
+    class StopAfterPromptEncodingError(Exception):
+        pass
+
+    pipeline = object.__new__(SanaImageToVideoPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.check_inputs = lambda **_kwargs: None
+    captured = {}
+
+    def capture_encode_prompt(*_args, **kwargs):
+        captured.update(kwargs)
+        raise StopAfterPromptEncodingError
+
+    pipeline.encode_prompt = capture_encode_prompt
+    with pytest.raises(StopAfterPromptEncodingError):
+        pipeline._generate_i2v(
+            image=Image.new("RGB", (320, 192)),
+            prompt="a robot walks",
+            negative_prompt="blurry",
+            height=192,
+            width=320,
+            frames=9,
+            num_inference_steps=2,
+            guidance_scale=4.5,
+            generator=torch.Generator(device="cpu").manual_seed(42),
+            latents=None,
+            use_resolution_binning=False,
+            max_sequence_length=300,
+        )
+
+    assert captured["complex_human_instruction"] == reference_default
+
+
+def test_sana_video_t2v_complex_instruction_has_no_mutable_default():
+    import inspect
+
+    from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
+        SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION,
+        SanaVideoPipeline,
+    )
+
+    parameter = inspect.signature(SanaVideoPipeline._generate).parameters["complex_human_instruction"]
+    assert parameter.default is None
+    assert isinstance(SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION, tuple)
+
+
+def test_diffusers_adapter_t2v_omitted_num_frames_uses_sana_default():
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+    adapter = object.__new__(DiffusersAdapterPipeline)
+    adapter._accept_call_kwargs = {"prompt", "negative_prompt", "frames", "generator"}
+    adapter._pipeline_utils = SanaVideoPipelineUtils()
+    adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None)
+
+    req = _make_request_batch(
+        {"prompt": "a robot walks", "negative_prompt": "blurry"},
+        seed=42,
+        generator_device="cpu",
+    )
+    assert req.sampling_params.num_frames == 1
+
+    kwargs = adapter._build_call_kwargs(req)
+
+    assert kwargs["frames"] == 81
+    assert "num_frames" not in kwargs
+    assert kwargs["generator"].initial_seed() == 42
+
+
+def test_diffusers_adapter_i2v_omitted_num_frames_uses_sana_default():
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+    adapter = object.__new__(DiffusersAdapterPipeline)
+    adapter._accept_call_kwargs = {"prompt", "negative_prompt", "image", "frames", "generator"}
+    adapter._pipeline_utils = SanaVideoPipelineUtils()
+    adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None)
+    image = Image.new("RGB", (320, 192))
+
+    req = _make_request_batch(
+        {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}},
+        seed=42,
+        generator_device="cpu",
+    )
+    assert req.sampling_params.num_frames == 1
+
+    kwargs = adapter._build_call_kwargs(req)
+
+    assert kwargs["image"] is image
+    assert kwargs["prompt"] == "a robot walks"
+    assert kwargs["negative_prompt"] == "blurry"
+    assert kwargs["frames"] == 81
+    assert kwargs["generator"].initial_seed() == 42
+
+
+@pytest.mark.parametrize(
+    ("sampling_overrides", "expected_frames"),
+    [
+        ({}, 49),
+        ({"num_frames": 9}, 9),
+    ],
+)
+def test_diffusers_adapter_sana_resolves_frame_precedence(sampling_overrides, expected_frames):
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+    adapter = object.__new__(DiffusersAdapterPipeline)
+    adapter._accept_call_kwargs = {"prompt", "frames", "generator"}
+    adapter._pipeline_utils = SanaVideoPipelineUtils()
+    adapter.od_config = SimpleNamespace(diffusers_call_kwargs={"frames": 49}, output_type=None)
+
+    kwargs = adapter._build_call_kwargs(
+        _make_request_batch(
+            "a robot walks",
+            seed=42,
+            generator_device="cpu",
+            **sampling_overrides,
+        )
+    )
+
+    assert kwargs["frames"] == expected_frames
+
+
+@pytest.mark.parametrize("pipeline_class_name", ["SanaVideoPipeline", "SanaImageToVideoPipeline"])
+def test_diffusers_adapter_selects_sana_pipeline_utils(pipeline_class_name):
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import (
+        SanaVideoPipelineUtils,
+        get_pipeline_utils,
+    )
+
+    assert isinstance(get_pipeline_utils(pipeline_class_name), SanaVideoPipelineUtils)
+
+
+def test_diffusers_adapter_resolves_requested_sana_i2v_pipeline():
+    from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline
+
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import (
+        SanaVideoPipelineUtils,
+        get_pipeline_utils_for_config,
+        resolve_diffusers_pipeline_class,
+    )
+
+    od_config = SimpleNamespace(
+        model_class_name="SanaImageToVideoPipeline",
+        diffusers_pipeline_cls=SanaVideoPipeline,
+    )
+
+    assert isinstance(get_pipeline_utils_for_config(od_config), SanaVideoPipelineUtils)
+    assert resolve_diffusers_pipeline_class(od_config) is SanaImageToVideoPipeline
+
+
+def test_diffusers_adapter_pipeline_override_is_sana_specific():
+    from diffusers import DiffusionPipeline
+
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import resolve_diffusers_pipeline_class
+
+    class OtherPipeline(DiffusionPipeline):
+        pass
+
+    od_config = SimpleNamespace(
+        model_class_name="DiffusersAdapterPipeline",
+        diffusers_pipeline_cls=OtherPipeline,
+    )
+
+    assert resolve_diffusers_pipeline_class(od_config) is OtherPipeline
+
+
+def test_diffusers_adapter_loads_resolved_pipeline_class(monkeypatch):
+    from vllm_omni.diffusion.models.diffusers_adapter import pipeline_diffusers_adapter
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+
+    loaded = []
+    expected_pipeline = object()
+
+    class ResolvedPipeline:
+        @classmethod
+        def from_pretrained(cls, model_id, **kwargs):
+            loaded.append((model_id, kwargs))
+            return expected_pipeline
+
+    monkeypatch.setattr(
+        pipeline_diffusers_adapter,
+        "resolve_diffusers_pipeline_class",
+        lambda _od_config: ResolvedPipeline,
+    )
+    adapter = object.__new__(DiffusersAdapterPipeline)
+    adapter.od_config = SimpleNamespace(diffusers_pipeline_cls=object)
+
+    result = adapter._load_pipeline_from_pretrained("sana-checkpoint", {"torch_dtype": torch.bfloat16})
+
+    assert result is expected_pipeline
+    assert loaded == [("sana-checkpoint", {"torch_dtype": torch.bfloat16})]
+
+
+def test_diffusers_adapter_sana_i2v_override_reports_image_capability():
+    from diffusers import SanaVideoPipeline
+
+    from vllm_omni.diffusion.io_support import supports_multimodal_input
+
+    od_config = SimpleNamespace(
+        diffusion_load_format="diffusers",
+        model_class_name="SanaImageToVideoPipeline",
+        diffusers_pipeline_cls=SanaVideoPipeline,
+    )
+
+    assert supports_multimodal_input(od_config) == (True, False)
+
+
+def test_diffusers_adapter_sana_i2v_warmup_supplies_image():
+    from diffusers import SanaVideoPipeline
+
+    from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        diffusion_load_format="diffusers",
+        model_class_name="SanaImageToVideoPipeline",
+        diffusers_pipeline_cls=SanaVideoPipeline,
+    )
+    captured_requests = []
+    engine.pre_process_func = lambda request: request
+
+    def capture_request(request):
+        captured_requests.append(request)
+        return SimpleNamespace(error=None)
+
+    engine.add_req_and_wait_for_response = capture_request
+
+    engine._dummy_run()
+
+    assert len(captured_requests) == 1
+    prompt = captured_requests[0].prompt
+    assert isinstance(prompt, dict)
+    assert isinstance(prompt["multi_modal_data"]["image"], Image.Image)
+
+
+def test_diffusers_adapter_disables_resolution_binning_for_warmup():
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+    from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+    adapter = object.__new__(DiffusersAdapterPipeline)
+    adapter._accept_call_kwargs = {"prompt", "frames", "use_resolution_binning", "generator"}
+    adapter._pipeline_utils = SanaVideoPipelineUtils()
+    adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None)
+
+    kwargs = adapter._build_call_kwargs(
+        _make_request_batch(
+            "dummy run",
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
+            generator_device="cpu",
+        )
+    )
+
+    assert kwargs["use_resolution_binning"] is False
+    assert kwargs["frames"] == 1
+
+
+def test_pipeline_is_torch_module_and_supports_eval():
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    nn.Module.__init__(pipeline)
+
+    assert isinstance(pipeline, nn.Module)
+    assert pipeline.eval() is pipeline
+    assert pipeline.training is False
+
+
+@pytest.mark.parametrize(
+    ("vae_scale", "patch_size", "valid_size", "invalid_size", "alignment"),
+    [
+        (8, (1, 2, 2), 624, 632, 16),
+        (32, (1, 1, 1), 704, 712, 32),
+    ],
+)
+def test_check_inputs_uses_variant_spatial_alignment(vae_scale, patch_size, valid_size, invalid_size, alignment):
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    pipeline.vae_scale_factor_spatial = vae_scale
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(patch_size=patch_size))
+
+    pipeline.check_inputs(prompt="test", height=valid_size, width=valid_size)
+    with pytest.raises(ValueError, match=f"divisible by {alignment}"):
+        pipeline.check_inputs(prompt="test", height=invalid_size, width=valid_size)
+
+
+@pytest.mark.parametrize(
+    ("requested_output_type", "expected_internal_output_type"),
+    [
+        (None, "raw"),
+        ("np", "raw"),
+        ("latent", "latent"),
+    ],
+)
+def test_forward_maps_omni_request_to_sana_generation_args(requested_output_type, expected_internal_output_type):
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(frames=torch.zeros(1, 3, 9, 192, 320))
+
+    pipeline._generate = fake_generate
+    req = _make_request_batch(
+        {"prompt": "a robot walks", "negative_prompt": "blurry"},
+        height=192,
+        width=320,
+        num_frames=9,
+        num_inference_steps=2,
+        guidance_scale=4.5,
+        seed=42,
+        output_type=requested_output_type,
+        extra_args={"motion_score": 30, "use_resolution_binning": False},
+    )
+
+    output = pipeline.forward(req)
+
+    assert output.output.shape == (1, 3, 9, 192, 320)
+    assert calls[0]["prompt"] == "a robot walks motion score: 30."
+    assert calls[0]["negative_prompt"] == "blurry"
+    assert calls[0]["height"] == 192
+    assert calls[0]["width"] == 320
+    assert calls[0]["frames"] == 9
+    assert calls[0]["num_inference_steps"] == 2
+    assert calls[0]["guidance_scale"] == 4.5
+    assert calls[0]["use_resolution_binning"] is False
+    assert calls[0]["generator"].initial_seed() == 42
+    assert calls[0]["output_type"] == expected_internal_output_type
+
+
+def test_sana_video_t2v_omitted_num_frames_uses_model_default():
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(frames=torch.zeros(1, 3, 81, 2, 2))
+
+    pipeline._generate = fake_generate
+    req = _make_request_batch(
+        "a robot walks",
+        height=192,
+        width=320,
+        num_inference_steps=1,
+        generator_device="cpu",
+    )
+
+    assert req.sampling_params.num_frames == 1
+    output = pipeline.forward(req)
+
+    assert output.output.shape == (1, 3, 81, 2, 2)
+    assert calls[0]["frames"] == 81
+
+
+def test_sana_video_t2v_dummy_run_keeps_single_frame():
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(frames=torch.zeros(1, 3, 1, 192, 320))
+
+    pipeline._generate = fake_generate
+
+    output = pipeline.forward(
+        _make_request_batch(
+            "dummy run",
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
+            height=192,
+            width=320,
+            num_frames=1,
+            num_inference_steps=1,
+            generator_device="cpu",
+        )
+    )
+
+    assert output.output.shape == (1, 3, 1, 192, 320)
+    assert calls[0]["frames"] == 1
+
+
+@pytest.mark.parametrize(
+    ("sample_size", "expected_height", "expected_width"),
+    [
+        (30, 480, 832),
+        (22, 704, 1280),
+    ],
+)
+def test_t2v_uses_variant_default_resolution(sample_size, expected_height, expected_width):
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=sample_size))
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(frames=torch.zeros(1, 3, 9, expected_height, expected_width))
+
+    pipeline._generate = fake_generate
+    output = pipeline.forward(
+        _make_request_batch(
+            "a robot walks",
+            num_frames=9,
+            num_inference_steps=1,
+            guidance_scale=1.0,
+            seed=42,
+            generator_device="cpu",
+        )
+    )
+
+    assert output.output.shape == (1, 3, 9, expected_height, expected_width)
+    assert calls[0]["height"] == expected_height
+    assert calls[0]["width"] == expected_width
+
+
+def test_forward_requires_exactly_one_nonempty_prompt():
+    from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+    pipeline = object.__new__(SanaVideoPipeline)
+
+    with pytest.raises(ValueError, match="Prompt is required"):
+        pipeline.forward(_make_request_batch(""))
+
+    first = OmniDiffusionRequest(
+        prompt="first",
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_id="first",
+    )
+    second = OmniDiffusionRequest(
+        prompt="second",
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_id="second",
+    )
+    with pytest.raises(ValueError, match="exactly one prompt"):
+        pipeline.forward(DiffusionRequestBatch([first, second]))
@@ -0,0 +1,602 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_TINY_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 12,
+    "num_layers": 1,
+    "num_cross_attention_heads": 2,
+    "cross_attention_head_dim": 12,
+    "cross_attention_dim": 24,
+    "caption_channels": 8,
+    "mlp_ratio": 2.0,
+    "patch_size": (1, 2, 2),
+    "sample_size": 4,
+    "rope_max_seq_len": 32,
+}
+
+_GOLDEN_PREFIX = torch.tensor(
+    [
+        -0.1159307063,
+        0.2003862262,
+        -0.3754127920,
+        -0.2403523028,
+        -1.2376221418,
+        -0.6513092518,
+        -0.3358457983,
+        -0.9894289970,
+        0.7858567238,
+        0.5089095831,
+        -0.5362522602,
+        -0.3732060790,
+    ]
+)
+
+
+@pytest.mark.parametrize(
+    ("elementwise_affine", "bias", "expected_state_keys"),
+    [
+        (False, False, set()),
+        (False, True, set()),
+        (True, False, {"weight"}),
+        (True, True, {"weight", "bias"}),
+    ],
+)
+@pytest.mark.parametrize(
+    ("input_dtype", "parameter_dtype"),
+    [
+        (torch.float16, torch.float16),
+        (torch.bfloat16, torch.bfloat16),
+        (torch.float32, torch.float32),
+        (torch.float64, torch.float64),
+        (torch.float16, torch.float32),
+        (torch.bfloat16, torch.float32),
+        (torch.float32, torch.float16),
+        (torch.float32, torch.bfloat16),
+    ],
+)
+def test_sana_rms_norm_matches_diffusers(
+    elementwise_affine,
+    bias,
+    expected_state_keys,
+    input_dtype,
+    parameter_dtype,
+):
+    from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaRMSNorm
+
+    reference = DiffusersRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to(
+        dtype=parameter_dtype
+    )
+    actual = SanaRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to(dtype=parameter_dtype)
+
+    assert actual.eps == reference.eps
+    assert actual.elementwise_affine == reference.elementwise_affine
+    assert actual.dim == reference.dim
+    assert set(actual.state_dict()) == expected_state_keys
+    assert set(actual.state_dict()) == set(reference.state_dict())
+
+    if elementwise_affine:
+        weight = torch.linspace(0.5, 1.5, 8, dtype=parameter_dtype)
+        actual.weight.data.copy_(weight)
+        reference.weight.data.copy_(weight)
+        if bias:
+            norm_bias = torch.linspace(-0.25, 0.25, 8, dtype=parameter_dtype)
+            actual.bias.data.copy_(norm_bias)
+            reference.bias.data.copy_(norm_bias)
+    else:
+        assert actual.weight is None
+        assert actual.bias is None
+
+    hidden_states = torch.linspace(-2.0, 2.0, 48, dtype=input_dtype).reshape(2, 3, 8)
+    expected = reference(hidden_states)
+    result = actual(hidden_states)
+
+    assert result.dtype == expected.dtype
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    ("num_channels", "flip_sin_to_cos", "downscale_freq_shift", "scale"),
+    [
+        (256, True, 0, 1),
+        (7, False, 1, 0.5),
+        (12, True, 0.5, 2),
+    ],
+)
+@pytest.mark.parametrize("timestep_dtype", [torch.int64, torch.float32, torch.float64])
+def test_sana_timesteps_matches_diffusers(
+    num_channels,
+    flip_sin_to_cos,
+    downscale_freq_shift,
+    scale,
+    timestep_dtype,
+):
+    from diffusers.models.embeddings import Timesteps
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaTimesteps
+
+    reference = Timesteps(
+        num_channels=num_channels,
+        flip_sin_to_cos=flip_sin_to_cos,
+        downscale_freq_shift=downscale_freq_shift,
+        scale=scale,
+    )
+    actual = SanaTimesteps(
+        num_channels=num_channels,
+        flip_sin_to_cos=flip_sin_to_cos,
+        downscale_freq_shift=downscale_freq_shift,
+        scale=scale,
+    )
+    timesteps = torch.tensor([0, 1, 500, 999], dtype=timestep_dtype)
+
+    assert actual.state_dict() == reference.state_dict() == {}
+    assert actual.num_channels == reference.num_channels
+    assert actual.flip_sin_to_cos == reference.flip_sin_to_cos
+    assert actual.downscale_freq_shift == reference.downscale_freq_shift
+    assert actual.scale == reference.scale
+    torch.testing.assert_close(actual(timesteps), reference(timesteps), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "sample_shape", "condition_shape"),
+    [
+        ({}, (2, 256), None),
+        (
+            {
+                "out_dim": 12,
+                "post_act_fn": "silu",
+                "cond_proj_dim": 6,
+                "sample_proj_bias": False,
+            },
+            (2, 256),
+            (2, 6),
+        ),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_sana_timestep_embedding_matches_diffusers(
+    kwargs,
+    sample_shape,
+    condition_shape,
+    dtype,
+):
+    from diffusers.models.embeddings import TimestepEmbedding
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaTimestepEmbedding,
+    )
+
+    torch.manual_seed(17)
+    reference = TimestepEmbedding(256, 24, **kwargs).to(dtype=dtype)
+    actual = SanaTimestepEmbedding(256, 24, **kwargs).to(dtype=dtype)
+    assert set(actual.state_dict()) == set(reference.state_dict())
+    actual.load_state_dict(reference.state_dict())
+
+    sample = torch.randn(sample_shape, dtype=dtype)
+    condition = torch.randn(condition_shape, dtype=dtype) if condition_shape else None
+    expected = reference(sample, condition)
+    result = actual(sample, condition)
+
+    assert result.shape == expected.shape
+    assert result.dtype == expected.dtype
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_sana_timestep_embedding_rejects_unsupported_activation():
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaTimestepEmbedding,
+    )
+
+    with pytest.raises(ValueError, match="act_fn='silu'"):
+        SanaTimestepEmbedding(8, 16, act_fn="gelu")
+    with pytest.raises(ValueError, match="post_act_fn"):
+        SanaTimestepEmbedding(8, 16, post_act_fn="gelu")
+
+    module = SanaTimestepEmbedding(8, 16)
+    with pytest.raises(ValueError, match="cond_proj_dim"):
+        module(torch.randn(1, 8), torch.randn(1, 4))
+
+
+@pytest.mark.parametrize("use_additional_conditions", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_sana_ada_layer_norm_single_matches_diffusers(
+    use_additional_conditions,
+    dtype,
+):
+    from diffusers.models.normalization import AdaLayerNormSingle
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaAdaLayerNormSingle,
+    )
+
+    torch.manual_seed(19)
+    reference = AdaLayerNormSingle(
+        24,
+        use_additional_conditions=use_additional_conditions,
+    ).to(dtype=dtype)
+    actual = SanaAdaLayerNormSingle(
+        24,
+        use_additional_conditions=use_additional_conditions,
+    ).to(dtype=dtype)
+    assert set(actual.state_dict()) == set(reference.state_dict())
+    actual.load_state_dict(reference.state_dict())
+
+    timestep = torch.tensor([10, 500])
+    added_cond_kwargs = None
+    if use_additional_conditions:
+        added_cond_kwargs = {
+            "resolution": torch.tensor([[480, 832], [704, 1280]]),
+            "aspect_ratio": torch.tensor([[832 / 480], [1280 / 704]]),
+        }
+
+    expected = reference(
+        timestep,
+        added_cond_kwargs=added_cond_kwargs,
+        batch_size=2,
+        hidden_dtype=dtype,
+    )
+    result = actual(
+        timestep,
+        added_cond_kwargs=added_cond_kwargs,
+        batch_size=2,
+        hidden_dtype=dtype,
+    )
+
+    assert len(result) == len(expected) == 2
+    for actual_tensor, expected_tensor in zip(result, expected):
+        assert actual_tensor.shape == expected_tensor.shape
+        assert actual_tensor.dtype == expected_tensor.dtype
+        torch.testing.assert_close(actual_tensor, expected_tensor, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("act_fn", ["gelu_tanh", "silu", "silu_fp32"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_sana_pixart_text_projection_matches_diffusers(act_fn, dtype):
+    from diffusers.models.embeddings import PixArtAlphaTextProjection
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaPixArtAlphaTextProjection,
+    )
+
+    torch.manual_seed(23)
+    reference = PixArtAlphaTextProjection(
+        in_features=8,
+        hidden_size=24,
+        out_features=12,
+        act_fn=act_fn,
+    ).to(dtype=dtype)
+    actual = SanaPixArtAlphaTextProjection(
+        in_features=8,
+        hidden_size=24,
+        out_features=12,
+        act_fn=act_fn,
+    ).to(dtype=dtype)
+    assert set(actual.state_dict()) == set(reference.state_dict())
+    actual.load_state_dict(reference.state_dict())
+
+    caption = torch.randn(2, 5, 8, dtype=dtype)
+    expected = reference(caption)
+    result = actual(caption)
+
+    assert result.shape == expected.shape
+    assert result.dtype == expected.dtype
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("freqs_dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize(("dim", "max_seq_len", "theta"), [(4, 8, 10000.0), (12, 32, 256.0)])
+def test_native_rope_matches_diffusers(dim, max_seq_len, theta, freqs_dtype):
+    from diffusers.models.embeddings import get_1d_rotary_pos_embed
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        _get_1d_rotary_pos_embed,
+    )
+
+    expected_cos, expected_sin = get_1d_rotary_pos_embed(
+        dim,
+        max_seq_len,
+        theta,
+        use_real=True,
+        repeat_interleave_real=True,
+        freqs_dtype=freqs_dtype,
+    )
+    actual_cos, actual_sin = _get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype)
+
+    torch.testing.assert_close(actual_cos, expected_cos, rtol=0, atol=0)
+    torch.testing.assert_close(actual_sin, expected_sin, rtol=0, atol=0)
+
+
+def test_tiny_transformer_matches_diffusers_and_frozen_output():
+    from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer
+    from diffusers.configuration_utils import ConfigMixin
+    from diffusers.models.modeling_utils import ModelMixin
+
+    from vllm_omni.diffusion.attention.layer import Attention
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaLinearAttention,
+        SanaRMSNorm,
+        SanaVideoTransformerConfig,
+        SanaVideoTransformerOutput,
+    )
+
+    torch.manual_seed(7)
+    reference = DiffusersTransformer(**_TINY_CONFIG).eval()
+    model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+    assert not isinstance(model, ModelMixin)
+    assert not isinstance(model, ConfigMixin)
+    assert isinstance(model.config, SanaVideoTransformerConfig)
+    assert set(model.state_dict()) == set(reference.state_dict())
+    model.load_state_dict(reference.state_dict())
+
+    block = model.transformer_blocks[0]
+    assert isinstance(block.attn1, SanaLinearAttention)
+    assert isinstance(block.attn2.attn, Attention)
+    assert block.attn2.attn.role == "cross"
+    assert isinstance(block.attn1.norm_q, SanaRMSNorm)
+    assert isinstance(block.attn1.norm_k, SanaRMSNorm)
+    assert isinstance(block.attn2.norm_q, SanaRMSNorm)
+    assert isinstance(block.attn2.norm_k, SanaRMSNorm)
+    assert isinstance(model.caption_norm, SanaRMSNorm)
+
+    torch.manual_seed(11)
+    hidden_states = torch.randn(1, 4, 3, 4, 4)
+    encoder_hidden_states = torch.randn(1, 5, 8)
+    encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]])
+    timestep = torch.tensor([500.0])
+
+    with torch.no_grad():
+        expected = reference(
+            hidden_states,
+            encoder_hidden_states,
+            timestep,
+            encoder_attention_mask=encoder_attention_mask,
+        ).sample
+        actual_output = model(
+            hidden_states,
+            encoder_hidden_states,
+            timestep,
+            encoder_attention_mask=encoder_attention_mask,
+        )
+        actual = actual_output.sample
+
+    assert isinstance(actual_output, SanaVideoTransformerOutput)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual.flatten()[: len(_GOLDEN_PREFIX)], _GOLDEN_PREFIX, rtol=1e-5, atol=1e-5)
+
+
+def test_encoder_attention_mask_uses_bool_sdpa_contract(monkeypatch):
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    torch.manual_seed(37)
+    model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+    captured = {}
+    omni_attention = model.transformer_blocks[0].attn2.attn
+
+    def fake_sdpa_forward(query, key, value, attn_metadata):
+        captured["query_shape"] = query.shape
+        captured["key_shape"] = key.shape
+        captured["mask"] = attn_metadata.attn_mask.detach().clone()
+        return torch.zeros_like(query)
+
+    def unexpected_backend_forward(*args, **kwargs):
+        raise AssertionError("masked SANA cross-attention must use SDPA")
+
+    monkeypatch.setattr(omni_attention.sdpa_fallback, "forward", fake_sdpa_forward)
+    monkeypatch.setattr(omni_attention, "forward", unexpected_backend_forward)
+
+    hidden_states = torch.randn(2, 4, 3, 4, 4)
+    encoder_hidden_states = torch.randn(2, 5, 8)
+    input_mask = torch.tensor(
+        [
+            [1, 1, 1, 0, 0],
+            [1, 1, 1, 1, 0],
+        ],
+        dtype=torch.int64,
+    )
+
+    with torch.no_grad():
+        output = model(
+            hidden_states,
+            encoder_hidden_states,
+            torch.tensor([100.0, 500.0]),
+            encoder_attention_mask=input_mask,
+        ).sample
+
+    assert output.shape == hidden_states.shape
+    assert captured["query_shape"][:2] == (2, 12)
+    assert captured["key_shape"][:2] == (2, 5)
+    assert captured["query_shape"][1] != captured["key_shape"][1]
+    assert captured["mask"].shape == (2, 5)
+    assert captured["mask"].dtype == torch.bool
+    assert torch.equal(captured["mask"], input_mask.bool())
+
+
+def test_encoder_attention_mask_rejects_additive_bias():
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+    additive_mask = torch.tensor([[0.0, 0.0, 0.0, -10000.0, -10000.0]])
+
+    with pytest.raises(TypeError, match="not a floating-point additive attention bias"):
+        model(
+            torch.randn(1, 4, 3, 4, 4),
+            torch.randn(1, 5, 8),
+            torch.tensor([500.0]),
+            encoder_attention_mask=additive_mask,
+        )
+
+
+def test_linear_attention_requires_rotary_embeddings():
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention
+
+    attention = SanaLinearAttention(
+        dim=24,
+        num_heads=2,
+        head_dim=12,
+        dropout=0.0,
+        bias=True,
+        qk_norm="rms_norm_across_heads",
+    )
+
+    with pytest.raises(ValueError, match="requires rotary_emb"):
+        attention(torch.randn(1, 4, 24))
+
+
+def test_linear_attention_bfloat16_matches_diffusers_mixed_precision():
+    from diffusers.models.attention_processor import Attention as DiffusersAttention
+    from diffusers.models.transformers.transformer_sana_video import SanaLinearAttnProcessor3_0
+
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention
+
+    torch.manual_seed(17)
+    reference = DiffusersAttention(
+        query_dim=24,
+        heads=2,
+        dim_head=12,
+        kv_heads=2,
+        qk_norm="rms_norm_across_heads",
+        dropout=0.0,
+        bias=True,
+        cross_attention_dim=None,
+        processor=SanaLinearAttnProcessor3_0(),
+    ).to(dtype=torch.bfloat16)
+    actual = SanaLinearAttention(
+        dim=24,
+        num_heads=2,
+        head_dim=12,
+        dropout=0.0,
+        bias=True,
+        qk_norm="rms_norm_across_heads",
+    ).to(dtype=torch.bfloat16)
+    actual.load_state_dict(reference.state_dict())
+
+    hidden_states = torch.randn(1, 5, 24, dtype=torch.bfloat16)
+    rotary_emb = (
+        torch.randn(1, 5, 1, 12, dtype=torch.bfloat16),
+        torch.randn(1, 5, 1, 12, dtype=torch.bfloat16),
+    )
+
+    with torch.no_grad():
+        expected = reference(hidden_states, rotary_emb=rotary_emb)
+        result = actual(hidden_states, rotary_emb=rotary_emb)
+
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_transformer_block_without_cross_attention_keeps_feed_forward_norm():
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerBlock
+
+    block = SanaVideoTransformerBlock(
+        dim=24,
+        num_attention_heads=2,
+        attention_head_dim=12,
+        num_cross_attention_heads=2,
+        cross_attention_head_dim=12,
+        cross_attention_dim=None,
+        mlp_ratio=2.0,
+    )
+
+    assert block.attn2 is None
+    assert isinstance(block.norm2, torch.nn.LayerNorm)
+
+
+def test_guidance_transformer_matches_diffusers_and_tuple_output():
+    from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer
+
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaTimestepEmbedding,
+        SanaTimesteps,
+    )
+
+    config = _TINY_CONFIG | {"guidance_embeds": True}
+    torch.manual_seed(29)
+    reference = DiffusersTransformer(**config).eval()
+    model = SanaVideoTransformer3DModel(**config).eval()
+    assert set(model.state_dict()) == set(reference.state_dict())
+    model.load_state_dict(reference.state_dict())
+    assert isinstance(model.time_embed.time_proj, SanaTimesteps)
+    assert isinstance(model.time_embed.timestep_embedder, SanaTimestepEmbedding)
+    assert isinstance(model.time_embed.guidance_condition_proj, SanaTimesteps)
+    assert isinstance(model.time_embed.guidance_embedder, SanaTimestepEmbedding)
+
+    torch.manual_seed(31)
+    hidden_states = torch.randn(1, 4, 3, 4, 4)
+    encoder_hidden_states = torch.randn(1, 5, 8)
+    encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]])
+    timestep = torch.tensor([500.0])
+    guidance = torch.tensor([6.0])
+
+    with torch.no_grad():
+        expected = reference(
+            hidden_states,
+            encoder_hidden_states,
+            timestep,
+            guidance=guidance,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=False,
+        )
+        actual = model(
+            hidden_states,
+            encoder_hidden_states,
+            timestep,
+            guidance=guidance,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=False,
+        )
+
+    assert isinstance(actual, tuple)
+    assert len(actual) == len(expected) == 1
+    torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0)
+
+
+def test_native_transformer_config_filters_diffusers_metadata():
+    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerConfig
+
+    config = SanaVideoTransformerConfig.from_dict(
+        _TINY_CONFIG | {"_class_name": "SanaVideoTransformer3DModel", "_diffusers_version": "0.38.0"}
+    )
+    assert config.patch_size == (1, 2, 2)
+
+    with pytest.raises(ValueError, match="unsupported_field"):
+        SanaVideoTransformerConfig.from_dict(_TINY_CONFIG | {"unsupported_field": True})
+
+
+def test_dual_vae_variant_resolution():
+    from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+        DistributedAutoencoderKLLTX2Video,
+    )
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+    from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+    from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
+        SanaVideoPipeline,
+        _resolve_vae_class_and_dtype,
+    )
+
+    assert not issubclass(SanaVideoPipeline, DiffusionPipeline)
+    assert issubclass(SanaVideoPipeline, ProgressBarMixin)
+
+    vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLWan", torch.bfloat16)
+    assert vae_class is DistributedAutoencoderKLWan
+    assert dtype is torch.float32
+
+    vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLLTX2Video", torch.bfloat16)
+    assert vae_class is DistributedAutoencoderKLLTX2Video
+    assert dtype is torch.bfloat16
+
+    with pytest.raises(ValueError, match="Unsupported SANA-Video VAE"):
+        _resolve_vae_class_and_dtype("UnknownVAE", torch.bfloat16)