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

# Load model & processor
# model_id = "llava-hf/llava-1.5-7b-hf" # Exceeding 16 GB Memory
model_id = "llava-hf/llava-1.5-7b-hf-int4"

processor = AutoProcessor.from_pretrained(model_id)
model = LlavaForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    low_cpu_mem_usage=True,
    device_map="auto",
)


# Main prediction function
def analyze_palm(image, question, history):
    if image is None or not question.strip():
        history.append((question, "Please provide both image and question."))
        return history, ""

    conversation = [
        {"role": "user", "content": [
            {"type": "text", "text": question},
            {"type": "image"}
        ]}
    ]
    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)

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

    history.append((question, response))
    return history, ""


# Build UI using Blocks
with gr.Blocks() as demo:
    gr.Markdown("## 🖐️ AI Palm Reader\nUpload a palm image and ask a question. Get a palmistry-style response.")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Palm Image")
            prompt_input = gr.Textbox(lines=2, label="Your Question", placeholder="What does my palm say?")
            submit_btn = gr.Button("Ask")

        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Palmistry Chat")

    state = gr.State([])

    submit_btn.click(
        fn=analyze_palm,
        inputs=[image_input, prompt_input, state],
        outputs=[chatbot, prompt_input]
    )

demo.launch()