"""
Minimal ZeroGPU test - Grounding DINO only
"""
import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
import spaces

GDINO_ID = "IDEA-Research/grounding-dino-tiny"

# Load processor only (lightweight)
processor = AutoProcessor.from_pretrained(GDINO_ID)
model = None  # Load inside GPU

@spaces.GPU(duration=15)
def detect(image, text_prompt, box_threshold=0.35, text_threshold=0.25):
    """Detect objects using Grounding DINO on ZeroGPU."""
    global model
    
    # Load model inside GPU context
    if model is None:
        print("Loading Grounding DINO...")
        model = AutoModelForZeroShotObjectDetection.from_pretrained(GDINO_ID)
        print("✓ Loaded")
    
    try:
        # Process image
        print(f"Processing prompt: {text_prompt}")
        pil_image = Image.open(image).convert("RGB")
        
        # Run detection
        inputs = processor(images=pil_image, text=text_prompt, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
        
        results = processor.post_process_grounded_object_detection(
            outputs,
            inputs.input_ids,
            box_threshold=box_threshold,
            text_threshold=text_threshold,
            target_sizes=[pil_image.size[::-1]]
        )[0]
        
        boxes = results["boxes"].cpu().numpy()
        labels = results["labels"]
        scores = results["scores"].cpu().numpy()
        
        detections = []
        for i in range(len(boxes)):
            detections.append({
                "label": labels[i],
                "score": float(scores[i]),
                "box": boxes[i].tolist()
            })
        
        return {
            "num_detections": len(detections),
            "detections": detections
        }
        
    except Exception as e:
        import traceback
        return {"error": f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"}

demo = gr.Interface(
    fn=detect,
    inputs=[
        gr.Image(type="filepath", label="Upload image"),
        gr.Textbox(value="chair . table . sofa", label="Objects to detect (dot-separated)"),
        gr.Slider(0, 1, value=0.35, label="Box threshold"),
        gr.Slider(0, 1, value=0.25, label="Text threshold")
    ],
    outputs=gr.JSON(label="Results"),
    title="🔍 Grounding DINO Test",
    api_name="detect",
    show_error=True
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)