from fastapi import FastAPI
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
import torch

app = FastAPI()

model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto", device_map="auto")
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")

@app.post("/predict")
async def predict(messages: list):
    # Processamento e inferência
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], padding=True, return_tensors="pt").to("cuda")  # Altere para "cpu" se não tiver GPU
    generated_ids = model.generate(**inputs, max_new_tokens=128)
    output_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
    return {"response": output_text}