import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load model dari Hugging Face hub untuk CPU
# Hapus torch_dtype=torch.float16 dan .to("cuda")
print("Memuat model ke CPU... proses ini mungkin butuh beberapa saat.")
pipe = StableDiffusionPipeline.from_pretrained("izzudd/sd-batik-llava")
print("Model berhasil dimuat!")

def generate_image(prompt):
    print(f"Mulai menghasilkan gambar untuk prompt: '{prompt}'")
    # Generate gambar di CPU
    image = pipe(prompt).images[0]
    print("Gambar berhasil dibuat!")
    return image

# Gradio UI
gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Masukkan Deskripsi Batik"),
    outputs=gr.Image(label="Hasil Gambar Batik"),
    title="Generator Batik AI (Versi CPU)",
    description="Masukkan deskripsi motif batik, dan AI akan menghasilkan gambarnya! Harap bersabar, proses di CPU akan memakan waktu beberapa menit."
).launch()