@@ -10,8 +10,11 @@ RUN apt-get -y update \
     && add-apt-repository ppa:deadsnakes/ppa && \
     apt-get update
 
+# cuda-cudart-dev provides the CUDA headers (cuda.h etc.) AOT-inductor compilation needs;
+# the runtime base image only ships the libraries.
 RUN apt install -y bash \
     build-essential \
+    cuda-cudart-dev-12-9 \
     git \
     git-lfs \
     curl \
 
@@ -24,7 +24,7 @@
 from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
 from ..attention_dispatch import dispatch_attention_fn
 from ..modeling_outputs import AutoencoderKLOutput
-from ..modeling_utils import ModelMixin
+from ..modeling_utils import ModelMixin, get_parameter_dtype
 from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution
 
 
@@ -857,6 +857,9 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutp
             The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the sampled latents with
             `latents_mean` / `latents_std` afterwards.
         """
+        # Every module is pinned to float32 by `_keep_in_fp32_modules`, so a pipeline running in a lower `torch_dtype`
+        # hands over lower-precision pixels; align them with the weights, like the audio autoencoder does.
+        x = x.to(get_parameter_dtype(self.encoder))
         if self.use_slicing and x.shape[0] > 1:
             moments = torch.cat([self._encode(x_slice) for x_slice in x.split(1)])
         else:
@@ -881,6 +884,7 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t
             [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
                 The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`.
         """
+        z = z.to(get_parameter_dtype(self.decoder))
         if self.use_slicing and z.shape[0] > 1:
             decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)])
         else:
 
@@ -524,6 +524,10 @@ class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin):
     """
 
     _supports_gradient_checkpointing = False
+    # `weight_norm` recomputes `weight` from `weight_g` / `weight_v` in a forward pre-hook, which runs before the
+    # leaf-level group offloading hook has onloaded them, so the convolution would see a CPU weight. Same reason the
+    # other weight-normalized audio autoencoders (`AutoencoderOobleck`, `Cosmos3AVAEAudioTokenizer`) opt out.
+    _supports_group_offloading = False
     # The released checkpoint is float32 and the DAC/BigVGAN stack (weight-normalized convolutions, Snake
     # activations) degrades audibly under bfloat16 (roughly 20 dB quieter decodes), so a pipeline-level
     # `torch_dtype=torch.bfloat16` must not downcast the weights.
 
@@ -46,12 +46,14 @@ class MiniMaxH3TransformerOutput(BaseOutput):
         sample (`torch.Tensor` of shape `(batch_size, num_video_tokens, in_channels * prod(patch_size))`):
             The video velocity prediction for the rows addressed by `video_indices`, in the same order. Conditioning
             rows are returned unmasked — masking them out before the scheduler step is the caller's job.
-        audio_sample (`torch.Tensor` of shape `(batch_size, num_audio_tokens, audio_in_channels)`):
-            The audio velocity prediction for the rows addressed by `audio_indices`, in the same order.
+        audio_sample (`torch.Tensor` of shape `(batch_size, num_audio_tokens, audio_in_channels)`, defaults to `None`):
+            The audio velocity prediction for the rows addressed by `audio_indices`, in the same order. `forward`
+            always populates it; it only defaults to `None` so that the output can be rebuilt from a plain dict of its
+            fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device.
     """
 
     sample: torch.Tensor
