diff --git a/src/diffusers/modular_pipelines/minimax_h3/decoders.py b/src/diffusers/modular_pipelines/minimax_h3/decoders.py index bdc7adf7c5e7..200a25e7255b 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/decoders.py +++ b/src/diffusers/modular_pipelines/minimax_h3/decoders.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import torch from ...configuration_utils import FrozenDict @@ -190,8 +191,25 @@ def __call__( video = components.vae.decode(latents, return_dict=False)[0] pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1) pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1) - video = (video.float() * pixel_std + pixel_mean).clamp(0, 1) - block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) + # Keep the FP32 pixel intermediates bounded, including when the final output lives on the CPU. + for start in range(0, video.shape[2], 16): + frames = (video[:, :, start : start + 16].float() * pixel_std + pixel_mean).clamp(0, 1) + frames = components.video_processor.postprocess_video(frames, output_type=block_state.output_type) + if start == 0: + if block_state.output_type == "pt": + videos = frames.new_empty((frames.shape[0], video.shape[2], *frames.shape[2:])) + elif block_state.output_type == "np": + videos = np.empty((frames.shape[0], video.shape[2], *frames.shape[2:]), dtype=frames.dtype) + else: + videos = [[] for _ in range(video.shape[0])] + + if block_state.output_type == "pil": + for batch_idx, batch_frames in enumerate(frames): + videos[batch_idx].extend(batch_frames) + else: + videos[:, start : start + frames.shape[1]] = frames + + block_state.videos = videos self.set_block_state(state, block_state) return components, state diff --git a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py index a5eb5095366e..af67e703c434 100644 --- a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py +++ b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py @@ -20,7 +20,9 @@ import safetensors.torch import torch from PIL import Image +from torch.utils._python_dispatch import TorchDispatchMode +from diffusers import AutoencoderKLMiniMaxH3 from diffusers.modular_pipelines import ( MiniMaxH3Blocks, MiniMaxH3ModularPipeline, @@ -31,6 +33,7 @@ MiniMaxH3VideoReference, ) from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep +from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3VideoDecodeStep from diffusers.modular_pipelines.minimax_h3.encoders import ( MiniMaxH3FL2VATextEncoderStep, MiniMaxH3KeyframeVaeEncoderStep, @@ -39,6 +42,7 @@ ) from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MINIMAX_H3_FPS from diffusers.utils import is_peft_available, logging +from diffusers.video_processor import VideoProcessor from ...testing_utils import CaptureLogger from ..testing_utils import ( @@ -963,6 +967,95 @@ class TestMiniMaxH3Ref2VAModularPipelineMemory(MiniMaxH3Ref2VAModularPipelineTes pass +class TestMiniMaxH3VideoDecode: + def get_pipeline(self, dtype, do_normalize=False): + torch.manual_seed(0) + vae = ( + AutoencoderKLMiniMaxH3( + latent_channels=4, + block_out_channels=(8, 16), + layers_per_block=1, + spatial_downsample_factors=(2, 2), + temporal_downsample_factors=(2, 2), + norm_num_groups=8, + decoder_num_layers=2, + decoder_num_attention_heads=2, + decoder_attention_head_dim=8, + decoder_num_register_tokens=2, + decoder_ffn_mult=2, + latents_mean=(0.0,) * 4, + latents_std=(1.0,) * 4, + ) + .to(dtype=dtype) + .eval() + ) + pipe = MiniMaxH3VideoDecodeStep().init_pipeline() + pipe.update_components(vae=vae, video_processor=VideoProcessor(do_normalize=do_normalize)) + return pipe + + @pytest.mark.parametrize( + "output_type,dtype,do_normalize", + [ + ("pt", torch.float32, False), + ("np", torch.float16, False), + ("pil", torch.bfloat16, False), + ("pt", torch.float16, True), + ("np", torch.bfloat16, True), + ("pil", torch.float32, True), + ], + ) + @torch.no_grad() + def test_video_postprocess_outputs(self, output_type, dtype, do_normalize): + pipe = self.get_pipeline(dtype, do_normalize) + latents = torch.randn(2, 4, 7, 2, 2, generator=torch.Generator().manual_seed(0)) + decoded = pipe.vae.decode(latents, return_dict=False)[0] + pixel_mean = torch.tensor(pipe.pixel_mean).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(pipe.pixel_std).view(1, -1, 1, 1, 1) + expected = pipe.video_processor.postprocess_video( + (decoded.float() * pixel_std + pixel_mean).clamp(0, 1), output_type=output_type + ) + + actual = pipe(latents=latents, output_type=output_type, output="videos") + + if output_type == "pt": + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.is_contiguous() + elif output_type == "np": + np.testing.assert_array_equal(actual, expected) + assert actual.dtype == expected.dtype + else: + assert len(actual) == len(expected) + for actual_video, expected_video in zip(actual, expected): + assert len(actual_video) == len(expected_video) + for actual_frame, expected_frame in zip(actual_video, expected_video): + np.testing.assert_array_equal(np.asarray(actual_frame), np.asarray(expected_frame)) + + def test_video_postprocess_fp32_storage_is_bounded(self): + pipe = self.get_pipeline(torch.float16) + latents = torch.randn(2, 4, 7, 2, 2, generator=torch.Generator().manual_seed(0)) + pixel_storage_bytes = [] + + class PixelStorageMode(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + result = func(*args, **(kwargs or {})) + if ( + isinstance(result, torch.Tensor) + and result.dtype == torch.float32 + and result.ndim == 5 + and result.shape[:2] == (2, 3) + and result.shape[-2:] == (8, 8) + ): + pixel_storage_bytes.append(result.untyped_storage().nbytes()) + return result + + with PixelStorageMode(): + videos = pipe(latents=latents, output_type="np", output="videos") + + assert videos.shape == (2, 22, 8, 8, 3) + assert pixel_storage_bytes + assert max(pixel_storage_bytes) <= 2 * 3 * 16 * 8 * 8 * 4 + + class TestMiniMaxH3Reference: """ The three reference dataclasses and the passes that normalize them, neither of which needs a checkpoint.