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

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

# Load model
pipe = pipeline(
    task="image-to-text",
    model="bczhou/tiny-llava-v1-hf",
    trust_remote_code=True,
    device_map="cpu"
)

# 👇 Clear info so the widget tree is stable
st.markdown("Use your **webcam** OR upload a file:")

# Try webcam first
image = st.camera_input("📷 Take a snapshot")

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

# Prompt
prompt = st.text_input("💬 Your question:", value="Describe this scene.")

# Run TinyLLaVA
if image is not None 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("Generating..."):
        result = pipe(query, img)
        answer = result[0]["generated_text"]

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