import streamlit as st
from diffusers import StableDiffusionImg2ImgPipeline
import torch
from PIL import Image
import os

# Inisialisasi pipeline model
@st.cache_resource
def load_model():
    model = StableDiffusionImg2ImgPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5", 
        torch_dtype=torch.float16
    )
    model.to("cuda")  # Gunakan GPU untuk performa optimal
    return model

# Header aplikasi
st.title("Stable Diffusion Image-to-Image Generator")
st.write("🎨 Unggah gambar dan tambahkan deskripsi untuk menghasilkan gambar baru!")

# Input teks dari pengguna
prompt = st.text_input("Masukkan teks deskripsi gambar:", "")
negative_prompt = st.text_input("Masukkan teks untuk menghindari elemen tertentu (Negative Prompt):", "")

# Opsi parameter tambahan
guidance_scale = st.slider("Guidance Scale (Semakin tinggi, semakin spesifik gambar)", 1.0, 20.0, 7.5)
strength = st.slider("Strength (Pengaruh gambar asli, 0=mirip, 1=berbeda)", 0.1, 1.0, 0.75)
steps = st.slider("Jumlah Langkah Inferensi", 10, 50, 30)

# Input gambar dari pengguna
uploaded_file = st.file_uploader("Unggah gambar awal (Input Image)", type=["jpg", "jpeg", "png"])

# Tombol untuk menghasilkan gambar
if st.button("Generate") and prompt and uploaded_file:
    with st.spinner("Sedang menghasilkan gambar..."):
        # Load model
        model = load_model()
        
        # Buka dan proses gambar
        init_image = Image.open(uploaded_file).convert("RGB")
        init_image = init_image.resize((512, 512))  # Sesuaikan ukuran gambar ke 512x512
        
        # Proses image-to-image
        image = model(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            image=init_image,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=steps
        ).images[0]  # Ambil gambar pertama dari hasil

        # Tampilkan hasil gambar
        st.image(image, caption="Gambar yang dihasilkan", use_column_width=True)

        # Simpan gambar ke folder assets
        if not os.path.exists("assets/"):
            os.makedirs("assets/")
        save_path = f"assets/{prompt.replace(' ', '_')}_neg_{negative_prompt.replace(' ', '_') if negative_prompt else 'none'}.png"
        image.save(save_path)
        st.success(f"Gambar telah disimpan di {save_path}")
else:
    if not prompt:
        st.warning("Masukkan deskripsi untuk memulai.")
    if not uploaded_file:
        st.warning("Unggah gambar untuk memulai.")

# Tampilan gambar sebelumnya
st.subheader("Gambar Sebelumnya")
if os.path.exists("assets/"):
    for img_file in os.listdir("assets/"):
        st.image(f"assets/{img_file}", caption=img_file, use_column_width=True)