Ayush AgarwalandHongsheng Liu
b968373c88
Signed-off-by: ayushag <ayushag@nvidia.com> Signed-off-by: Ayush Agarwal <ayushag@nvidia.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Literal, get_args
|
|
|
|
import torch
|
|
|
|
Solver = Literal["ode", "sde"]
|
|
|
|
|
|
@dataclass
|
|
class DMD2Config:
|
|
"""Inference-time contract for a FastGen DMD2-distilled checkpoint."""
|
|
|
|
num_inference_steps: int = 4
|
|
denoising_timesteps: list[int] | None = None
|
|
solver: Solver = "ode"
|
|
guidance_scale: float = 1.0
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.solver not in get_args(Solver):
|
|
raise ValueError(f"DMD2Config.solver must be one of {list(get_args(Solver))}, got {self.solver!r}")
|
|
|
|
@classmethod
|
|
def from_model_index(cls, model_index: dict) -> DMD2Config:
|
|
"""Read the `dmd2_config` block from a model_index.json dict. Missing block → defaults."""
|
|
block = model_index.get("dmd2_config", {})
|
|
solver = block.get("solver", cls.solver)
|
|
if isinstance(solver, str):
|
|
solver = solver.strip().lower()
|
|
return cls(
|
|
num_inference_steps=block.get("num_inference_steps", cls.num_inference_steps),
|
|
denoising_timesteps=block.get("denoising_timesteps"),
|
|
solver=solver,
|
|
guidance_scale=block.get("guidance_scale", cls.guidance_scale),
|
|
)
|
|
|
|
def resolve_timesteps(self) -> list[int]:
|
|
if self.denoising_timesteps is not None:
|
|
return list(self.denoising_timesteps)
|
|
# Uniformly spaced timesteps from 999 down toward 0, excluding the final 0.
|
|
ts = torch.linspace(999, 0, self.num_inference_steps + 1)[:-1]
|
|
return ts.to(torch.int32).tolist()
|