Fast ODE-based Sampling for Diffusion Models in Around 5 Steps

[remyx-ai]

Model/Pipeline/Scheduler description

Add an AMEDScheduler implementing the Approximate Mean-Direction (AMED) solver from Fast ODE-based Sampling for Diffusion Models in Around 5 Steps (Zhenyu Zhou, Defang Chen, Can Wang, Chun Chen).

Each sampling step takes a single Euler hop to the geometric-mean intermediate noise level, evaluates the model there, and applies that mean-direction estimate across the whole step. This cancels the leading truncation error of an Euler step and produces high-quality samples in around five steps.

The proposal targets the training-free instantiation of the solver: the output-scaling coefficient c from the paper's plugin network defaults to 1.0 and is exposed through a scale_factor config argument for users who distill their own value. It reuses the noise-schedule construction and helpers of HeunDiscreteScheduler and can be selected for any Karras-style pipeline with AMEDScheduler.from_config(pipe.scheduler.config).

As a SchedulerMixin/ConfigMixin subclass selected via from_config, it is immediately reachable by all 273 existing self.scheduler.step() call sites with no pipeline changes, and it mirrors the existing solver-family pattern in the repo.

Open source status

• The model implementation is available.
• The model weights are available (Only relevant if addition is not a scheduler). — N/A: this is a scheduler.

Provide useful links for the implementation

• Reference implementation: https://github.com/zju-pi/diff-sampler — license Apache-2.0 (permissive)
• Paper: https://arxiv.org/abs/2312.00094v3
• Authors: Zhenyu Zhou, Defang Chen, Can Wang, Chun Chen

---

Why this is an Issue and not a PR

A coding agent wrote a working draft of the scheduler (new scheduling_amed.py, registry wiring in src/diffusers/__init__.py and schedulers/__init__.py, and an amed.md doc page) before the downgrade gate fired, so it is shared here as a starting point rather than a submitted PR. The draft can be applied locally with git apply from the block in the Discovery-context section below.

Drafted by Outrider — paper: arXiv:2312.00094.

Discovery context

Provenance

• Recommended paper: Fast ODE-based Sampling for Diffusion Models in Around 5 Steps
• Confidence: moderate (Remyx relevance 0.65)
• Research interest: diffusers
• License: Apache-2.0 (class: permissive, compat: 1.00, source: github) 🟢 Permissive — safe to adopt.

Why this candidate (selected from the lookback pool)

diffusers schedulers are the cleanest contract anchor in the repo: a new fast ODE solver subclasses SchedulerMixin/ConfigMixin and is selected via scheduler.from_config, immediately reachable by all 273 existing self.scheduler.step() call sites with no pipeline changes. [29] is the only code-bearing candidate (Apache-2.0) that fits this module, and the scheduler/few-step-sampling acceleration theme is an explicit interest (it was surfaced by the 'scheduler solver few-step sampling ODE diffusion acceleration' refine query). It mirrors the existing solver family pattern exactly, making it a clean addition.

Why this paper is interesting for the team

Surfaced by Outrider deep-search refine query scheduler solver few-step sampling ODE diffusion training-free acceleration against /search/assets. The engine's normal ranking did not place this paper in the interest's broad pool — it's here because the audit pass identified an under-represented theme this paper covers.

Proposed implementation (draft)

Apply locally with git apply after saving the block below.

Diff (476 lines)

