app.py

import os
import io
import gradio as gr
import torch
from huggingface_hub import login
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
from qwen_vl_utils import process_vision_info
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse

login(token=os.environ["HF_TOKEN"])

model_id = "mendicant04/DermoGPT-RL"

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

print("Loading model...")
model = Qwen3VLForConditionalGeneration.from_pretrained(
    model_id,
    quantization_config=quant_config,
    device_map="auto"
)
processor = AutoProcessor.from_pretrained(model_id)
print("Model loaded!")

def analyze(image: Image.Image, prompt: str = "What's happening with my skin?"):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt}
            ],
        }
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, _ = process_vision_info(messages)
    inputs = processor(
        text=[text], images=image_inputs, return_tensors="pt"
    ).to("cuda")

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs, max_new_tokens=256, do_sample=True, temperature=0.2
        )
    output_ids = generated_ids[0, inputs.input_ids.shape[1]:]
    return processor.decode(output_ids, skip_special_tokens=True)

# ── FastAPI ───────────────────────────────────────────────────────────────────
fapi = FastAPI()

@fapi.get("/health")
def health():
    return {"status": "ok"}

@fapi.post("/analyze")
async def analyze_endpoint(
    image: UploadFile = File(...),
    prompt: str = Form(default="What's happening with my skin?")
):
    image_bytes = await image.read()
    pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    result = analyze(pil_image, prompt)
    return JSONResponse({"analysis": result})

# ── Gradio mounted ONCE onto FastAPI ─────────────────────────────────────────
demo = gr.Interface(
    fn=analyze,
    inputs=[
        gr.Image(type="pil", label="Skin Image"),
        gr.Textbox(value="What's happening with my skin?", label="Prompt")
    ],
    outputs=gr.Textbox(label="Analysis"),
    title="DermoGPT-RL",
    description="Upload a skin image to get an AI-powered dermatological analysis."
)

gr.mount_gradio_app(fapi, demo, path="/ui")