Merge branch 'main' into use-fixture-modular-tests

* Fix LTX-2 image-to-video generation failure in two stages generation

In LTX-2's two-stage image-to-video generation task, specifically after
the upsampling step, a shape mismatch occurs between the `latents` and
the `conditioning_mask`, which causes an error in function
`_create_noised_state`.

Fix it by creating the `conditioning_mask` based on the shape of the
`latents`.

* Add unit test for LTX-2 i2v two stages inference with upsampler

* Downscaling the upsampler in LTX-2 image-to-video unit test

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

docs/source/en/modular_diffusers/custom_blocks.md

@@ -332,49 +332,4 @@ Make your custom block work with Mellon's visual interface. See the [Mellon Cust
 Browse the [Modular Diffusers Custom Blocks](https://huggingface.co/collections/diffusers/modular-diffusers-custom-blocks) collection for inspiration and ready-to-use blocks.
 
 </hfoption>
-</hfoptions>
-
-## Dependencies
-
-Declaring package dependencies in custom blocks prevents runtime import errors later on. Diffusers validates the dependencies and returns a warning if a package is missing or incompatible.
-
-Set a `_requirements` attribute in your block class, mapping package names to version specifiers.
-
-```py
-from diffusers.modular_pipelines import PipelineBlock
-
-class MyCustomBlock(PipelineBlock):
-    _requirements = {
-        "transformers": ">=4.44.0",
-        "sentencepiece": ">=0.2.0"
-    }
-```
-
-When there are blocks with different requirements, Diffusers merges their requirements.
-
-```py
-from diffusers.modular_pipelines import SequentialPipelineBlocks
-
-class BlockA(PipelineBlock):
-    _requirements = {"transformers": ">=4.44.0"}
-    # ...
-
-class BlockB(PipelineBlock):
-    _requirements = {"sentencepiece": ">=0.2.0"}
-    # ...
-
-pipe = SequentialPipelineBlocks.from_blocks_dict({
-    "block_a": BlockA,
-    "block_b": BlockB,
-})
-```
-
-When this block is saved with [`~ModularPipeline.save_pretrained`], the requirements are saved to the `modular_config.json` file. When this block is loaded, Diffusers checks each requirement against the current environment. If there is a mismatch or a package isn't found, Diffusers returns the following warning.
-
-```md
-# missing package
-xyz-package was specified in the requirements but wasn't found in the current environment.
-
-# version mismatch
-xyz requirement 'specific-version' is not satisfied by the installed version 'actual-version'. Things might work unexpected.
-```
+</hfoptions>

src/diffusers/commands/custom_blocks.py

@@ -89,6 +89,8 @@ class CustomBlocksCommand(BaseDiffusersCLICommand):
         # automap = self._create_automap(parent_class=parent_class, child_class=child_class)
         # with open(CONFIG, "w") as f:
         #     json.dump(automap, f)
+        with open("requirements.txt", "w") as f:
+            f.write("")
 
     def _choose_block(self, candidates, chosen=None):
         for cls, base in candidates:

src/diffusers/models/attention_dispatch.py

@@ -733,7 +733,7 @@ def _wrapped_flash_attn_3(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     # Hardcoded for now because pytorch does not support tuple/int type hints
     window_size = (-1, -1)
-    out, lse, *_ = flash_attn_3_func(
+    result = flash_attn_3_func(
         q=q,
         k=k,
         v=v,
@@ -750,7 +750,9 @@ def _wrapped_flash_attn_3(
         pack_gqa=pack_gqa,
         deterministic=deterministic,
         sm_margin=sm_margin,
+        return_attn_probs=True,
     )
+    out, lse, *_ = result
     lse = lse.permute(0, 2, 1)
     return out, lse
 
@@ -2701,7 +2703,7 @@ def _flash_varlen_attention_3(
     key_packed = torch.cat(key_valid, dim=0)
     value_packed = torch.cat(value_valid, dim=0)
 
-    out, lse, *_ = flash_attn_3_varlen_func(
+    result = flash_attn_3_varlen_func(
         q=query_packed,
         k=key_packed,
         v=value_packed,
@@ -2711,7 +2713,13 @@ def _flash_varlen_attention_3(
         max_seqlen_k=max_seqlen_k,
         softmax_scale=scale,
         causal=is_causal,
+        return_attn_probs=return_lse,
     )
+    if isinstance(result, tuple):
+        out, lse, *_ = result
+    else:
+        out = result
+        lse = None
     out = out.unflatten(0, (batch_size, -1))
 
     return (out, lse) if return_lse else out

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -40,7 +40,6 @@ from .modular_pipeline_utils import (
     InputParam,
     InsertableDict,
     OutputParam,
-    _validate_requirements,
     combine_inputs,
     combine_outputs,
     format_components,
@@ -291,7 +290,6 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin):
 
     config_name = "modular_config.json"
     model_name = None
-    _requirements: dict[str, str] | None = None
     _workflow_map = None
 
     @classmethod
@@ -406,9 +404,6 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin):
                 "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file."
             )
 
-        if "requirements" in config and config["requirements"] is not None:
-            _ = _validate_requirements(config["requirements"])
-
         class_ref = config["auto_map"][cls.__name__]
         module_file, class_name = class_ref.split(".")
         module_file = module_file + ".py"
@@ -433,13 +428,8 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin):
         module = full_mod.rsplit(".", 1)[-1].replace("__dynamic__", "")
         parent_module = self.save_pretrained.__func__.__qualname__.split(".", 1)[0]
         auto_map = {f"{parent_module}": f"{module}.{cls_name}"}
+
         self.register_to_config(auto_map=auto_map)
-
-        # resolve requirements
-        requirements = _validate_requirements(getattr(self, "_requirements", None))
-        if requirements:
-            self.register_to_config(requirements=requirements)
-
         self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
         config = dict(self.config)
         self._internal_dict = FrozenDict(config)
@@ -1250,14 +1240,6 @@ class SequentialPipelineBlocks(ModularPipelineBlocks):
             expected_configs=self.expected_configs,
         )
 
-    @property
-    def _requirements(self) -> dict[str, str]:
-        requirements = {}
-        for block_name, block in self.sub_blocks.items():
-            if getattr(block, "_requirements", None):
-                requirements[block_name] = block._requirements
-        return requirements
-
 
 class LoopSequentialPipelineBlocks(ModularPipelineBlocks):
     """

src/diffusers/modular_pipelines/modular_pipeline_utils.py

@@ -22,12 +22,10 @@ from typing import Any, Literal, Type, Union, get_args, get_origin
 
 import PIL.Image
 import torch
-from packaging.specifiers import InvalidSpecifier, SpecifierSet
 
 from ..configuration_utils import ConfigMixin, FrozenDict
 from ..loaders.single_file_utils import _is_single_file_path_or_url
 from ..utils import DIFFUSERS_LOAD_ID_FIELDS, is_torch_available, logging
-from ..utils.import_utils import _is_package_available
 
 
 if is_torch_available():
@@ -974,89 +972,6 @@ def make_doc_string(
     return output
 
 
-def _validate_requirements(reqs):
-    if reqs is None:
-        normalized_reqs = {}
-    else:
-        if not isinstance(reqs, dict):
-            raise ValueError(
-                "Requirements must be provided as a dictionary mapping package names to version specifiers."
-            )
-        normalized_reqs = _normalize_requirements(reqs)
-
-    if not normalized_reqs:
-        return {}
-
-    final: dict[str, str] = {}
-    for req, specified_ver in normalized_reqs.items():
-        req_available, req_actual_ver = _is_package_available(req)
-        if not req_available:
-            logger.warning(f"{req} was specified in the requirements but wasn't found in the current environment.")
-
-        if specified_ver:
-            try:
-                specifier = SpecifierSet(specified_ver)
-            except InvalidSpecifier as err:
-                raise ValueError(f"Requirement specifier '{specified_ver}' for {req} is invalid.") from err
-
-            if req_actual_ver == "N/A":
-                logger.warning(
-                    f"Version of {req} could not be determined to validate requirement '{specified_ver}'. Things might work unexpected."
-                )
-            elif not specifier.contains(req_actual_ver, prereleases=True):
-                logger.warning(
-                    f"{req} requirement '{specified_ver}' is not satisfied by the installed version {req_actual_ver}. Things might work unexpected."
-                )
-
-        final[req] = specified_ver
-
-    return final
-
-
-def _normalize_requirements(reqs):
-    if not reqs:
-        return {}
-
-    normalized: "OrderedDict[str, str]" = OrderedDict()
-
-    def _accumulate(mapping: dict[str, Any]):
-        for pkg, spec in mapping.items():
-            if isinstance(spec, dict):
-                # This is recursive because blocks are composable. This way, we can merge requirements
-                # from multiple blocks.
-                _accumulate(spec)
-                continue
-
-            pkg_name = str(pkg).strip()
-            if not pkg_name:
-                raise ValueError("Requirement package name cannot be empty.")
-
-            spec_str = "" if spec is None else str(spec).strip()
-            if spec_str and not spec_str.startswith(("<", ">", "=", "!", "~")):
-                spec_str = f"=={spec_str}"
-
-            existing_spec = normalized.get(pkg_name)
-            if existing_spec is not None:
-                if not existing_spec and spec_str:
-                    normalized[pkg_name] = spec_str
-                elif existing_spec and spec_str and existing_spec != spec_str:
-                    try:
-                        combined_spec = SpecifierSet(",".join(filter(None, [existing_spec, spec_str])))
-                    except InvalidSpecifier:
-                        logger.warning(
-                            f"Conflicting requirements for '{pkg_name}' detected: '{existing_spec}' vs '{spec_str}'. Keeping '{existing_spec}'."
-                        )
-                    else:
-                        normalized[pkg_name] = str(combined_spec)
-                continue
-
-            normalized[pkg_name] = spec_str
-
-    _accumulate(reqs)
-
-    return normalized
-
-
 def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> list[InputParam]:
     """
     Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current

src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py

@@ -699,9 +699,13 @@ class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraL
         mask_shape = (batch_size, 1, num_frames, height, width)
 
         if latents is not None:
-            conditioning_mask = latents.new_zeros(mask_shape)
-            conditioning_mask[:, :, 0] = 1.0
             if latents.ndim == 5:
+                # conditioning_mask needs to the same shape as latents in two stages generation.
+                batch_size, _, num_frames, height, width = latents.shape
+                mask_shape = (batch_size, 1, num_frames, height, width)
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
                 latents = self._normalize_latents(
                     latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
                 )
@@ -710,6 +714,9 @@ class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraL
                 latents = self._pack_latents(
                     latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
                 )
+            else:
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
             conditioning_mask = self._pack_latents(
                 conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
             ).squeeze(-1)

tests/modular_pipelines/flux/test_modular_pipeline_flux.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import random
-import tempfile
 
 import numpy as np
 import PIL
@@ -129,18 +128,16 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
 
         return inputs
 
-    def test_save_from_pretrained(self):
+    def test_save_from_pretrained(self, tmp_path):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            base_pipe.save_pretrained(tmpdirname)
-
-            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
-            pipe.load_components(torch_dtype=torch.float32)
-            pipe.to(torch_device)
-            pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
+        base_pipe.save_pretrained(tmp_path)
+        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
+        pipe.load_components(torch_dtype=torch.float32)
+        pipe.to(torch_device)
+        pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
 
         pipes.append(pipe)
 
@@ -212,18 +209,16 @@ class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
 
         return inputs
 
-    def test_save_from_pretrained(self):
+    def test_save_from_pretrained(self, tmp_path):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            base_pipe.save_pretrained(tmpdirname)
-
-            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
-            pipe.load_components(torch_dtype=torch.float32)
-            pipe.to(torch_device)
-            pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
+        base_pipe.save_pretrained(tmp_path)
+        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
+        pipe.load_components(torch_dtype=torch.float32)
+        pipe.to(torch_device)
+        pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
 
         pipes.append(pipe)
 

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -1,7 +1,4 @@
 import gc
-import json
-import os
-import tempfile
 from typing import Callable
 
 import pytest
@@ -10,7 +7,6 @@ import torch
 import diffusers
 from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks
 from diffusers.guiders import ClassifierFreeGuidance
-from diffusers.modular_pipelines import SequentialPipelineBlocks
 from diffusers.modular_pipelines.modular_pipeline_utils import (
     ComponentSpec,
     ConfigSpec,
@@ -20,13 +16,7 @@ from diffusers.modular_pipelines.modular_pipeline_utils import (
 )
 from diffusers.utils import logging
 
-from ..testing_utils import (
-    CaptureLogger,
-    backend_empty_cache,
-    numpy_cosine_similarity_distance,
-    require_accelerator,
-    torch_device,
-)
+from ..testing_utils import backend_empty_cache, numpy_cosine_similarity_distance, require_accelerator, torch_device
 
 
 class ModularPipelineTesterMixin:
@@ -337,16 +327,15 @@ class ModularPipelineTesterMixin:
 
         assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
 
-    def test_save_from_pretrained(self):
+    def test_save_from_pretrained(self, tmp_path):
         pipes = []
         base_pipe = self.get_pipeline().to(torch_device)
         pipes.append(base_pipe)
 
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            base_pipe.save_pretrained(tmpdirname)
-            pipe = ModularPipeline.from_pretrained(tmpdirname).to(torch_device)
-            pipe.load_components(torch_dtype=torch.float32)
-            pipe.to(torch_device)
+        base_pipe.save_pretrained(tmp_path)
+        pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
+        pipe.load_components(torch_dtype=torch.float32)
+        pipe.to(torch_device)
 
         pipes.append(pipe)
 
@@ -409,56 +398,6 @@ class ModularGuiderTesterMixin:
         assert max_diff > expected_max_diff, "Output with CFG must be different from normal inference"
 
 
-class TestCustomBlockRequirements:
-    def get_dummy_block_pipe(self):
-        class DummyBlockOne:
-            # keep two arbitrary deps so that we can test warnings.
-            _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"}
-
-        class DummyBlockTwo:
-            # keep two dependencies that will be available during testing.
-            _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"}
-
-        pipe = SequentialPipelineBlocks.from_blocks_dict(
-            {"dummy_block_one": DummyBlockOne, "dummy_block_two": DummyBlockTwo}
-        )
-        return pipe
-
-    def test_custom_requirements_save_load(self):
-        pipe = self.get_dummy_block_pipe()
-        with tempfile.TemporaryDirectory() as tmpdir:
-            pipe.save_pretrained(tmpdir)
-            config_path = os.path.join(tmpdir, "modular_config.json")
-            with open(config_path, "r") as f:
-                config = json.load(f)
-
-        assert "requirements" in config
-        requirements = config["requirements"]
-
-        expected_requirements = {
-            "xyz": ">=0.8.0",
-            "abc": ">=10.0.0",
-            "transformers": ">=4.44.0",
-            "diffusers": ">=0.2.0",
-        }
-        assert expected_requirements == requirements
-
-    def test_warnings(self):
-        pipe = self.get_dummy_block_pipe()
-        with tempfile.TemporaryDirectory() as tmpdir:
-            logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline_utils")
-            logger.setLevel(30)
-
-            with CaptureLogger(logger) as cap_logger:
-                pipe.save_pretrained(tmpdir)
-
-            template = "{req} was specified in the requirements but wasn't found in the current environment"
-            msg_xyz = template.format(req="xyz")
-            msg_abc = template.format(req="abc")
-            assert msg_xyz in str(cap_logger.out)
-            assert msg_abc in str(cap_logger.out)
-
-
 class TestModularModelCardContent:
     def create_mock_block(self, name="TestBlock", description="Test block description"):
         class MockBlock:

tests/modular_pipelines/test_modular_pipelines_custom_blocks.py

@@ -14,7 +14,6 @@
 
 import json
 import os
-import tempfile
 from collections import deque
 from typing import List
 
@@ -153,25 +152,24 @@ class TestModularCustomBlocks:
         output_prompt = output.values["output_prompt"]
         assert output_prompt.startswith("Modular diffusers + ")
 
-    def test_custom_block_saving_loading(self):
+    def test_custom_block_saving_loading(self, tmp_path):
         custom_block = DummyCustomBlockSimple()
 
-        with tempfile.TemporaryDirectory() as tmpdir:
-            custom_block.save_pretrained(tmpdir)
-            assert any("modular_config.json" in k for k in os.listdir(tmpdir))
+        custom_block.save_pretrained(tmp_path)
+        assert any("modular_config.json" in k for k in os.listdir(tmp_path))
 
-            with open(os.path.join(tmpdir, "modular_config.json"), "r") as f:
-                config = json.load(f)
-            auto_map = config["auto_map"]
-            assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
+        with open(os.path.join(tmp_path, "modular_config.json"), "r") as f:
+            config = json.load(f)
+        auto_map = config["auto_map"]
+        assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"}
 
-            # For now, the Python script that implements the custom block has to be manually pushed to the Hub.
-            # This is why, we have to separately save the Python script here.
-            code_path = os.path.join(tmpdir, "test_modular_pipelines_custom_blocks.py")
-            with open(code_path, "w") as f:
-                f.write(CODE_STR)
+        # For now, the Python script that implements the custom block has to be manually pushed to the Hub.
+        # This is why, we have to separately save the Python script here.
+        code_path = os.path.join(tmp_path, "test_modular_pipelines_custom_blocks.py")
+        with open(code_path, "w") as f:
+            f.write(CODE_STR)
 
-            loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmpdir, trust_remote_code=True)
+        loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmp_path, trust_remote_code=True)
 
         pipe = loaded_custom_block.init_pipeline()
         prompt = "Diffusers is nice"

tests/pipelines/ltx2/test_ltx2_image2video.py

@@ -24,7 +24,8 @@ from diffusers import (
     LTX2ImageToVideoPipeline,
     LTX2VideoTransformer3DModel,
 )
-from diffusers.pipelines.ltx2 import LTX2TextConnectors
+from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline, LTX2TextConnectors
+from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
 
 from ...testing_utils import enable_full_determinism
@@ -174,6 +175,15 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
 
         return components
 
+    def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1):
+        upsampler = LTX2LatentUpsamplerModel(
+            in_channels=in_channels,
+            mid_channels=mid_channels,
+            num_blocks_per_stage=num_blocks_per_stage,
+        )
+
+        return upsampler
+
     def get_dummy_inputs(self, device, seed=0):
         if str(device).startswith("mps"):
             generator = torch.manual_seed(seed)
@@ -287,5 +297,60 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
         assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4)
         assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4)
 
