65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from __future__ import annotations
|
|
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
ClassVar,
|
|
Protocol,
|
|
runtime_checkable,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput
|
|
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
|
|
|
|
|
|
@runtime_checkable
|
|
class SupportImageInput(Protocol):
|
|
support_image_input: ClassVar[bool] = True
|
|
color_format: ClassVar[str] = "RGB" # Default color format
|
|
|
|
|
|
@runtime_checkable
|
|
class SupportAudioInput(Protocol):
|
|
support_audio_input: ClassVar[bool] = True
|
|
|
|
|
|
@runtime_checkable
|
|
class SupportAudioOutput(Protocol):
|
|
support_audio_output: ClassVar[bool] = True
|
|
|
|
|
|
@runtime_checkable
|
|
class SupportsStepExecution(Protocol):
|
|
"""State-driven step-level execution protocol for diffusion pipelines.
|
|
|
|
Pipelines should split request-level ``forward()`` into:
|
|
``prepare_encode()`` (one-time request setup), ``denoise_step()``
|
|
(one denoise forward), ``step_scheduler()`` (one scheduler update),
|
|
and ``post_decode()`` (final decode).
|
|
"""
|
|
|
|
supports_step_execution: ClassVar[bool] = True
|
|
|
|
def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState:
|
|
"""Prepare request-level inputs and return initialized state."""
|
|
|
|
def denoise_step(self, state: DiffusionRequestState, **kwargs: Any) -> torch.Tensor | None:
|
|
"""Run one denoise step."""
|
|
|
|
def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
|
|
"""Run one scheduler step."""
|
|
|
|
def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput:
|
|
"""Decode output after denoise loop."""
|
|
|
|
|
|
def supports_step_execution(pipeline: object) -> bool:
|
|
"""Return whether `pipeline` implements :class:`SupportsStepExecution`."""
|
|
|
|
return isinstance(pipeline, SupportsStepExecution)
|