diff --git a/docs/source/en/api/schedulers/amed.md b/docs/source/en/api/schedulers/amed.md
new file mode 100644
index 0000000..8ff69d7
--- /dev/null
+++ b/docs/source/en/api/schedulers/amed.md
@@ -0,0 +1,23 @@
+<!--Copyright 2025 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# AMEDScheduler
+
+The Approximate Mean-Direction (AMED) solver is from [Fast ODE-based Sampling for Diffusion Models in Around 5 Steps](https://huggingface.co/papers/2312.00094) by Zhenyu Zhou, Defang Chen, Can Wang, and Chun Chen. Each sampling step takes a single Euler hop to the geometric-mean intermediate noise level, evaluates the model there, and applies that mean-direction estimate across the whole step. This cancels the leading truncation error of an Euler step and produces high-quality samples in around five steps.
+
+This implementation is the training-free instantiation of the solver: the output-scaling coefficient `c` from the paper's plugin network defaults to `1.0` and is exposed through the `scale_factor` config argument for users who distill their own value. It reuses the noise-schedule construction and helpers of [`HeunDiscreteScheduler`] and can be selected for any Karras-style pipeline with `AMEDScheduler.from_config(pipe.scheduler.config)`.
+
+## AMEDScheduler
+[[autodoc]] AMEDScheduler
+
+## AMEDSchedulerOutput
+[[autodoc]] schedulers.scheduling_amed.AMEDSchedulerOutput
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
index da77fa6..1c7918a 100644
--- a/src/diffusers/__init__.py
+++ b/src/diffusers/__init__.py
@@ -382,6 +382,7 @@ else:
     _import_structure["quantizers"] = ["DiffusersQuantizer"]
     _import_structure["schedulers"].extend(
         [
+            "AMEDScheduler",
             "AmusedScheduler",
             "BlockRefinementScheduler",
             "BlockRefinementSchedulerOutput",
@@ -1237,6 +1238,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
         )
         from .quantizers import DiffusersQuantizer
         from .schedulers import (
+            AMEDScheduler,
             AmusedScheduler,
             BlockRefinementScheduler,
             BlockRefinementSchedulerOutput,
diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py
index 447586c..fecf726 100644
--- a/src/diffusers/schedulers/__init__.py
+++ b/src/diffusers/schedulers/__init__.py
@@ -39,6 +39,7 @@ except OptionalDependencyNotAvailable:
 
 else:
     _import_structure["deprecated"] = ["KarrasVeScheduler", "ScoreSdeVpScheduler"]
+    _import_structure["scheduling_amed"] = ["AMEDScheduler"]
     _import_structure["scheduling_amused"] = ["AmusedScheduler"]
     _import_structure["scheduling_block_refinement"] = ["BlockRefinementScheduler", "BlockRefinementSchedulerOutput"]
     _import_structure["scheduling_consistency_decoder"] = ["ConsistencyDecoderScheduler"]
@@ -146,6 +147,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
         from ..utils.dummy_pt_objects import *  # noqa F403
     else:
         from .deprecated import KarrasVeScheduler, ScoreSdeVpScheduler
+        from .scheduling_amed import AMEDScheduler
         from .scheduling_amused import AmusedScheduler
         from .scheduling_block_refinement import BlockRefinementScheduler, BlockRefinementSchedulerOutput
         from .scheduling_consistency_decoder import ConsistencyDecoderScheduler
diff --git a/src/diffusers/schedulers/scheduling_amed.py b/src/diffusers/schedulers/scheduling_amed.py
new file mode 100644
index 0000000..4eba31c
--- /dev/null
+++ b/src/diffusers/schedulers/scheduling_amed.py
@@ -0,0 +1,288 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Adapted from "Fast ODE-based Sampling for Diffusion Models in Around 5 Steps"
+# (Zhou et al., https://huggingface.co/papers/2312.00094). The Approximate
+# Mean-Direction (AMED) solver replaces the endpoint-averaging of a second-order
+# Heun step with a single model evaluation at an intermediate noise level chosen
+# to approximate the mean ODE direction over the step. The paper learns the
+# intermediate position and an output-scaling coefficient with a tiny plugin
+# network; this scheduler implements the training-free instantiation, where the
+# intermediate noise level is the geometric mean of the step endpoints and the
+# scaling coefficient defaults to 1.0 (and is exposed via ``scale_factor`` for
+# users who distill their own value).
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from ..configuration_utils import register_to_config
+from ..utils import BaseOutput
+from .scheduling_heun_discrete import HeunDiscreteScheduler, betas_for_alpha_bar
+
+
+@dataclass
+class AMEDSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next mod

[smellslikeml]

Closing to test landing-zone-to-PR on this fork. The successor PR will add a new scheduler class implementing the ~5-step ODE sampler — plugging into the existing 54-scheduler registry alongside DPMSolverMultistepScheduler / LCMScheduler / TCDScheduler. See feat/lead-implementation-spec + REMYX-195.