[Model] Add TP and CFG parallelism to SANA-Video 2B (#5861)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

tests/model_tests/diffusion/case_filtering.py

@@ -76,6 +76,8 @@ def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts
             test_info.supported_tasks,
             test_info.check_multi_output and test_group is None,
             test_info.check_determinism and test_group is None,
+            # The divergence check compares generated outputs, which only the offline test does.
+            *([] if online else [test_info.check_i2v_t2v_divergence]),
             id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
             marks=get_test_group_marks(model_name, test_group, test_info.marks),
         )

tests/model_tests/diffusion/config_types.py

@@ -76,6 +76,11 @@ class DiffusionModelTestOpts(NamedTuple):
     check_multi_output: bool = True  # Runs multiple generations in one request
     check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
 
+    # For IMAGE_TO_VIDEO, also run a T2V request and assert the outputs differ, to catch a
+    # silently dropped input image. Only valid for unified text/image pipelines (e.g. LTX2);
+    # set False for pipelines that mandate an image and fail closed without one (SANA-Video I2V).
+    check_i2v_t2v_divergence: bool = True
+
 
 ### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
 ACC_OMNI_KWARGS = {

tests/model_tests/diffusion/diff_model_builders.py

@@ -1,3 +1,4 @@
+import json
 from functools import partial
 from pathlib import Path
 
@@ -49,15 +50,17 @@ def tiny_ltx2_builder() -> str:
     return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
 
 
-def tiny_sana_video_builder() -> str:
+def _build_tiny_sana_video(pipeline_name: str) -> str:
     """Build a tiny 480p SANA-Video model without downloading model weights.
 
     The tokenizer is the only component loaded from the upstream repository.
     All modules with weights are initialized locally from intentionally small
-    configs, and the scheduler is constructed from its weight-free config.
+    configs, and the scheduler is constructed from its weight-free config. The
+    I2V pipeline shares this diffusers layout, so it reuses the same components
+    and only rewrites model_index.json's ``_class_name`` to its own class.
     """
     model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers"
-    model_dir = _get_tiny_model_path("SanaVideoPipeline")
+    model_dir = _get_tiny_model_path(pipeline_name)
 
     tokenizer = GemmaTokenizerFast.from_pretrained(model_id, subfolder="tokenizer")
     scheduler = DPMSolverMultistepScheduler(
@@ -112,9 +115,22 @@ def tiny_sana_video_builder() -> str:
         scheduler=scheduler,
     )
     pipeline.to(dtype=torch.bfloat16).save_pretrained(model_dir)
+    if pipeline_name != "SanaVideoPipeline":
+        index_path = Path(model_dir) / "model_index.json"
+        index = json.loads(index_path.read_text())
+        index["_class_name"] = pipeline_name
+        index_path.write_text(json.dumps(index, indent=2))
     return model_dir
 
 
+def tiny_sana_video_builder() -> str:
+    return _build_tiny_sana_video("SanaVideoPipeline")
+
+
+def tiny_sana_video_i2v_builder() -> str:
+    return _build_tiny_sana_video("SanaImageToVideoPipeline")
+
+
 def _shrink_flux_clip_text_encoder(config: dict) -> dict:
     config["num_hidden_layers"] = 2
     return config

tests/model_tests/diffusion/model_settings.py

@@ -53,6 +53,22 @@ DIFFUSION_TEST_SETTINGS = {
         model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
         builder=diff_model_builders.tiny_sana_video_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
+        extra_test_groups=[
+            [DiffusionAccs.TENSOR_PARALLEL],
+            [DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+        ],
+        check_multi_output=False,
+    ),
+    "SanaImageToVideoPipeline": DiffusionModelTestOpts(
+        model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+        builder=diff_model_builders.tiny_sana_video_i2v_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_VIDEO],
+        extra_test_groups=[
+            [DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+        ],
+        check_i2v_t2v_divergence=False,
         check_multi_output=False,
     ),
     "QwenImagePipeline": DiffusionModelTestOpts(

tests/model_tests/diffusion/task_runners.py

@@ -220,16 +220,20 @@ def run_and_validate_text_to_video_request(omni: Omni):
     _validate_video(_run_offline_t2v(omni))
 
 
-def run_and_validate_image_to_video_request(omni: Omni):
+def run_and_validate_image_to_video_request(omni: Omni, check_t2v_divergence: bool = True):
     """Run and validate an image to video request.
 
     LTX2Pipeline uses a unified text/image entry and can fall back to T2V if
     the image never reaches conditioning, so also run T2V with the same seed
     and prompt and assert the outputs differ; a regression that silently
     drops the input image would otherwise still pass a "did we get a valid
-    video back" check alone.
+    video back" check alone. Pipelines that mandate an image and fail closed
+    without one (SANA-Video I2V) pass check_t2v_divergence=False: a dropped
+    image raises instead of silently degrading, so the T2V comparison is moot.
     """
     i2v_video = _validate_video_frames(_get_offline_videos(_run_offline_i2v(omni)))[0]
+    if not check_t2v_divergence:
+        return
     t2v_video = _validate_video_frames(_get_offline_videos(_run_offline_t2v(omni)))[0]
     assert not np.array_equal(i2v_video, t2v_video), (
         "I2V output is identical to T2V output with the same seed and prompt; "

tests/model_tests/diffusion/test_alignment.py

@@ -61,12 +61,6 @@ EXCLUDED_MODELS = [
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
     "LingBotVideoPipeline",
-    # The common harness does not yet define an image-to-video task or supply
-    # an input image. SANA I2V is covered by
-    # tests/diffusion/models/sana_video/test_pipeline_sana_video.py and
-    # tests/e2e/online_serving/test_sana_video_expansion.py. SANA T2V is exercised through
-    # the common tiny-model setting in model_settings.py.
-    "SanaImageToVideoPipeline",
     # Requires camera/action assets and persistent AR state; covered by its
     # dedicated contract tests and real-checkpoint E2E smoke instead.
     "LingBotWorldCausalDMDPipeline",

tests/model_tests/diffusion/test_common_offline.py

@@ -24,7 +24,7 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
 
 
 @pytest.mark.parametrize(
-    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism,check_i2v_t2v_divergence",
     get_parametrized_options(DIFFUSION_TEST_SETTINGS),
 )
 def test_pipeline_on_supported_tasks(
@@ -33,6 +33,7 @@ def test_pipeline_on_supported_tasks(
     supported_tasks: list[DiffusionTasks],
     check_multioutput: bool,
     check_determinism: bool,
+    check_i2v_t2v_divergence: bool,
     tiny_model_paths: dict[str, str],
     subtests,
 ):
@@ -61,7 +62,7 @@ def test_pipeline_on_supported_tasks(
                 elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
                     run_and_validate_text_to_video_request(omni)
                 elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
-                    run_and_validate_image_to_video_request(omni)
+                    run_and_validate_image_to_video_request(omni, check_t2v_divergence=check_i2v_t2v_divergence)
                 else:
                     raise ValueError(f"Task type {task_type} is not yet supported")