app.py

from spaces import GPU
GPU()  # ✅ Required for ZeroGPU activation

import gradio as gr
import torch
from PIL import Image

model = None
processor = None

def load_model():
    global model, processor
    if model is None:
        if not torch.cuda.is_available():
            return "GPU not ready — please retry shortly."

        # 🔁 Lazy import to avoid cold start crash
        import unsloth
        from unsloth import FastVisionModel
        from transformers import AutoProcessor

        model_id = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit"
        model, processor = FastVisionModel.from_pretrained(
            model_id,
            load_in_4bit=True,
            device_map="auto",
            torch_dtype=torch.float16
        )

def chat(image, prompt):
    if not prompt:
        return "Please enter a prompt."
    if image is None:
        return "Please upload an image."

    if not torch.cuda.is_available():
        return "⏳ GPU not yet active. Please wait 10–30 seconds and retry."

    load_model()
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=100)
    response = processor.decode(outputs[0], skip_special_tokens=True)
    return response

# ✅ No `app = Interface()` (that causes ASGI errors in Spaces)
gr.Interface(
    fn=chat,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(lines=2, label="Prompt")
    ],
    outputs="text",
    title="Unsloth LLaMA 3.2 Vision Chatbot (ZeroGPU)",
    description="Lazy-loading Vision + Text model with Unsloth on ZeroGPU"
).launch()