24 lines
761 B
Python
24 lines
761 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from diffusers import FlowMatchEulerDiscreteScheduler
|
|
|
|
|
|
class DMD2EulerScheduler(FlowMatchEulerDiscreteScheduler):
|
|
"""Euler scheduler that always uses the fixed DMD2 training timestep schedule."""
|
|
|
|
def __init__(self, *args, dmd2_timesteps: list[int], **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._dmd2_timesteps = dmd2_timesteps
|
|
|
|
def set_timesteps(
|
|
self,
|
|
num_inference_steps: int | None = None,
|
|
device: str | torch.device | None = None,
|
|
**kwargs,
|
|
) -> None:
|
|
super().set_timesteps(timesteps=self._dmd2_timesteps, device=device)
|