+    def test_two_stages_inference_with_upsampler(self):
+        device = "cpu"
+
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe.to(device)
+        pipe.set_progress_bar_config(disable=None)
+
+        inputs = self.get_dummy_inputs(device)
+        inputs["output_type"] = "latent"
+        first_stage_output = pipe(**inputs)
+        video_latent = first_stage_output.frames
+        audio_latent = first_stage_output.audio
+
+        self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16))
+        self.assertEqual(audio_latent.shape, (1, 2, 5, 2))
+        self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels)
+
+        upsampler = self.get_dummy_upsample_component(in_channels=video_latent.shape[1])
+        upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler)
+        upscaled_video_latent = upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0]
+        self.assertEqual(upscaled_video_latent.shape, (1, 4, 3, 32, 32))
+
+        inputs["latents"] = upscaled_video_latent
+        inputs["audio_latents"] = audio_latent
+        inputs["output_type"] = "pt"
+        second_stage_output = pipe(**inputs)
+        video = second_stage_output.frames
+        audio = second_stage_output.audio
+
+        self.assertEqual(video.shape, (1, 5, 3, 64, 64))
+        self.assertEqual(audio.shape[0], 1)
+        self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels)
+
+        # fmt: off
+        expected_video_slice = torch.tensor(
+            [
+                0.4497, 0.6757, 0.4219, 0.7686, 0.4525, 0.6483, 0.3969, 0.7404, 0.3541, 0.3039, 0.4592, 0.3521, 0.3665, 0.2785, 0.3336, 0.3079
+            ]
+        )
+        expected_audio_slice = torch.tensor(
+            [
+                0.0271, 0.0492, 0.1249, 0.1126, 0.1661, 0.1060, 0.1717, 0.0944, 0.0672, -0.0069, 0.0688, 0.0097, 0.0808, 0.1231, 0.0986, 0.0739
+            ]
+        )
+        # fmt: on
+
+        video = video.flatten()
+        audio = audio.flatten()
+        generated_video_slice = torch.cat([video[:8], video[-8:]])
+        generated_audio_slice = torch.cat([audio[:8], audio[-8:]])
+
+        assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4)
+        assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4)
+
     def test_inference_batch_single_identical(self):
         self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=2e-2)