Fix MPS float64 crash in MiniMax-H3 PrepareLayoutStep

src/diffusers/modular_pipelines/minimax_h3/before_denoise.py

@@ -17,7 +17,7 @@
 
 from ...schedulers import MiniMaxH3Scheduler
 from ...utils import logging
-from ...utils.torch_utils import randn_tensor
+from ...utils.torch_utils import maybe_adjust_dtype_for_device, randn_tensor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
 from .modular_pipeline import (
@@ -441,7 +441,10 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -
             components.video_tag,
             block_state.keyframe_anchors,
         )
-        block_state.position_ids = position_ids.to(device)
+        # `position_ids` is built in float64 for cumsum precision over long sequences (see `_temporal_position_grid`);
+        # MPS/NPU have no float64 support, so only those devices get the float32 downcast the rope embedding applies
+        # anyway (`MiniMaxH3RotaryPosEmbed.forward`).
+        block_state.position_ids = position_ids.to(device, dtype=maybe_adjust_dtype_for_device(torch.float64, device))
         block_state.token_tags = token_tags.to(device)
         block_state.video_indices = video_indices.to(device)
         block_state.audio_indices = audio_indices.to(device)
@@ -765,7 +768,10 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -
             components.audio_tag,
             components.video_tag,
         )
-        block_state.position_ids = position_ids.to(device)
+        # `position_ids` is built in float64 for cumsum precision over long sequences (see `_temporal_position_grid`);
+        # MPS/NPU have no float64 support, so only those devices get the float32 downcast the rope embedding applies
+        # anyway (`MiniMaxH3RotaryPosEmbed.forward`).
+        block_state.position_ids = position_ids.to(device, dtype=maybe_adjust_dtype_for_device(torch.float64, device))
         block_state.token_tags = token_tags.to(device)
         block_state.video_indices = video_indices.to(device)
         block_state.audio_indices = audio_indices.to(device)

tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py

@@ -27,6 +27,7 @@
     MiniMaxH3ImageReference,
     MiniMaxH3VideoReference,
 )
+from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3PrepareLayoutStep
 from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
 from diffusers.modular_pipelines.minimax_h3.encoders import (
     MiniMaxH3FL2VATextEncoderStep,
@@ -35,7 +36,9 @@
     MiniMaxH3TextEncoderStep,
 )
 from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MINIMAX_H3_FPS
+from diffusers.utils.torch_utils import maybe_adjust_dtype_for_device
 
+from ...testing_utils import torch_device
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -894,3 +897,64 @@ def test_video_reference_resampled_to_the_model_frame_rate(self):
             frames, fps=float(MINIMAX_H3_FPS), num_frames=124, **canvas
         )
         assert np.shares_memory(untouched, frames)
+
+
+class TestMiniMaxH3PositionIdsDevice:
+    """
+    `MiniMaxH3PrepareLayoutStep.build_packed_sequence` lays `position_ids` out in float64: the temporal axis is a
+    `cumsum` over a `5/3`-scaled span (see `_temporal_position_grid`), and float32 drifts perceptibly over the
+    thousands of latent frames a long video packs. MPS (and NPU) have no float64 support, so the layout steps have to
+    downcast when they move `position_ids` onto the execution device — see
+    https://github.com/huggingface/diffusers/issues/14639. Neither this nor the checkpoint-free
+    [`TestMiniMaxH3Reference`] above needs a checkpoint or an accelerator to run.
+    """
+
+    def test_position_ids_are_float64_on_the_host(self):
+        r"""The packed layout is always built in float64, regardless of what device it later moves to."""
+        text_token_tags = torch.ones(4, dtype=torch.int64)
+        position_ids, *_ = MiniMaxH3PrepareLayoutStep.build_packed_sequence(
+            text_token_tags,
+            num_latent_frames=2,
+            latent_height=16,
+            latent_width=16,
+            num_audio_latents=2,
+            patch_size=(1, 8, 8),
+            audio_channels=1,
+            audio_tag=2,
+            video_tag=0,
+        )
+        assert position_ids.dtype == torch.float64
+
+    @pytest.mark.parametrize(
+        "device_type, expected_dtype",
+        [("cpu", torch.float64), ("cuda", torch.float64), ("mps", torch.float32)],
+    )
+    def test_device_transfer_dtype_matches_the_execution_device(self, device_type, expected_dtype):
+        r"""
+        The dtype the layout steps pass to `position_ids.to(device, dtype=...)` is float32 exactly on the devices
+        that cannot hold float64, and float64 (i.e. a no-op cast) everywhere else — including on a machine that has
+        none of these backends, since building a `torch.device` object needs no hardware. NPU is also downcast by
+        `maybe_adjust_dtype_for_device`, but constructing a `torch.device("npu")` needs the `torch_npu` plugin
+        installed, so it is exercised only by that helper's own device-agnostic logic, not by an actual device here.
+        """
+        assert maybe_adjust_dtype_for_device(torch.float64, torch.device(device_type)) == expected_dtype
+
+    @pytest.mark.skipif(torch_device != "mps", reason="exercises the actual MPS float64 tensor-move limitation")
+    def test_position_ids_move_to_mps_as_float32(self):
+        r"""On real MPS hardware, the layout step's device transfer no longer tries to materialize a float64 tensor."""
+        text_token_tags = torch.ones(4, dtype=torch.int64)
+        position_ids, *_ = MiniMaxH3PrepareLayoutStep.build_packed_sequence(
+            text_token_tags,
+            num_latent_frames=2,
+            latent_height=16,
+            latent_width=16,
+            num_audio_latents=2,
+            patch_size=(1, 8, 8),
+            audio_channels=1,
+            audio_tag=2,
+            video_tag=0,
+        )
+        device = torch.device(torch_device)
+        moved = position_ids.to(device, dtype=maybe_adjust_dtype_for_device(torch.float64, device))
+        assert moved.dtype == torch.float32
+        assert moved.device.type == "mps"