import gradio as gr
from diffusers import DiffusionPipeline
import torch
import os
import time

# Konfigurasi model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "cagliostrolab/animagine-xl-3.1"

pipe = DiffusionPipeline.from_pretrained(
    model_repo_id, 
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, 
    use_safetensors=True,
)
pipe.to(device)

# Fungsi inference
def infer(prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
    try:
        # Generate image
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=int(width),
            height=int(height),
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_inference_steps),
        ).images[0]

        # Simpan hasil gambar di folder output dengan nama unik berdasarkan timestamp
        os.makedirs("./output", exist_ok=True)
        output_path = f"./output/generated_image_{int(time.time())}.png"
        image.save(output_path)

        return image
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio interface
with gr.Blocks() as demo:
    # Pesan pemberitahuan jika menggunakan CPU
    gr.Markdown(
        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
    )
    
    gr.Markdown("## Text-to-Image Generator with animagine-xl-3.1")

    # Output gambar di atas
    result_image = gr.Image(label="Generated Image", elem_id="result-image")

    # Input parameter di bawah
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Masukkan prompt Anda di sini",
                value="1girl, souryuu asuka langley, neon genesis evangelion, solo, upper body, v, smile, looking at viewer, outdoors, night",
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Masukkan negative prompt untuk menghindari elemen tidak diinginkan",
                value="nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"
            )

            # Accordion untuk pengaturan lanjutan
            with gr.Accordion("Advanced Settings", open=False):
                width = gr.Dropdown(
                    label="Width",
                    choices=["256", "512", "768", "832", "896", "1024"],
                    value="832",
                )
                height = gr.Dropdown(
                    label="Height",
                    choices=["256", "512", "768", "832", "896", "1216", "1024"],
                    value="1216",
                )
                guidance_scale = gr.Dropdown(
                    label="Guidance Scale",
                    choices=[str(i / 10) for i in range(0, 201, 10)],  # 0.0 to 20.0
                    value="7.0",
                )
                num_inference_steps = gr.Dropdown(
                    label="Number of Inference Steps",
                    choices=[str(i) for i in range(1, 101)],  # 1 to 100
                    value="28",
                )

            run_button = gr.Button("Generate Image")

    # Hubungkan fungsi infer ke UI
    run_button.click(
        fn=infer,
        inputs=[prompt, negative_prompt, width, height, guidance_scale, num_inference_steps],
        outputs=result_image,
    )

# Jalankan aplikasi
if __name__ == "__main__":
    demo.launch()