from fastapi import FastAPI, Query
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
import requests
from PIL import Image
from io import BytesIO

app = FastAPI()

# Load model and processor
checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"

# Check for Metal GPU support on macOS
device = "mps" if torch.backends.mps.is_available() else "cpu"

processor = AutoProcessor.from_pretrained(
    checkpoint,
    min_pixels=256 * 28 * 28,
    max_pixels=1280 * 28 * 28
)

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    checkpoint,
    torch_dtype=torch.float16 if device == "mps" else torch.bfloat16,  # Use float16 on Apple Metal
    device_map={"": 0} if device == "mps" else "cpu",
    attn_implementation="flash_attention_2",  # If it supports Mac
)


# Function to load and resize images (reduces processing time)
def load_and_resize_image(image_url):
    response = requests.get(image_url)
    image = Image.open(BytesIO(response.content)).convert("RGB")
    image = image.resize((512, 512))  # Resize to 512x512 to speed up processing
    return image


@app.get("/")
def read_root():
    return {"message": "API is live. Use the /predict endpoint."}


@app.get("/predict")
def predict(image_url: str = Query(...), prompt: str = Query(...)):
    messages = [
        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
    ]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    # Process image
    image_inputs = [load_and_resize_image(image_url)]
    video_inputs = None

    # Process inputs
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        truncation=True,  # Ensures token limit
        max_length=512,  # Prevents excessive memory usage
        return_tensors="pt",
    ).to(device)

    # Generate response
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=64)  # Reduced for faster inference

    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
    output_texts = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    return {"response": output_texts[0]}