import os
import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM

# 1. Determine the available hardware device
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

print(f"Loading model on device: {device}")

# 2. Initialize Florence-2 model and processor
# Using trust_remote_code=True as required by the Florence architecture
model_id = 'microsoft/Florence-2-base'
florence_model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    trust_remote_code=True,
    torch_dtype=torch_dtype
).to(device).eval()

florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

def generate_caption(image):
    if image is None:
        return "Please upload an image first."
        
    # Ensure the input is a valid PIL Image
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
        
    # Execute the detailed captioning task
    prompt_task = "<MORE_DETAILED_CAPTION>"
    inputs = florence_processor(text=prompt_task, images=image, return_tensors="pt").to(device)
    
    # Match data type to the model precision
    if torch.cuda.is_available():
        inputs = {k: v.to(torch.float16) if k == "pixel_values" else v for k, v in inputs.items()}

    with torch.no_grad():
        generated_ids = florence_model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )
        
    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
    parsed_answer = florence_processor.post_process_generation(
        generated_text, 
        task=prompt_task, 
        image_size=(image.width, image.height)
    )
    
    base_caption = parsed_answer[prompt_task]
    
    # Appends styling parameters for prompt generation formatting
    final_prompt = f"{base_caption}. Detailed features, high quality, highly descriptive cinematic lighting."
    return final_prompt

# 3. Construct the Gradio Interface UI
with gr.Blocks() as demo:
    gr.Markdown("# Detailed Image Captioning Tool")
    gr.Markdown("Upload an image to generate highly descriptive text prompts using Florence-2.")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Upload Image", type="pil")
            run_btn = gr.Button("Generate Detailed Prompt", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Generated Output", lines=5, show_copy_button=True)
            
    run_btn.click(fn=generate_caption, inputs=input_img, outputs=output_text)

# Launch the app container
if __name__ == "__main__":
    demo.launch()