# Copyright 2024 Ronan Le Meillat
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Import necessary libraries
import gradio as gr
from transformers import AutoProcessor, Idefics3ForConditionalGeneration, image_utils
import torch

# Determine the device (GPU or CPU) to run the model on
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(f"Using device: {device}")  # Log the device being used

# Define the model ID and base model path
model_id = "eltorio/IDEFICS3_ROCO"
base_model_path = "HuggingFaceM4/Idefics3-8B-Llama3"  # or change to local path

# Initialize the processor from the base model path
processor = AutoProcessor.from_pretrained(base_model_path, trust_remote_code=True)

# Initialize the model from the base model path and set the torch dtype to bfloat16
model = Idefics3ForConditionalGeneration.from_pretrained(
    base_model_path, torch_dtype=torch.bfloat16
).to(device)  # Move the model to the specified device

# Load the adapter from the model ID and automatically map it to the device
model.load_adapter(model_id, device_map="auto")

# Define a function to infer a description from an image
def infere(image):
    """
    Generate a description of a medical image.

    Args:
    - image (PIL Image): The medical image to describe.

    Returns:
    - generated_texts (List[str]): A list containing the generated description.
    """

    # Define a chat template for the model to respond to
    messages = [
        {
            "role": "system",
            "content": [
                {"type": "text", "text": "You are a valuable medical doctor and you are looking at an image of your patient."},
            ]
        },
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "What do we see in this image?"},
            ]
        },
    ]

    # Apply the chat template and add a generation prompt
    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)

    # Preprocess the input image and text
    inputs = processor(text=prompt, images=[image], return_tensors="pt")

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

    # Generate a description with the model
    generated_ids = model.generate(**inputs, max_new_tokens=100)

    # Decode the generated IDs into text
    generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)

    return generated_texts

# Define the title, description, and device description for the Gradio interface
title = f"<a href='https://huggingface.co/eltorio/IDEFICS3_ROCO'>IDEFICS3_ROCO</a>: AI X-ray Image Recognition <span style='color:red'><b>running on {device}</b></span>"
desc = "This model generates a description of a medical image."

device_desc = f"This model is running on {device} 🚀." if device == torch.device('cuda') else f"🐢 This model is running on {device} it will be very (very) slow. If you can donate some GPU time it will be usable 🐢. <a href='https://huggingface.co/eltorio/IDEFICS3_ROCO/discussions'>Please contact us.</a>"

# Define the long description for the Gradio interface
long_desc = f"This demo is based on the <a href='https://huggingface.co/eltorio/IDEFICS3_ROCO'>IDEFICS3_ROCO model</a>, which is a multimodal model that can generate text from images. It has been fine-tuned on <a href='https://huggingface.co/datasets/eltorio/ROCO-radiology'>eltorio/ROCO-radiology</a>&nbsp;a dataset of medical images and can generate descriptions of medical images. Try uploading an image of a medical image and see what the model generates!<br><b>{device_desc}</b><br> 2024 - Ronan Le Meillat"

# Create a Gradio interface with the infere function and specified title and descriptions
radiotest = gr.Interface(fn=infere, inputs="image", outputs="text", title=title, 
                description=desc, article=long_desc)

# Launch the Gradio interface and share it
radiotest.launch(share=True)