import streamlit as st
import torch
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image

# Load the model and processor
model_id = "brucewayne0459/paligemma_derm"
processor = AutoProcessor.from_pretrained(model_id)
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, device_map={"": 0})
model.eval()

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# Streamlit app
st.title("Skin Condition Identifier")
st.write("Upload an image and provide a text prompt to identify the skin condition.")

# File uploader for image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

# Text input for prompt
input_text = st.text_input("Enter your prompt:", "Identify the skin condition?")

# Process and display the result when the button is clicked
if uploaded_file is not None and st.button("Analyze"):
    try:
        # Open the uploaded image
        input_image = Image.open(uploaded_file).convert("RGB")
        st.image(input_image, caption="Uploaded Image", use_column_width=True)

        # Prepare inputs
        inputs = processor(
            text=input_text,
            images=input_image,
            return_tensors="pt",
            padding="longest"
        ).to(device)

        # Generate output
        max_new_tokens = 50
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
        
        # Decode output
        decoded_output = processor.decode(outputs[0], skip_special_tokens=True)

        # Display result
        st.success("Analysis Complete!")
        st.write("**Model Output:**", decoded_output)
    except Exception as e:
        st.error(f"Error: {str(e)}")