diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f73cf3c8d..c52038465 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -23,6 +23,7 @@ import vllm.envs as envs from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from PIL import Image +from pydantic import BaseModel, Field from starlette.datastructures import State from starlette.routing import Route from vllm import SamplingParams @@ -117,6 +118,30 @@ from vllm_omni.lora.utils import stable_lora_int_id logger = init_logger(__name__) router = APIRouter() +profiler_router = APIRouter() + + +def _should_enable_profiler_endpoints(args: Namespace) -> bool: + # Check upstream vLLM's profiler_config + profiler_config = getattr(args, "profiler_config", None) + if profiler_config is not None: + # profiler_config exists, check if profiler is set + profiler = getattr(profiler_config, "profiler", None) + if profiler is not None: + return True + + # TODO: remove this env after refactoring torch profiler to CLI args + env_value = os.environ.get("VLLM_TORCH_PROFILER_DIR") + return env_value is not None + + +class ProfileRequest(BaseModel): + """Request model for profiling endpoints.""" + + stages: list[int] | None = Field( + default=None, + description="List of stage IDs to profile. If None, profiles all stages.", + ) def _remove_route_from_router( @@ -246,6 +271,11 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, await omni_init_app_state(engine_client, app.state, args) + # Conditionally register profiler endpoints based on config or env var + if _should_enable_profiler_endpoints(args): + logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!") + app.include_router(profiler_router) + vllm_config = await engine_client.get_vllm_config() # Check if pure diffusion mode (vllm_config will be None) @@ -897,8 +927,113 @@ async def list_voices(raw_request: Request): if handler is None: return base(raw_request).create_error_response(message="The model does not support Speech API") + # Get all speakers (both model built-in and uploaded) speakers = sorted(handler.supported_speakers) if handler.supported_speakers else [] - return JSONResponse(content={"voices": speakers}) + + # Get uploaded speakers details + uploaded_speakers = [] + if hasattr(handler, "uploaded_speakers"): + for voice_name, info in handler.uploaded_speakers.items(): + uploaded_speakers.append( + { + "name": info.get("name", voice_name), + "consent": info.get("consent", ""), + "created_at": info.get("created_at", 0), + "file_size": info.get("file_size", 0), + "mime_type": info.get("mime_type", ""), + } + ) + + return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers}) + + +@router.post( + "/v1/audio/voices", + responses={ + HTTPStatus.OK.value: {"model": dict}, + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +async def upload_voice( + raw_request: Request, + audio_sample: UploadFile = File(...), + consent: str = Form(...), + name: str = Form(...), +): + """Upload a new voice sample for voice cloning. + + Uploads an audio file that can be used as a reference for voice cloning + in Base task TTS requests. The voice can then be referenced by name + in subsequent TTS requests. + + Args: + audio_sample: Audio file (max 10MB) + consent: Consent recording ID + name: Name for the new voice + raw_request: Raw FastAPI request + + Returns: + JSON response with voice information + """ + handler = Omnispeech(raw_request) + if handler is None: + return base(raw_request).create_error_response(message="The model does not support Speech API") + + try: + # Upload the voice + result = await handler.upload_voice(audio_sample, consent, name) + + return JSONResponse(content={"success": True, "voice": result}) + + except ValueError as e: + return base(raw_request).create_error_response(message=str(e)) + except Exception as e: + logger.exception(f"Failed to upload voice: {e}") + return base(raw_request).create_error_response(message=f"Failed to upload voice: {str(e)}") + + +@router.delete( + "/v1/audio/voices/{name}", + responses={ + HTTPStatus.OK.value: {"model": dict}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +async def delete_voice(name: str, raw_request: Request): + """Delete an uploaded voice. + + Deletes the voice sample and associated metadata. Also removes any + cached voice clone prompts for this voice. + + Args: + name: Name of the voice to delete + raw_request: Raw FastAPI request + + Returns: + JSON response indicating success or failure + """ + handler = Omnispeech(raw_request) + if handler is None: + return base(raw_request).create_error_response(message="The model does not support Speech API") + + try: + # Delete the voice + success = await handler.delete_voice(name) + if not success: + return JSONResponse( + content={"success": False, "error": f"Voice '{name}' not found"}, + status_code=HTTPStatus.NOT_FOUND.value, + ) + + return JSONResponse(content={"success": True, "message": f"Voice '{name}' deleted successfully"}) + + except ValueError as e: + return base(raw_request).create_error_response(message=str(e)) + except Exception as e: + logger.exception(f"Failed to delete voice '{name}': {e}") + return base(raw_request).create_error_response(message=f"Failed to delete voice: {str(e)}") @router.websocket("/v1/audio/speech/stream") @@ -1951,3 +2086,55 @@ async def download_video(video_id: str) -> FileResponse: raise HTTPException(status_code=404, detail="Generated video file not found on disk") return FileResponse(path=full_path, media_type=job.media_type, filename=job.file_name) + + +@profiler_router.post("/start_profile") +async def start_profile(raw_request: Request, request: ProfileRequest | None = None): + """Start profiling for the engine. + + Args: + request: Optional request body with stages to profile. + - stages: List of stage IDs to profile. If None, profiles all stages. + + Example: + POST /start_profile + {"stages": [0, 1]} # Profile only stages 0 and 1 + """ + try: + stages = request.stages if request else None + logger.info("Starting profiler for stages: %s", stages if stages else "all") + engine_client = raw_request.app.state.engine_client + result = await engine_client.start_profile(stages=stages) + logger.info("Profiler started.") + return JSONResponse(content=result) + except Exception as e: + logger.exception("Failed to start profiler: %s", e) + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to start profiler: {str(e)}" + ) + + +@profiler_router.post("/stop_profile") +async def stop_profile(raw_request: Request, request: ProfileRequest | None = None): + """Stop profiling for the engine. + + Args: + request: Optional request body with stages to stop profiling. + - stages: List of stage IDs to stop profiling. If None, stops all stages. + + Example: + POST /stop_profile + {"stages": [0, 1]} # Stop profiling only stages 0 and 1 + """ + try: + stages = request.stages if request else None + logger.info("Stopping profiler for stages: %s", stages if stages else "all") + engine_client = raw_request.app.state.engine_client + result = await engine_client.stop_profile(stages=stages) + logger.info("Profiler stopped.") + return JSONResponse(content=result) + except Exception as e: + logger.exception("Failed to stop profiler: %s", e) + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to stop profiler: {str(e)}" + )