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

# Model configuration
MODEL_ID = "FreedomIntelligence/HuatuoGPT-Vision-7B"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load model and processor
@gr.cache_resource
def load_model():
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        device_map="auto" if DEVICE == "cuda" else None,
    )
    if DEVICE == "cpu":
        model = model.to(DEVICE)
    return processor, model

def process_image_and_text(image, text_prompt, temperature=0.7, max_new_tokens=512):
    """
    Process an image and text prompt with HuatuoGPT-Vision-7B model.

    Args:
        image: PIL Image or path to image
        text_prompt: Text prompt/question about the image
        temperature: Sampling temperature
        max_new_tokens: Maximum tokens to generate

    Returns:
        Generated text response
    """
    if image is None:
        return "Please provide an image."

    if not text_prompt or text_prompt.strip() == "":
        return "Please provide a text prompt or question."

    try:
        processor, model = load_model()

        # Convert to PIL Image if needed
        if isinstance(image, str):
            image = Image.open(image).convert("RGB")
        elif not isinstance(image, Image.Image):
            image = Image.fromarray(image).convert("RGB")

        # Prepare inputs
        inputs = processor(
            text=text_prompt,
            images=image,
            return_tensors="pt"
        )

        # Move inputs to device
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=temperature > 0,
                top_p=0.9,
            )

        # Decode response
        response = processor.decode(outputs[0], skip_special_tokens=True)

        # Clean up response - remove the input prompt if it's in the output
        if text_prompt in response:
            response = response.split(text_prompt)[-1].strip()

        return response

    except Exception as e:
        return f"Error processing request: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="HuatuoGPT-Vision-7B Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏥 HuatuoGPT-Vision-7B Demo

    A powerful vision-language model specialized for medical and scientific image understanding.

    **Model Details:**
    - **Architecture:** LLaVA-Qwen2 (7.9B parameters)
    - **Training Data:** PubMedVision dataset (medical/scientific images)
    - **Capabilities:** Image understanding, medical analysis, visual question answering
    - **License:** Apache 2.0

    ---
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input")
            image_input = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard", "webcam"]
            )

            text_input = gr.Textbox(
                label="Question/Prompt",
                placeholder="Ask a question about the image, e.g., 'What is shown in this medical scan?'",
                lines=3
            )

            with gr.Row():
                temperature = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.7,
                    step=0.05,
                    info="Higher = more creative, Lower = more deterministic"
                )
                max_tokens = gr.Slider(
                    label="Max Tokens",
                    minimum=64,
                    maximum=1024,
                    value=512,
                    step=64,
                    info="Maximum length of response"
                )

            submit_btn = gr.Button("🚀 Generate Response", variant="primary", size="lg")
            clear_btn = gr.Button("🔄 Clear", size="sm")

        with gr.Column(scale=1):
            gr.Markdown("### Output")
            output_text = gr.Textbox(
                label="Model Response",
                lines=10,
                interactive=False,
                show_copy_button=True
            )

    # Examples
    gr.Markdown("### 📋 Example Prompts")
    gr.Examples(
        examples=[
            ["What does this image show?", None],
            ["Describe the main features visible in this image.", None],
            ["What medical condition might this image indicate?", None],
            ["Analyze the anatomical structures in this image.", None],
        ],
        inputs=[text_input, image_input],
        label="Try these prompts (upload your own image first)",
        cache_examples=False
    )

    # Event handlers
    submit_btn.click(
        fn=process_image_and_text,
        inputs=[image_input, text_input, temperature, max_tokens],
        outputs=output_text
    )

    def clear_all():
        return None, "", "", ""

    clear_btn.click(
        fn=clear_all,
        outputs=[image_input, text_input, output_text]
    )

    # Info section
    gr.Markdown("""
    ---

    ### ℹ️ About HuatuoGPT-Vision-7B

    HuatuoGPT-Vision is a specialized vision-language model built on top of LLaVA and Qwen-2.
    It's trained on PubMedVision, a large-scale dataset of biomedical images and text pairs.

    **Best used for:**
    - Medical image analysis
    - Scientific image understanding
    - Visual question answering on medical topics
    - Biomedical document analysis

    **Model Card:** [FreedomIntelligence/HuatuoGPT-Vision-7B](https://huggingface.co/FreedomIntelligence/HuatuoGPT-Vision-7B)

    **Paper:** [arxiv:2406.19280](https://arxiv.org/abs/2406.19280)
    """)

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