Deprecate output_type="latent" in modular decoders

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
 

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 "

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)

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)
 

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

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:

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

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