[Refactor]: clean up diffusion scheduler interfaces (#5216)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

vllm_omni/diffusion/models/interface.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 
     from vllm_omni.diffusion.data import DiffusionOutput
     from vllm_omni.diffusion.worker.input_batch import InputBatch
-    from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+    from vllm_omni.diffusion.worker.utils import StepRequestState
 
 
 @runtime_checkable
@@ -54,7 +54,7 @@ class SupportsStepExecution(Protocol):
 
     supports_step_execution: ClassVar[bool] = True
 
-    def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState:
+    def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState:
         """Prepare request-level inputs and return initialized state."""
         ...
 
@@ -62,11 +62,11 @@ class SupportsStepExecution(Protocol):
         """Run one denoise forward on the runner-assembled batch."""
         ...
 
-    def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
+    def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
         """Run one scheduler step."""
         ...
 
-    def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput:
+    def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput:
         """Decode output after denoise loop or at a partial chunk boundary."""
         ...