#!/usr/bin/env python3 """ Uncensored chat LLM API — loads an instruct model once, serves an OpenAI-compatible /v1/chat/completions endpoint. Runs on the AMD MI50 (gfx906) via ROCm 5.7 + torch 2.3.1, mirroring the joycaption service pattern. Model fits in fp16 inside the 32GB VRAM (≈12B ceiling). Override with env MODEL_ID. Endpoints: POST /v1/chat/completions — OpenAI-compatible chat (used by gen_poses.py) GET /v1/models — list the loaded model GET /health — health check + GPU info Env: MODEL_ID HuggingFace repo id (default below) PORT listen port (default 8001) HSA_OVERRIDE_GFX_VERSION=9.0.6 set by start script for gfx906 """ import asyncio import os import subprocess import time import uuid from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager import torch from fastapi import FastAPI, HTTPException from pydantic import BaseModel from transformers import AutoModelForCausalLM, AutoTokenizer torch.set_num_threads(1) torch.set_num_interop_threads(1) gpu_executor = ThreadPoolExecutor(max_workers=1) # Mistral [INST] template — fallback for models without chat_template in tokenizer_config _MISTRAL_TEMPLATE = ( "{{ bos_token }}" "{% for message in messages %}" "{% if message['role'] == 'system' %}" "{{ '[INST] ' + message['content'] + '\n\n' }}" "{% elif message['role'] == 'user' %}" "{% if not loop.first or messages[0]['role'] != 'system' %}{{ '[INST] ' }}{% endif %}" "{{ message['content'] + ' [/INST]' }}" "{% elif message['role'] == 'assistant' %}" "{{ ' ' + message['content'] + eos_token }}" "{% endif %}" "{% endfor %}" ) MODEL_ID = os.environ.get("MODEL_ID", "dphn/Dolphin3.0-Mistral-24B") # For 24B fp16 (~48GB): split across 32GB VRAM + CPU RAM via device_map="auto". # Override with MAX_GPU_MEM / MAX_CPU_MEM env vars if needed. MAX_GPU_MEM = os.environ.get("MAX_GPU_MEM", "27GiB") # 32GiB total - 4GiB sys - 1GiB headroom MAX_CPU_MEM = os.environ.get("MAX_CPU_MEM", "100GiB") # 113GB available state: dict = {} @asynccontextmanager async def lifespan(app: FastAPI): # Boost GPU to high performance mode (avoids power-saving clock throttle) subprocess.run(["/opt/rocm/bin/rocm-smi", "--setperflevel", "high"], capture_output=True) print(f"Loading model {MODEL_ID} (gpu≤{MAX_GPU_MEM} cpu≤{MAX_CPU_MEM})...", flush=True) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) if not tokenizer.chat_template: print("No chat_template found — applying Mistral [INST] fallback.", flush=True) tokenizer.chat_template = _MISTRAL_TEMPLATE model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16, device_map="auto", max_memory={0: MAX_GPU_MEM, "cpu": MAX_CPU_MEM}, ) model.eval() state["tokenizer"] = tokenizer state["model"] = model # Warm up: tiny generation to trigger kernel compilation print("Warming up...", flush=True) _run_chat([{"role": "user", "content": "hi"}], max_tokens=4, temperature=0.0, top_p=1.0, stop=None) print(f"Model ready on: {next(model.parameters()).device}", flush=True) yield gpu_executor.shutdown(wait=False, cancel_futures=True) state.clear() app = FastAPI(title="Uncensored Chat LLM API", lifespan=lifespan) def _run_chat(messages, max_tokens, temperature, top_p, stop): tokenizer = state["tokenizer"] model = state["model"] input_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ).to(model.device) prompt_tokens = input_ids.shape[1] do_sample = temperature is not None and temperature > 0 gen_kwargs = dict( max_new_tokens=max_tokens, do_sample=do_sample, use_cache=True, pad_token_id=tokenizer.eos_token_id, ) if do_sample: gen_kwargs["temperature"] = temperature gen_kwargs["top_p"] = top_p t0 = time.perf_counter() with torch.inference_mode(): output_ids = model.generate(input_ids, **gen_kwargs) dt = time.perf_counter() - t0 new_ids = output_ids[0][prompt_tokens:] text = tokenizer.decode(new_ids, skip_special_tokens=True) finish = "stop" if stop: stops = [stop] if isinstance(stop, str) else list(stop) cut = min((text.find(s) for s in stops if s and s in text), default=-1) if cut != -1: text = text[:cut] completion_tokens = int(new_ids.shape[0]) if completion_tokens >= max_tokens: finish = "length" del input_ids, output_ids return { "text": text.strip(), "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "finish_reason": finish, "generate_s": round(dt, 2), } async def run_chat(*args): return await asyncio.get_running_loop().run_in_executor(gpu_executor, _run_chat, *args) class ChatMessage(BaseModel): role: str content: str class ChatRequest(BaseModel): model: str | None = None messages: list[ChatMessage] max_tokens: int = 512 temperature: float = 0.8 top_p: float = 0.95 stop: list[str] | str | None = None @app.post("/v1/chat/completions") async def chat_completions(req: ChatRequest): if not req.messages: raise HTTPException(400, "messages must not be empty") msgs = [{"role": m.role, "content": m.content} for m in req.messages] r = await run_chat(msgs, req.max_tokens, req.temperature, req.top_p, req.stop) return { "id": f"chatcmpl-{uuid.uuid4().hex[:24]}", "object": "chat.completion", "created": int(time.time()), "model": MODEL_ID, "choices": [ { "index": 0, "message": {"role": "assistant", "content": r["text"]}, "finish_reason": r["finish_reason"], } ], "usage": { "prompt_tokens": r["prompt_tokens"], "completion_tokens": r["completion_tokens"], "total_tokens": r["prompt_tokens"] + r["completion_tokens"], }, "timing": {"generate_s": r["generate_s"]}, } @app.get("/v1/models") async def list_models(): return {"object": "list", "data": [{"id": MODEL_ID, "object": "model"}]} @app.get("/health") async def health(): gpu = {} if torch.cuda.is_available(): gpu = { "name": torch.cuda.get_device_name(0), "vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1), "vram_used_gb": round(torch.cuda.memory_allocated(0) / 1e9, 1), } return {"status": "ok", "model": MODEL_ID, "loaded": "model" in state, "gpu": gpu} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8001)))