reafator pipeline stage/step pipeline (#1368)

Signed-off-by: asukaqaq-s <1311722138@qq.com>

vllm_omni/diffusion/models/interface.py

@@ -1,11 +1,21 @@
 # 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):
@@ -21,3 +31,34 @@ class SupportAudioInput(Protocol):
 @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)