import streamlit as st
from PIL import Image
from diffusers import DDPMPipeline # مدل سبک‌تر (تنها diffusers استفاده می‌شود)
import torch
import io

# تنظیمات برنامه Streamlit
st.set_page_config(
    page_title="AI Photo Enhancer (CPU)",
    page_icon="✨",
    layout="centered"
)
st.title("✨ ارتقاء و ترمیم چهره (مدل سبک CPU)")
st.markdown("---")

# تنظیم مدل (کش کردن برای سرعت)
@st.cache_resource
def load_pipeline():
    # مدل سبک مخصوص چهره و CPU
    pipe = DDPMPipeline.from_pretrained(
        "google/ddpm-ema-celebahq-256", 
    )
    return pipe

pipe = load_pipeline()

# تابع اصلی تولید چهره جدید
def generate_face():
    # اجرای Generation
    result_image = pipe(
        batch_size=1
    ).images[0]
    return result_image

# رابط کاربری Streamlit
st.markdown("این مدل برای کار روی CPU بهینه شده است و **فقط** می‌تواند تصاویر جدیدی از چهره تولید کند، نه ترمیم عکس ورودی.")

if st.button("تولید چهره جدید (رایگان)"):
    with st.spinner("در حال اجرای هوش مصنوعی و تولید تصویر... (صبر کنید)"):
        try:
            enhanced_image = generate_face() 
            st.success("✅ عملیات با موفقیت انجام شد!")
            st.image(enhanced_image, caption="چهره تولید شده (خروجی CPU)", use_column_width=True)
            
            # دکمه دانلود
            buf = io.BytesIO()
            enhanced_image.save(buf, format="PNG")

            st.download_button(
                label="دانلود عکس تولید شده",
                data=buf.getvalue(),
                file_name="generated_cpu_photo.png",
                mime="image/png"
            )
        except Exception as e:
            st.error(f"خطا در اجرای مدل: {e}")