import torch
import gradio as gr
from PIL import Image
from transformers import Blip2ForConditionalGeneration, AutoProcessor
from peft import PeftModel, PeftConfig

# Your fine-tuned model on Hugging Face
peft_model_id = "MargiPandya/blip2-bdd100k"

# Load PEFT configuration
config = PeftConfig.from_pretrained(peft_model_id)

# Load the base BLIP-2 model (use full precision if 8-bit not supported)
model = Blip2ForConditionalGeneration.from_pretrained(
    config.base_model_name_or_path,
    torch_dtype=torch.float32,
    device_map=None,
)

# Load the fine-tuned adapter
model = PeftModel.from_pretrained(model, peft_model_id)

# Load processor
# processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
processor = AutoProcessor.from_pretrained(config.base_model_name_or_path)

# Move to CPU (for Spaces with no GPU)
device = "cpu"
model.to(device)
model.eval()

# Image captioning function
def caption_image(image):
    if image is None:
        return "Please upload an image."
    # Ensure image is processed and input tensors are on same device as the model
    inputs = processor(images=image, text = "", return_tensors="pt")
    # Move all tensors to correct device
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        # Safe parameter recommendations for BLIP-2
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=False,
            # temperature=0.9,
            # top_p=0.9,
            # repetition_penalty=1.5,
        )
    caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return caption.strip()

# Gradio Interface
app = gr.Interface(
    fn=caption_image,
    inputs=gr.Image(type="pil", label="Upload a driving scene image"),
    outputs=gr.Textbox(label="Generated Caption"),
    title="BLIP-2 BDD100K Scene Captioner",
    description="Upload a driving scene image and get an automatically generated caption.",
    theme="soft",
)

if __name__ == "__main__":
    app.launch()