import gradio as gr
import subprocess
import torch
import os
import base64
import io
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM
from daggr import FnNode, Graph

try:
    subprocess.run('pip install flash-attn --no-build-isolation', 
                   env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, 
                   check=True, shell=True)
except Exception as e:
    print(f"Flash-attn not installed: {e}")

device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading Florence-2 models on {device}...")
try:
    # Base Model
    model_base = AutoModelForCausalLM.from_pretrained(
        'microsoft/Florence-2-base', trust_remote_code=True
    ).to(device).eval()
    proc_base = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)

    # Large Model
    model_large = AutoModelForCausalLM.from_pretrained(
        'microsoft/Florence-2-large', trust_remote_code=True
    ).to(device).eval()
    proc_large = AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True)
    print("✅ Models loaded.")
except Exception as e:
    print(f"❌ Error loading models: {e}")

def load_any_image(img_input):
    """
    Detects if the input is a file path, a Base64 string, or a PIL object.
    """
    if isinstance(img_input, Image.Image):
        return img_input.convert("RGB")
    
    if isinstance(img_input, str):
        # Check if it is a Base64 Data URI
        if img_input.startswith("data:image"):
            base64_data = img_input.split(",")[1]
            img_bytes = base64.b64decode(base64_data)
            return Image.open(io.BytesIO(img_bytes)).convert("RGB")
        # Otherwise treat as a standard file path
        return Image.open(img_input).convert("RGB")
    
    # Fallback for numpy arrays
    return Image.fromarray(img_input).convert("RGB")

def describe_image(uploaded_image, model_choice):
    if uploaded_image is None:
        return "Please upload an image."

    try:
        # Fix the "File name too long" error by handling string inputs correctly
        image = load_any_image(uploaded_image)
        
        # Select Model
        if model_choice == "Florence-2-large":
            model, processor = model_large, proc_large
        else:
            model, processor = model_base, proc_base

        prompt = "<MORE_DETAILED_CAPTION>"
        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
        
        with torch.no_grad():
            generated_ids = model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=1024,
                num_beams=3,
                do_sample=False
            )

        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
        parsed_answer = processor.post_process_generation(
            generated_text, 
            task=prompt, 
            image_size=(image.width, image.height)
        )

        return parsed_answer[prompt]
    
    except Exception as e:
        return f"Error processing image: {str(e)}"

caption_node = FnNode(
    fn=describe_image,
    inputs={
        "uploaded_image": gr.Image(
            label="Upload Image", 
            type="filepath" 
        ),
        "model_choice": gr.Radio(
            choices=["Florence-2-base", "Florence-2-large"], 
            value="Florence-2-large", 
            label="Model Version"
        ),
    },
    outputs={
        "caption": gr.Textbox(label="Generated Detailed Caption", lines=6, interactive=True),
    },
)

graph = Graph(
    name="Florence-2 Image Captioning",
    nodes=[caption_node]
)

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