diff --git a/docs/source/en/api/pipelines/echo.md b/docs/source/en/api/pipelines/echo.md index 47b186e00163..605ab44541b2 100644 --- a/docs/source/en/api/pipelines/echo.md +++ b/docs/source/en/api/pipelines/echo.md @@ -86,8 +86,8 @@ 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. To get the denormalized VAE tensors +without decoding, run the pipeline without the decode step. ## EchoModularPipeline diff --git a/src/diffusers/modular_pipelines/cosmos/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py index 688c74629d77..5117848b63be 100644 --- a/src/diffusers/modular_pipelines/cosmos/decoders.py +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -3,7 +3,7 @@ from ...configuration_utils import FrozenDict from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan -from ...utils import logging +from ...utils import deprecate, logging from ...video_processor import VideoProcessor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam @@ -48,23 +48,30 @@ def __call__( self, components: Cosmos3OmniModularPipeline, state: PipelineState ) -> tuple[Cosmos3OmniModularPipeline, PipelineState]: block_state = self.get_block_state(state) + if block_state.output_type == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, run " + "the pipeline without the decode step and use `output='latent'`.", + ) + block_state.videos = block_state.latents + self.set_block_state(state, block_state) + return components, state 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 " @@ -245,13 +252,24 @@ def __call__( ) -> tuple[Cosmos3OmniModularPipeline, PipelineState]: block_state = self.get_block_state(state) device = components._execution_device - + if block_state.output_type == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, use " + "`output='latents'` or run the pipeline without the decode step.", + ) + block_state.videos = torch.cat(block_state.output_chunks, dim=2)[:, :, : block_state.total_frames] + block_state.sound = None + block_state.sampling_rate = None + self.set_block_state(state, block_state) + return components, state 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..43392d87e251 100644 --- a/src/diffusers/modular_pipelines/echo/decoders.py +++ b/src/diffusers/modular_pipelines/echo/decoders.py @@ -19,6 +19,7 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video from ...pipelines.ltx2.vocoder import LTX2Vocoder +from ...utils import deprecate from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..modular_pipeline import ModularPipelineBlocks, PipelineState @@ -90,19 +91,24 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - vae = components.vae - - latents = block_state.latents - if block_state.output_type == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, use " + "`output='latents'` or run the pipeline without the decode step.", + ) block_state.videos = _denormalize_latents( - latents, + block_state.latents, components.latents_mean, components.latents_std, components.vae_scaling_factor, ) self.set_block_state(state, block_state) return components, state + vae = components.vae + + latents = block_state.latents if not vae.config.timestep_conditioning: timestep = None @@ -184,6 +190,12 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ) if block_state.output_type == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, use " + "`output='latents'` or run the pipeline without the decode step.", + ) block_state.audio = audio_latents else: audio_latents = audio_latents.to(audio_vae.dtype) diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 573adbb7ff2f..a40cb2e0e131 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/src/diffusers/modular_pipelines/flux/decoders.py @@ -136,25 +136,32 @@ 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 == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, use " + "`output='latents'` or run the pipeline without the decode step.", + ) + block_state.images = block_state.latents + self.set_block_state(state, block_state) + return components, state vae = components.vae - if not block_state.output_type == "latent": - latents = block_state.latents - if latents.ndim == 3: - deprecate( - "packed latents", - "0.47.0", - "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group " - "now unpacks them. Pass latents of shape `[B, C, H, W]` instead.", - ) - 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 + latents = block_state.latents + if latents.ndim == 3: + deprecate( + "packed latents", + "0.47.0", + "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group " + "now unpacks them. Pass latents of shape `[B, C, H, W]` instead.", ) - else: - block_state.images = 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..b940f01b57d4 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py @@ -18,7 +18,7 @@ from ...configuration_utils import FrozenDict from ...image_processor import VaeImageProcessor from ...models import AutoencoderKL -from ...utils import logging +from ...utils import deprecate, logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import StableDiffusion3ModularPipeline @@ -65,16 +65,23 @@ 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) - 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 + if block_state.output_type == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, use " + "`output='latents'` or run the pipeline without the decode step.", ) - else: block_state.images = block_state.latents + self.set_block_state(state, block_state) + return components, state + vae = components.vae + + 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..390fb1cf8c58 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py @@ -87,48 +87,55 @@ def upcast_vae(components): @torch.no_grad() def __call__(self, components, state: PipelineState) -> tuple[StableDiffusionXLModularPipeline, PipelineState]: block_state = self.get_block_state(state) - - 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 + if block_state.output_type == "latent": + deprecate( + "output_type='latent'", + "0.47.0", + "Passing `output_type='latent'` to the decode step is deprecated. To get the latents, use " + "`output='latents'` or run the pipeline without the decode step.", ) - block_state.has_latents_std = ( - hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None + block_state.images = block_state.latents + self.set_block_state(state, block_state) + return components, state + + 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) ) - 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 + 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 - 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 5a2723b84625..244ad91fddee 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -154,7 +154,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 cf3a38a66b7c..665c42ebb9e8 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -91,7 +91,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) ) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 0d076477ebbe..9fc648399184 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -93,6 +93,18 @@ def test_decode_accepts_packed_latents_with_deprecation(self): assert images.shape == (1, 3, 8, 8) + def test_decode_output_type_latent_is_deprecated(self): + # `output_type="latent"` still returns the latents unchanged for a deprecation window. + decoder = self.pipeline_blocks_class().sub_blocks["decode"].init_pipeline(self.pretrained_model_name_or_path) + decoder.load_components(dtype=torch.float32) + decoder.to(torch_device) + + latents = torch.randn(1, 1, 8, 8, generator=self.get_generator()).to(torch_device) + with pytest.warns(FutureWarning, match="output_type='latent'"): + images = decoder(latents=latents, output_type="latent", output="images") + + assert torch.equal(images, latents) + def test_float16_inference(self): super().test_float16_inference(9e-2)