Remove output_type="latent" from modular decoders

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
 

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 "

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

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)
 

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

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:

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",
         }
 
 

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",
         }
 
 

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)
         )