-    audio_sample: torch.Tensor
+    audio_sample: torch.Tensor | None = None
 
 
 def _apply_rotary_emb(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
 
@@ -20,23 +20,19 @@
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin
+from ..testing_utils import (
+    BaseModelTesterConfig,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TrainingTesterMixin,
+    run_nondeterministic,
+)
 from .testing_utils import NewAutoencoderTesterMixin
 
 
 enable_full_determinism()
 
 
-def _run_nondeterministic(fn):
-    # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation;
-    # temporarily relax the requirement for tests that do backward passes.
-    torch.use_deterministic_algorithms(False)
-    try:
-        fn()
-    finally:
-        torch.use_deterministic_algorithms(True)
-
-
 class AutoencoderKLKVAEVideoTesterConfig(BaseModelTesterConfig):
     @property
     def model_class(self):
@@ -91,14 +87,16 @@ def test_gradient_checkpointing_is_applied(self):
         expected_set = {"KVAECachedEncoder3D", "KVAECachedDecoder3D"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
+    # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation, so every test with a backward
+    # pass runs with determinism relaxed.
     def test_training(self):
-        _run_nondeterministic(super().test_training)
+        run_nondeterministic(super().test_training)
 
     def test_training_with_ema(self):
-        _run_nondeterministic(super().test_training_with_ema)
+        run_nondeterministic(super().test_training_with_ema)
 
     def test_mixed_precision_training(self):
-        _run_nondeterministic(super().test_mixed_precision_training)
+        run_nondeterministic(super().test_mixed_precision_training)
 
     @pytest.mark.skip(
         "Gradient checkpointing recomputes the forward pass, but the model uses a stateful cache_dict "
@@ -112,8 +110,10 @@ def test_gradient_checkpointing_equivalence(self):
 class TestAutoencoderKLKVAEVideoMemory(AutoencoderKLKVAEVideoTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for AutoencoderKLKVAEVideo."""
 
+    # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation, so every test with a backward
+    # pass runs with determinism relaxed.
     def test_layerwise_casting_training(self):
-        _run_nondeterministic(super().test_layerwise_casting_training)
+        run_nondeterministic(super().test_layerwise_casting_training)
 
 
 class TestAutoencoderKLKVAEVideoSlicingTiling(AutoencoderKLKVAEVideoTesterConfig, NewAutoencoderTesterMixin):
 
@@ -25,8 +25,8 @@
     BaseModelTesterConfig,
     MemoryTesterMixin,
     ModelTesterMixin,
-    TorchCompileTesterMixin,
     TrainingTesterMixin,
+    run_nondeterministic,
 )
 from .testing_utils import NewAutoencoderTesterMixin
 
@@ -134,6 +134,18 @@ def test_encode_decode_temporal_geometry(self):
 class TestAutoencoderKLMiniMaxH3Memory(AutoencoderKLMiniMaxH3TesterConfig, MemoryTesterMixin):
     """Memory optimization tests for the MiniMax-H3 video autoencoder."""
 
+    @pytest.mark.skip(
+        "`_keep_in_fp32_modules` pins every module of this autoencoder, so layerwise casting has nothing to cast and "
+        "the memory footprint does not change."
+    )
+    def test_layerwise_casting_memory(self):
+        pass
+
+    # The encoder's reflect spatial padding has no deterministic CUDA backward
+    # (reflection_pad3d_backward_out_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_layerwise_casting_training(self):
+        run_nondeterministic(super().test_layerwise_casting_training)
+
 
 class TestAutoencoderKLMiniMaxH3Training(AutoencoderKLMiniMaxH3TesterConfig, TrainingTesterMixin):
     """Training tests for the MiniMax-H3 video autoencoder."""
@@ -143,13 +155,23 @@ def test_gradient_checkpointing_is_applied(self):
             expected_set={"MiniMaxH3VideoDownBlock3d", "MiniMaxH3VideoViTDecoder3d"}
         )
 
+    # The encoder's reflect spatial padding has no deterministic CUDA backward
+    # (reflection_pad3d_backward_out_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_training(self):
+        run_nondeterministic(super().test_training)
 
-class TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, AttentionTesterMixin):
-    """Attention processor tests for the MiniMax-H3 video autoencoder."""
+    def test_training_with_ema(self):
+        run_nondeterministic(super().test_training_with_ema)
 
+    def test_mixed_precision_training(self):
+        run_nondeterministic(super().test_mixed_precision_training)
 
-class TestAutoencoderKLMiniMaxH3TorchCompile(AutoencoderKLMiniMaxH3TesterConfig, TorchCompileTesterMixin):
-    """Torch compile tests for the MiniMax-H3 video autoencoder."""
+    def test_gradient_checkpointing_equivalence(self):
+        run_nondeterministic(super().test_gradient_checkpointing_equivalence)
+
+
+class TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, AttentionTesterMixin):
+    """Attention processor tests for the MiniMax-H3 video autoencoder."""
 
 
 class TestAutoencoderKLMiniMaxH3SlicingTiling(AutoencoderKLMiniMaxH3TesterConfig, NewAutoencoderTesterMixin):
 
@@ -25,8 +25,8 @@
     BaseModelTesterConfig,
     MemoryTesterMixin,
     ModelTesterMixin,
-    TorchCompileTesterMixin,
     TrainingTesterMixin,
+    run_nondeterministic,
 )
 
 
@@ -130,10 +130,33 @@ def test_encode_pads_to_the_hop_length(self):
 class TestAutoencoderKLMiniMaxH3AudioMemory(AutoencoderKLMiniMaxH3AudioTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for the MiniMax-H3 audio autoencoder."""
 
+    @pytest.mark.skip(
+        "`_keep_in_fp32_modules` pins every module of this autoencoder, so layerwise casting has nothing to cast and "
+        "the memory footprint does not change."
+    )
+    def test_layerwise_casting_memory(self):
+        pass
+
+    # The causal-attention projection's `F.adaptive_avg_pool1d` has no deterministic CUDA backward
+    # (adaptive_avg_pool2d_backward_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_layerwise_casting_training(self):
+        run_nondeterministic(super().test_layerwise_casting_training)
+
 
 class TestAutoencoderKLMiniMaxH3AudioTraining(AutoencoderKLMiniMaxH3AudioTesterConfig, TrainingTesterMixin):
     """Training tests for the MiniMax-H3 audio autoencoder."""
 
+    # The causal-attention projection's `F.adaptive_avg_pool1d` has no deterministic CUDA backward
+    # (adaptive_avg_pool2d_backward_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_training(self):
+        run_nondeterministic(super().test_training)
+
+    def test_training_with_ema(self):
+        run_nondeterministic(super().test_training_with_ema)
+
+    def test_mixed_precision_training(self):
+        run_nondeterministic(super().test_mixed_precision_training)
+
 
 class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTesterConfig, AttentionTesterMixin):
     """Attention processor tests for the MiniMax-H3 audio autoencoder."""
@@ -144,7 +167,3 @@ class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTester
     )
     def test_attention_processor_count_mismatch_raises_error(self):
         pass
-
-
-class TestAutoencoderKLMiniMaxH3AudioTorchCompile(AutoencoderKLMiniMaxH3AudioTesterConfig, TorchCompileTesterMixin):
-    """Torch compile tests for the MiniMax-H3 audio autoencoder."""
 
@@ -19,23 +19,19 @@
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin
+from ..testing_utils import (
+    BaseModelTesterConfig,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TrainingTesterMixin,
+    run_nondeterministic,
+)
 from .testing_utils import NewAutoencoderTesterMixin
 
 
 enable_full_determinism()
 
 
-def _run_nondeterministic(fn):
-    # avg_pool3d_backward_cuda has no deterministic CUDA implementation;
-    # temporarily relax the requirement for tests that do backward passes.
-    torch.use_deterministic_algorithms(False)
-    try:
-        fn()
-    finally:
-        torch.use_deterministic_algorithms(True)
-
-
 class AutoencoderVidTokTesterConfig(BaseModelTesterConfig):
     @property
     def model_class(self):
@@ -90,24 +86,28 @@ def test_gradient_checkpointing_is_applied(self):
         expected_set = {"VidTokEncoder3D", "VidTokDecoder3D"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
+    # avg_pool3d_backward_cuda has no deterministic CUDA implementation, so every test with a backward pass runs
+    # with determinism relaxed.
     def test_training(self):
-        _run_nondeterministic(super().test_training)
+        run_nondeterministic(super().test_training)
 
     def test_training_with_ema(self):
-        _run_nondeterministic(super().test_training_with_ema)
+        run_nondeterministic(super().test_training_with_ema)
 
     def test_mixed_precision_training(self):
-        _run_nondeterministic(super().test_mixed_precision_training)
+        run_nondeterministic(super().test_mixed_precision_training)
 
     def test_gradient_checkpointing_equivalence(self):
-        _run_nondeterministic(super().test_gradient_checkpointing_equivalence)
+        run_nondeterministic(super().test_gradient_checkpointing_equivalence)
 
 
 class TestAutoencoderVidTokMemory(AutoencoderVidTokTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for AutoencoderVidTok."""
 
+    # avg_pool3d_backward_cuda has no deterministic CUDA implementation, so every test with a backward pass runs
+    # with determinism relaxed.
     def test_layerwise_casting_training(self):
-        _run_nondeterministic(super().test_layerwise_casting_training)
+        run_nondeterministic(super().test_layerwise_casting_training)
 
 
 class TestAutoencoderVidTokSlicingTiling(AutoencoderVidTokTesterConfig, NewAutoencoderTesterMixin):
 
@@ -48,6 +48,7 @@
 )
 from .single_file import SingleFileTesterMixin
 from .training import TrainingTesterMixin
+from .utils import run_nondeterministic
 
 
 __all__ = [
@@ -94,6 +95,7 @@
     "QuantoCompileTesterMixin",
     "QuantoConfigMixin",
     "QuantoTesterMixin",
+    "run_nondeterministic",
     "SDNQCompileTesterMixin",
     "SDNQConfigMixin",
     "SDNQTesterMixin",
 
@@ -82,7 +82,7 @@ def test_torch_compile_recompilation_and_graph_break(self):
 
     @torch.no_grad()
     def test_torch_compile_repeated_blocks(self, recompile_limit=1):
-        if self.model_class._repeated_blocks is None:
+        if not self.model_class._repeated_blocks:
             pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.")
 
         init_dict = self.get_init_dict()