[feat] Let the OpenAI server plug in a non-CUDA generator

Author: Ishxn20

fastvideo/entrypoints/openai/api_server.py

@@ -2,7 +2,7 @@
 # (https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py)
 
 from contextlib import asynccontextmanager
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 import os
 
 import uvicorn
@@ -18,7 +18,8 @@
     clear_state,
     set_state,
 )
-from fastvideo.entrypoints.openai.serving_engine import OpenAIServingEngine
+from fastvideo.entrypoints.openai.serving_engine import OpenAIServingEngine, ServingGenerator
+from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest
 from fastvideo.entrypoints.video_generator import VideoGenerator
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -60,10 +61,13 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     served_model_name: str | None = app.state.served_model_name
     default_request: GenerationRequest | None = getattr(app.state, "default_request", None)
 
-    logger.info("Loading model from %s ...", args.model_path)
-    generator = VideoGenerator.from_fastvideo_args(args)
-    serving_engine = OpenAIServingEngine(generator)
-    logger.info("Model loaded successfully.")
+    logger.info("Initializing %s generation runtime for %s ...", app.state.runtime, args.model_path)
+    # A non-CUDA runtime (e.g. MLX) supplies its own generator_factory instead
+    # of the default VideoGenerator/Executor path -- see MLXWanGenerator.
+    factory = app.state.generator_factory
+    generator = factory() if factory is not None else VideoGenerator.from_fastvideo_args(args)
+    serving_engine = OpenAIServingEngine(generator, app.state.video_request_validator)
+    logger.info("Generation runtime ready.")
 
     set_state(
         generator,
@@ -91,8 +95,18 @@ def create_app(
     output_dir: str = DEFAULT_OUTPUT_DIR,
     default_request: GenerationRequest | None = None,
     served_model_name: str | None = None,
+    *,
+    generator_factory: Callable[[], ServingGenerator] | None = None,
+    video_request_validator: Callable[[VideoGenerationRequest], None] | None = None,
+    runtime: str = "cuda",
 ) -> FastAPI:
-    """Build the FastAPI application with all routers mounted"""
+    """Build the FastAPI application with all routers mounted.
+
+    ``generator_factory``/``video_request_validator``/``runtime`` let a
+    non-CUDA backend (MLX) plug in its own generator and request rules
+    without a separate app-building path. Defaults reproduce the original
+    CUDA-only behavior exactly.
+    """
 
     app = FastAPI(
         title="FastVideo OpenAI-Compatible API",
@@ -103,6 +117,9 @@ def create_app(
     app.state.output_dir = output_dir
     app.state.default_request = default_request
     app.state.served_model_name = served_model_name
+    app.state.generator_factory = generator_factory
+    app.state.video_request_validator = video_request_validator
+    app.state.runtime = runtime
 
     app.add_middleware(
         CORSMiddleware,
@@ -148,7 +165,10 @@ async def openai_validation_error(_request: Request, exc: RequestValidationError
 
     app.include_router(common_router)
     app.include_router(video_router)
-    app.include_router(image_router)
+    # The MLX runtime only wires video-with-audio generation (see
+    # fastvideo/mlx_runtime/); image generation has no MLX backend yet.
+    if runtime != "mlx":
+        app.include_router(image_router)
 
     @app.get("/health")
     async def health():

fastvideo/entrypoints/openai/serving_engine.py

@@ -5,14 +5,28 @@
 
 import asyncio
 from collections.abc import Awaitable, Callable
-from typing import Any, TypeVar
+from typing import Any, Protocol, TypeVar
 
 from fastvideo.api.schema import GenerationRequest
-from fastvideo.entrypoints.video_generator import VideoGenerator
+from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest
 
 _T = TypeVar("_T")
 
 
+class ServingGenerator(Protocol):
+    """The minimal shape OpenAIServingEngine needs from a generator.
+
+    VideoGenerator (CUDA/mp/ray) satisfies this structurally already; an MLX
+    generator (e.g. MLXWanGenerator) implements just these two methods.
+    """
+
+    def generate(self, request: GenerationRequest) -> Any:
+        ...
+
+    def shutdown(self) -> None:
+        ...
+
+
 class OpenAIServingEngine:
     """Own generator lifecycle and serialize access to its mutable pipeline.
 
@@ -24,16 +38,24 @@ class OpenAIServingEngine:
     the lock without changing the transport contract.
     """
 
-    def __init__(self, generator: VideoGenerator) -> None:
+    def __init__(self,
+                 generator: ServingGenerator,
+                 video_request_validator: Callable[[VideoGenerationRequest], None] | None = None) -> None:
         self._generator = generator
+        self._video_request_validator = video_request_validator
         self._generation_lock = asyncio.Lock()
         self._closed = False
         self._unhealthy_reason: str | None = None
 
     @property
-    def generator(self) -> VideoGenerator:
+    def generator(self) -> ServingGenerator:
         return self._generator
 
+    def validate_video_request(self, request: VideoGenerationRequest) -> None:
+        """Runtime-specific request checks (e.g. MLX's supported-field allowlist)."""
+        if self._video_request_validator is not None:
+            self._video_request_validator(request)
+
     @property
     def closed(self) -> bool:
         return self._closed
@@ -127,4 +149,4 @@ async def shutdown(self) -> None:
             await asyncio.to_thread(self._generator.shutdown)
 
 
-__all__ = ["OpenAIServingEngine"]
+__all__ = ["OpenAIServingEngine", "ServingGenerator"]

fastvideo/entrypoints/openai/state.py

@@ -11,21 +11,20 @@
 
 if TYPE_CHECKING:
     from fastvideo.api.schema import GenerationRequest
-    from fastvideo.entrypoints.openai.serving_engine import OpenAIServingEngine
-    from fastvideo.entrypoints.video_generator import VideoGenerator
+    from fastvideo.entrypoints.openai.serving_engine import OpenAIServingEngine, ServingGenerator
     from fastvideo.fastvideo_args import FastVideoArgs
 
 DEFAULT_OUTPUT_DIR = "outputs"
 
-_generator: VideoGenerator | None = None
+_generator: ServingGenerator | None = None
 _serving_engine: OpenAIServingEngine | None = None
 _fastvideo_args: FastVideoArgs | None = None
 _output_dir: str = DEFAULT_OUTPUT_DIR
 _served_model_name: str | None = None
 _default_request: GenerationRequest | None = None
 
 
-def get_generator() -> VideoGenerator:
+def get_generator() -> ServingGenerator:
     """Return the global VideoGenerator instance (set during startup)."""
     assert _generator is not None, "Server not initialized — generator is None"
     return _generator
@@ -62,7 +61,7 @@ def get_default_request() -> GenerationRequest | None:
 
 
 def set_state(
-    generator: VideoGenerator,
+    generator: ServingGenerator,
     serving_engine: OpenAIServingEngine,
     fastvideo_args: FastVideoArgs,
     output_dir: str,

fastvideo/entrypoints/openai/video_api.py

@@ -331,6 +331,9 @@ async def _parse_video_request(raw_request: Request) -> VideoGenerationRequest:
 
 async def _adapt_request(request_id: str, request: VideoGenerationRequest) -> GenerationRequest:
     try:
+        # Runtime-specific checks (e.g. MLX's supported-field allowlist) run
+        # before the CUDA-oriented model/LoRA validation below.
+        get_serving_engine().validate_video_request(request)
         validate_model_and_lora(request, get_server_args(), get_served_model_name())
         await prepare_reference_media(request_id, request, get_output_dir())
         return await asyncio.to_thread(