# import gradio as gr
#
# def greet(name):
#     return "Hello " + name + "!!"
#
# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
# demo.launch()

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

# Load LLaVA model
model_id = "llava-hf/llava-1.5-7b-hf"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id,
                                             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                                             device_map="auto")


def chat_with_llava(image, question, history=[]):
    if image is None or not question.strip():
        return history + [[question, "Please provide both an image and a prompt."]]

    # Preprocess input
    inputs = processor(text=question, images=image, return_tensors="pt").to(model.device)

    # Generate output
    output = model.generate(**inputs, max_new_tokens=512)
    answer = processor.decode(output[0], skip_special_tokens=True)

    # Append to chat history
    history.append([question, answer])
    return history


# Create Gradio chatbot interface
chat_interface = gr.ChatInterface(
    fn=chat_with_llava,
    inputs=[gr.Image(type="pil", label="Palm Image"),
            gr.Textbox(label="Your Question", placeholder="What does my palm say about my future?")],
    title="🖐️ AI Palm Reader",
    description="Upload a palm image and ask a question. The LLaVA model will respond like a palmistry expert!",
)

chat_interface.launch()