diff --git a/docs/source/en/api/pipelines/echo.md b/docs/source/en/api/pipelines/echo.md index 47b186e00163..9d22d6724763 100644 --- a/docs/source/en/api/pipelines/echo.md +++ b/docs/source/en/api/pipelines/echo.md @@ -86,8 +86,7 @@ Modular workflows can cache the VAE encoder's normalized, unpacked tensors. Vide The core `denoise` block packs those tensors, expands conditioning for `num_videos_per_prompt`, and unpacks denoised outputs back to the same VAE form. Pass initial `latents` and `audio_latents` in that unpacked form too. Decoders take normalized VAE tensors and denormalize immediately before -decoding. `output=["latents", "audio_latents"]` returns normalized VAE tensors. With `output_type="latent"` and -`output=["videos", "audio"]`, you get the denormalized VAE tensors without decoding. +decoding. `output=["latents", "audio_latents"]` returns normalized VAE tensors. ## EchoModularPipeline diff --git a/src/diffusers/modular_pipelines/cosmos/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py index 688c74629d77..e4b54d4827b5 100644 --- a/src/diffusers/modular_pipelines/cosmos/decoders.py +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -48,23 +48,25 @@ def __call__( self, components: Cosmos3OmniModularPipeline, state: PipelineState ) -> tuple[Cosmos3OmniModularPipeline, PipelineState]: block_state = self.get_block_state(state) + if block_state.output_type not in ("pil", "np", "pt"): + raise ValueError( + f"`output_type` must be one of 'pil', 'np' or 'pt', got {block_state.output_type!r}. To get the " + "latents, use `output='latents'` or run the pipeline without the decode step." + ) device = components._execution_device - if block_state.output_type == "latent": - block_state.videos = block_state.latents - else: - in_dtype = block_state.latents.dtype - vae_dtype = components.vae.dtype - mean = components._vae_latents_mean.to(device=block_state.latents.device, dtype=vae_dtype) - inv_std = components._vae_latents_inv_std.to(device=block_state.latents.device, dtype=vae_dtype) - z_raw = block_state.latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1) - decoded = components.vae.decode(z_raw).sample.to(in_dtype) - block_state.videos = components.video_processor.postprocess_video( - decoded, output_type=block_state.output_type - )[0] - - if components.requires_safety_checker and block_state.output_type != "latent": + in_dtype = block_state.latents.dtype + vae_dtype = components.vae.dtype + mean = components._vae_latents_mean.to(device=block_state.latents.device, dtype=vae_dtype) + inv_std = components._vae_latents_inv_std.to(device=block_state.latents.device, dtype=vae_dtype) + z_raw = block_state.latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1) + decoded = components.vae.decode(z_raw).sample.to(in_dtype) + block_state.videos = components.video_processor.postprocess_video( + decoded, output_type=block_state.output_type + )[0] + + if components.requires_safety_checker: if getattr(components, "safety_checker", None) is None: raise ValueError( "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " @@ -244,14 +246,19 @@ def __call__( self, components: Cosmos3OmniModularPipeline, state: PipelineState ) -> tuple[Cosmos3OmniModularPipeline, PipelineState]: block_state = self.get_block_state(state) - device = components._execution_device + device = components._execution_device + if block_state.output_type not in ("pil", "np", "pt"): + raise ValueError( + f"`output_type` must be one of 'pil', 'np' or 'pt', got {block_state.output_type!r}. To get the " + "latents, use `output='latents'` or run the pipeline without the decode step." + ) decoded = torch.cat(block_state.output_chunks, dim=2)[:, :, : block_state.total_frames] block_state.videos = components.video_processor.postprocess_video( decoded, output_type=block_state.output_type )[0] - if components.requires_safety_checker and block_state.output_type != "latent": + if components.requires_safety_checker: if getattr(components, "safety_checker", None) is None: raise ValueError( "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " diff --git a/src/diffusers/modular_pipelines/echo/decoders.py b/src/diffusers/modular_pipelines/echo/decoders.py index 846f1ec559d5..feeddef24711 100644 --- a/src/diffusers/modular_pipelines/echo/decoders.py +++ b/src/diffusers/modular_pipelines/echo/decoders.py @@ -90,20 +90,15 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + if block_state.output_type not in ("pil", "np", "pt"): + raise ValueError( + f"`output_type` must be one of 'pil', 'np' or 'pt', got {block_state.output_type!r}. To get the " + "latents, use `output='latents'` or run the pipeline without the decode step." + ) vae = components.vae latents = block_state.latents - if block_state.output_type == "latent": - block_state.videos = _denormalize_latents( - latents, - components.latents_mean, - components.latents_std, - components.vae_scaling_factor, - ) - self.set_block_state(state, block_state) - return components, state - if not vae.config.timestep_conditioning: timestep = None else: @@ -165,7 +160,6 @@ def inputs(self) -> list[tuple[str, Any]]: required=True, description="Normalized audio VAE latents of shape (B, C, L, M).", ), - InputParam.template("output_type", default="pil"), ] @property @@ -183,12 +177,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: components.audio_latents_std, ) - if block_state.output_type == "latent": - block_state.audio = audio_latents - else: - audio_latents = audio_latents.to(audio_vae.dtype) - generated_mel_spectrograms = audio_vae.decode(audio_latents, return_dict=False)[0] - block_state.audio = components.vocoder(generated_mel_spectrograms.to(components.vocoder.dtype)) + audio_latents = audio_latents.to(audio_vae.dtype) + generated_mel_spectrograms = audio_vae.decode(audio_latents, return_dict=False)[0] + block_state.audio = components.vocoder(generated_mel_spectrograms.to(components.vocoder.dtype)) self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 796bff258a6b..82fb336f7f99 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/src/diffusers/modular_pipelines/flux/decoders.py @@ -92,18 +92,20 @@ def intermediate_outputs(self) -> list[str]: @torch.no_grad() def __call__(self, components, state: PipelineState) -> tuple[FluxModularPipeline, PipelineState]: block_state = self.get_block_state(state) + if block_state.output_type not in ("pil", "np", "pt"): + raise ValueError( + f"`output_type` must be one of 'pil', 'np' or 'pt', got {block_state.output_type!r}. To get the " + "latents, use `output='latents'` or run the pipeline without the decode step." + ) vae = components.vae - if not block_state.output_type == "latent": - latents = block_state.latents - latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor) - latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor - block_state.images = vae.decode(latents, return_dict=False)[0] - block_state.images = components.image_processor.postprocess( - block_state.images, output_type=block_state.output_type - ) - else: - block_state.images = block_state.latents + latents = block_state.latents + latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor) + latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor + block_state.images = vae.decode(latents, return_dict=False)[0] + block_state.images = components.image_processor.postprocess( + block_state.images, output_type=block_state.output_type + ) self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py b/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py index f2665ed06e47..63dcafb12ff8 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py @@ -65,16 +65,18 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components, state: PipelineState) -> tuple[StableDiffusion3ModularPipeline, PipelineState]: block_state = self.get_block_state(state) + if block_state.output_type not in ("pil", "np", "pt"): + raise ValueError( + f"`output_type` must be one of 'pil', 'np' or 'pt', got {block_state.output_type!r}. To get the " + "latents, use `output='latents'` or run the pipeline without the decode step." + ) vae = components.vae - if not block_state.output_type == "latent": - latents = (block_state.latents / vae.config.scaling_factor) + vae.config.shift_factor - block_state.images = vae.decode(latents, return_dict=False)[0] - block_state.images = components.image_processor.postprocess( - block_state.images, output_type=block_state.output_type - ) - else: - block_state.images = block_state.latents + latents = (block_state.latents / vae.config.scaling_factor) + vae.config.shift_factor + block_state.images = vae.decode(latents, return_dict=False)[0] + block_state.images = components.image_processor.postprocess( + block_state.images, output_type=block_state.output_type + ) self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py index ea6fdfa1df98..72f68c1d8fd4 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py @@ -87,48 +87,50 @@ def upcast_vae(components): @torch.no_grad() def __call__(self, components, state: PipelineState) -> tuple[StableDiffusionXLModularPipeline, PipelineState]: block_state = self.get_block_state(state) + if block_state.output_type not in ("pil", "np", "pt"): + raise ValueError( + f"`output_type` must be one of 'pil', 'np' or 'pt', got {block_state.output_type!r}. To get the " + "latents, use `output='latents'` or run the pipeline without the decode step." + ) - if not block_state.output_type == "latent": - latents = block_state.latents - # make sure the VAE is in float32 mode, as it overflows in float16 - block_state.needs_upcasting = components.vae.dtype == torch.float16 and components.vae.config.force_upcast - - if block_state.needs_upcasting: - self.upcast_vae(components) - latents = latents.to(next(iter(components.vae.post_quant_conv.parameters())).dtype) - elif latents.dtype != components.vae.dtype: - if torch.backends.mps.is_available(): - # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 - components.vae = components.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - block_state.has_latents_mean = ( - hasattr(components.vae.config, "latents_mean") and components.vae.config.latents_mean is not None + latents = block_state.latents + # make sure the VAE is in float32 mode, as it overflows in float16 + block_state.needs_upcasting = components.vae.dtype == torch.float16 and components.vae.config.force_upcast + + if block_state.needs_upcasting: + self.upcast_vae(components) + latents = latents.to(next(iter(components.vae.post_quant_conv.parameters())).dtype) + elif latents.dtype != components.vae.dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + components.vae = components.vae.to(latents.dtype) + + # unscale/denormalize the latents + # denormalize with the mean and std if available and not None + block_state.has_latents_mean = ( + hasattr(components.vae.config, "latents_mean") and components.vae.config.latents_mean is not None + ) + block_state.has_latents_std = ( + hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None + ) + if block_state.has_latents_mean and block_state.has_latents_std: + block_state.latents_mean = ( + torch.tensor(components.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) ) - block_state.has_latents_std = ( - hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None + block_state.latents_std = ( + torch.tensor(components.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) ) - if block_state.has_latents_mean and block_state.has_latents_std: - block_state.latents_mean = ( - torch.tensor(components.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - block_state.latents_std = ( - torch.tensor(components.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = ( - latents * block_state.latents_std / components.vae.config.scaling_factor + block_state.latents_mean - ) - else: - latents = latents / components.vae.config.scaling_factor + latents = ( + latents * block_state.latents_std / components.vae.config.scaling_factor + block_state.latents_mean + ) + else: + latents = latents / components.vae.config.scaling_factor - block_state.images = components.vae.decode(latents, return_dict=False)[0] + block_state.images = components.vae.decode(latents, return_dict=False)[0] - # cast back to fp16 if needed - if block_state.needs_upcasting: - components.vae.to(dtype=torch.float16) - else: - block_state.images = block_state.latents + # cast back to fp16 if needed + if block_state.needs_upcasting: + components.vae.to(dtype=torch.float16) # apply watermark if available if hasattr(components, "watermark") and components.watermark is not None: diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index c3bcb267758f..3f276707b6e9 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -153,7 +153,7 @@ def get_dummy_inputs(self, seed=0): "height": 32, "width": 32, "num_frames": 5, - "output_type": "latent", + "output_type": "pt", } diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py index 86bc6ad4d6f2..1899a790a523 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -90,7 +90,7 @@ def get_dummy_inputs(self, seed=0): "height": 32, "width": 32, "num_frames": 5, - "output_type": "latent", + "output_type": "pt", } diff --git a/tests/modular_pipelines/echo/test_modular_pipeline_echo.py b/tests/modular_pipelines/echo/test_modular_pipeline_echo.py index 68a18d2e5078..e0f53b2d80ce 100644 --- a/tests/modular_pipelines/echo/test_modular_pipeline_echo.py +++ b/tests/modular_pipelines/echo/test_modular_pipeline_echo.py @@ -19,7 +19,7 @@ import torch from diffusers.modular_pipelines.echo.before_denoise import EchoInputsStep, EchoPrepareConditioningStep -from diffusers.modular_pipelines.echo.decoders import EchoAudioDecoderStep +from diffusers.modular_pipelines.echo.decoders import _denormalize_audio_latents from diffusers.modular_pipelines.echo.encoders import EchoVaeEncoderStep from diffusers.modular_pipelines.echo.modular_blocks_echo import EchoDecoderStep @@ -73,15 +73,13 @@ def test_video_and_audio_outputs(self): def test_seeded_dmd_renoising_is_deterministic(self): pipe = self.get_pipeline() inputs = self.get_dummy_inputs(seed=1) - inputs["output_type"] = "latent" - first = pipe(**inputs, output=["videos", "audio"]) + first = pipe(**inputs, output=["latents", "audio_latents"]) inputs = self.get_dummy_inputs(seed=1) - inputs["output_type"] = "latent" - second = pipe(**inputs, output=["videos", "audio"]) + second = pipe(**inputs, output=["latents", "audio_latents"]) - assert torch.equal(first["videos"], second["videos"]) - assert torch.equal(first["audio"], second["audio"]) + assert torch.equal(first["latents"], second["latents"]) + assert torch.equal(first["audio_latents"], second["audio_latents"]) def test_audio_decoder_supports_mixed_precision(self): pipe = self.get_pipeline() @@ -94,7 +92,6 @@ def test_audio_decoder_supports_mixed_precision(self): def test_user_latents_are_not_modified(self): pipe = self.get_pipeline() inputs = self.get_dummy_inputs() - inputs["output_type"] = "latent" latent_num_frames = 1 + (inputs["num_frames"] - 1) // pipe.vae_temporal_compression_ratio latent_height = inputs["height"] // pipe.vae_spatial_compression_ratio latent_width = inputs["width"] // pipe.vae_spatial_compression_ratio @@ -115,7 +112,6 @@ def test_raw_audio_memory(self, waveform_lengths): None if length is None else torch.rand(2, length) for length in waveform_lengths ] inputs["memory_audio_sample_rates"] = 16000 - inputs["output_type"] = "latent" output = pipe(**inputs, output=["videos", "audio"]) @@ -196,10 +192,9 @@ def test_audio_encoder_decoder_use_channel_mel_statistics(self, dtype): expected = (raw_latents.transpose(1, 2).flatten(2) - mean) / std torch.testing.assert_close(normalized.transpose(1, 2).flatten(2), expected.float(), rtol=0, atol=0) - decoder = EchoAudioDecoderStep().init_pipeline(self.pretrained_model_name_or_path) - decoder.load_components(dtype=torch.float32) - decoder.update_components(audio_vae=normalized_vae) - decoded = decoder(audio_latents=normalized.to(dtype), output_type="latent", output="audio") + decoded = _denormalize_audio_latents( + normalized.to(dtype), normalized_vae.latents_mean, normalized_vae.latents_std + ) expected_decoded = ( (expected * std + mean).unflatten(2, (raw_latents.shape[1], raw_latents.shape[3])).transpose(1, 2) )