import os

# ✅ Set writable cache dir in /tmp
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
os.environ["HF_HOME"] = "/tmp/hf_cache"

import streamlit as st
from diffusers import DDPMPipeline
import torch

# Streamlit page config
st.set_page_config(page_title="DDPM - Score-Based Model", layout="centered")
st.title("🌀 DDPM-Based Image Generation")

st.markdown("""
This demo uses a DDPM (Denoising Diffusion Probabilistic Model) pipeline to generate images
from noise using Hugging Face 🤗 `diffusers` library. You are using the pretrained `google/ddpm-cifar10-32` model.
""")

@st.cache_resource
def load_ddpm_pipeline():
    model_id = "google/ddpm-cifar10-32"
    pipe = DDPMPipeline.from_pretrained(model_id)
    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    return pipe

def generate_ddpm_images(pipe, num_images=1):
    images = []
    for _ in range(num_images):
        image = pipe().images[0]
        images.append(image)
    return images

num_images = st.slider("Select number of images to generate:", 1, 4, 1)
generate_button = st.button("🚀 Generate Images")

if generate_button:
    with st.spinner("Generating images using DDPM model..."):
        pipe = load_ddpm_pipeline()
        images = generate_ddpm_images(pipe, num_images)
    st.success("Images generated using DDPM!")
    for idx, img in enumerate(images):
        st.image(img, caption=f"Generated Image {idx + 1}", use_column_width=True)