import streamlit as st
from transformers import pipeline
from PIL import Image

# Page config
st.set_page_config(page_title="TinyLLaVA Snapshot", layout="centered")
st.title("📸 TinyLLaVA — Snapshot Q&A (Hugging Face Spaces Safe)")

# Load TinyLLaVA pipeline
st.info("Loading TinyLLaVA... Please wait ⏳")
pipe = pipeline(
    task="image-to-text",
    model="bczhou/tiny-llava-v1-hf",
    trust_remote_code=True,
    device_map="cpu"
)
st.success("✅ Model loaded!")

# Widgets
st.markdown("**1️⃣ Take a snapshot OR upload an image**")
image = st.camera_input("📷 Click 'Take Photo' below:")

if not image:
    image = st.file_uploader("📂 Or upload an image file", type=["jpg", "png", "jpeg"])

st.markdown("**2️⃣ Enter your question:**")
prompt = st.text_input("💬", value="Describe this scene.")

# Debug info
st.write(f"📸 Image uploaded: {'Yes' if image else 'No'}")
st.write(f"💬 Prompt: {prompt}")

# If both present
if image and prompt:
    img = Image.open(image).convert("RGB")
    st.image(img, caption="Your Image", use_column_width=True)

    query = f"USER: <image>\n{prompt}\nASSISTANT:"

    with st.spinner("TinyLLaVA is generating..."):
        result = pipe(query, img)
        answer = result[0]["generated_text"]

    st.subheader("📝 Answer")
    st.write(answer)