Python function

run_denoising_step

run_denoising_step()

max.pipelines.modeling.base.run_denoising_step(step, cache_state, cache_config, device, compute_fn, taylorseer=None)

source

Execute one denoising step with caching logic.

This is a standalone version of DiffusionPipeline.run_denoising_step that does not require inheritance. The caller provides a compute_fn callback that runs the transformer and returns the raw result tuple.

Parameters:

• step (int) – Current step index.
• cache_state (DenoisingCacheState) – Per-request mutable cache state for this stream.
• cache_config (DenoisingCacheConfig) – Cache configuration.
• device (Device) – Target device.
• compute_fn (Callable[[], tuple[Tensor, ...]]) – Callable that runs the transformer and returns the result tuple. The tuple format depends on the cache mode: (noise_pred,) for standard, (new_residual, noise_pred) for FBCache.
• taylorseer (TaylorSeer | None) – Optional TaylorSeer instance. Required when cache_config.taylorseer is True.

Returns:

noise_pred tensor for this step.

Return type:

Tensor

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit