server.py

import asyncio
import threading
import io
import queue
import json
import argparse
from contextlib import asynccontextmanager

import mlx.core as mx
import numpy as np
import soundfile as sf
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from mlx_vlm import load
from mlx_vlm.models.qwen3_omni_moe.omni_utils import prepare_omni_inputs

MODEL_PATH = None
SAMPLE_RATE = None

model = None
processor = None
abort_flag = threading.Event()
model_lock = threading.Lock()
chunk_counter = 0
request_count = 0
request_count_lock = threading.Lock()

_DEBUG = False

def log(tag: str, msg: str):
    if _DEBUG:
        print(f"[{tag}] {msg}")


class ChatMessage(BaseModel):
    role: str
    content: list


class ChatRequest(BaseModel):
    messages: list[ChatMessage]
    speaker: str = "Ethan"
    thinker_max_new_tokens: int = 2048
    thinker_temperature: float = 0.7
    thinker_top_p: float = 0.9
    thinker_repetition_penalty: float = 1.1
    thinker_repetition_context_size: int = 64
    talker_max_new_tokens: int = 2048
    talker_temperature: float = 0.9
    talker_top_p: float = 1.0
    chunk_size: int = 50
    initial_chunk_size: int = 25


@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, processor
    model, processor = load(MODEL_PATH, trust_remote_code=True)
    yield


app = FastAPI(lifespan=lifespan)


def stream_generate(conversation: list, gen_config: dict, output_queue: queue.Queue):
    global request_count, chunk_counter
    
    with request_count_lock:
        request_count += 1
        log("SERVER", f"Request started, count={request_count}")
    
    try:
        log("SERVER", "Waiting for model lock...")
        with model_lock:
            log("SERVER", "Model lock acquired")
            chunk_counter = 0
            
            log("SERVER", "Preparing model inputs...")
            model_inputs, _ = prepare_omni_inputs(processor, conversation)
            log("SERVER", f"Model inputs ready, input_ids shape: {model_inputs['input_ids'].shape}")
            
            if abort_flag.is_set():
                log("SERVER", "Aborted before generation")
                output_queue.put(None)
                return
            
            generate_kwargs = {
                "input_ids": model_inputs["input_ids"],
                "pixel_values": model_inputs.get("pixel_values"),
                "pixel_values_videos": model_inputs.get("pixel_values_videos"),
                "image_grid_thw": model_inputs.get("image_grid_thw"),
                "video_grid_thw": model_inputs.get("video_grid_thw"),
                "input_features": model_inputs.get("input_features"),
                "feature_attention_mask": model_inputs.get("feature_attention_mask"),
                "audio_feature_lengths": model_inputs.get("audio_feature_lengths"),
                **gen_config,
            }
            generate_kwargs = {k: v for k, v in generate_kwargs.items() if v is not None}
            
            log("SERVER", "Starting generate_stream_interleaved...")
            for event_type, data in model.generate_stream_interleaved(**generate_kwargs):
                if abort_flag.is_set():
                    log("SERVER", "Aborted during generation")
                    break
                if event_type == "text":
                    text = processor.decode(data)
                    log("SERVER", f"Text: {text[-50:]}")
                    output_queue.put(("text", text))
                elif event_type == "audio":
                    chunk_counter += 1
                    audio_data = np.array(data.reshape(-1))
                    log("SERVER", f"Audio #{chunk_counter}: {len(audio_data)} samples")
                    output_queue.put(("audio", audio_data))
            
            output_queue.put(None)
            log("SERVER", f"Generation complete, {chunk_counter} audio chunks")
    except Exception as e:
        log("SERVER", f"Error: {e}")
        import traceback
        traceback.print_exc()
        output_queue.put(None)
    finally:
        with request_count_lock:
            request_count -= 1
            log("SERVER", f"Request ended, count={request_count}")


async def stream_generator(conversation: list, gen_config: dict):
    global abort_flag
    abort_flag.clear()
    
    output_queue = queue.Queue()
    loop = asyncio.get_event_loop()
    audio_counter = 0
    
    gen_thread = threading.Thread(target=stream_generate, args=(conversation, gen_config, output_queue), daemon=True)
    gen_thread.start()
    
    while True:
        if abort_flag.is_set():
            break
        try:
            item = await loop.run_in_executor(None, output_queue.get, True, 0.1)
        except queue.Empty:
            if not gen_thread.is_alive() and output_queue.empty():
                break
            continue
        
        if item is None:
            break
        
        event_type, data = item
        if event_type == "text":
            text_json = json.dumps({"type": "text", "content": data}).encode("utf-8")
            yield b"__TEXT__" + text_json + b"__END_TEXT__"
        elif event_type == "audio":
            audio_counter += 1
            buf = io.BytesIO()
            sf.write(buf, data, SAMPLE_RATE, format="WAV")
            yield buf.getvalue()
    
    gen_thread.join(timeout=1.0)
    log("SERVER", f"Stream ended, {audio_counter} audio chunks")


@app.post("/chat")
async def chat(request: ChatRequest):
    log("SERVER", "========== NEW CHAT REQUEST ==========")
    log("SERVER", f"{len(request.messages)} messages")
    if _DEBUG:
        for i, m in enumerate(request.messages):
            print(f"[SERVER] [{i}] {m.role}: {[c.get('type') for c in m.content]}")
    conversation = [{"role": m.role, "content": m.content} for m in request.messages]
    gen_config = {
        "speaker": request.speaker,
        "thinker_max_new_tokens": request.thinker_max_new_tokens,
        "thinker_temperature": request.thinker_temperature,
        "thinker_top_p": request.thinker_top_p,
        "thinker_repetition_penalty": request.thinker_repetition_penalty,
        "thinker_repetition_context_size": request.thinker_repetition_context_size,
        "talker_max_new_tokens": request.talker_max_new_tokens,
        "talker_temperature": request.talker_temperature,
        "talker_top_p": request.talker_top_p,
        "chunk_size": request.chunk_size,
        "initial_chunk_size": request.initial_chunk_size,
    }
    return StreamingResponse(stream_generator(conversation, gen_config), media_type="audio/wav")


@app.post("/abort")
async def abort():
    global abort_flag
    abort_flag.set()
    log("SERVER", "Abort requested")
    return {"status": "abort_requested", "server_code": 1 if request_count > 0 else 0}


@app.get("/status")
async def status():
    return {"server_code": 1 if request_count > 0 else 0}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, required=True)
    parser.add_argument("--sample-rate", type=int, default=24000)
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()
    
    MODEL_PATH = args.model_path
    SAMPLE_RATE = args.sample_rate
    _DEBUG = args.debug
    
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)