diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 9a5e395ae..2bc3074b5 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -4,7 +4,7 @@ import multiprocessing as mp import time import weakref -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import Any @@ -361,3 +361,8 @@ class DiffusionEngine: def close(self) -> None: self._finalizer() + + def abort(self, request_id: str | Iterable[str]) -> None: + # TODO implement it + logger.warning("DiffusionEngine abort is not implemented yet") + pass diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 0a4976a37..1a7f1174c 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -10,7 +10,7 @@ enabling concurrent request handling and streaming generation. import asyncio import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Iterable from concurrent.futures import ThreadPoolExecutor from dataclasses import fields from typing import Any @@ -263,6 +263,10 @@ class AsyncOmniDiffusion: except Exception: pass + async def abort(self, request_id: str | Iterable[str]) -> None: + """Abort a request.""" + self.engine.abort(request_id) + @property def is_running(self) -> bool: """Check if the engine is running."""