import torch
from diffusers import DiffusionPipeline
import gradio as gr
import numpy as np
import random

# Inisialisasi model
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16

repo = "stabilityai/stable-diffusion-3-medium-diffusers"
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium")

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1344

def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    image = pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator
    ).images[0]

    return image, seed

examples = [
    ["Astronaut in the forest, cool color palette, calm colors, detailed, 8k"],
    ["An astronaut riding a green horse"],
    ["A delicious piece of ceviche cheesecake"]
]

# CSS untuk antarmuka yang lebih menarik
css = """
body {
    font-family: 'Arial', sans-serif;
    background-color: #f0f0f0;
    color: #333;
    margin: 0;
    padding: 0;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    height: 100vh;
}

#col-container {
    max-width: 800px;
    width: 100%;
    margin: 20px;
    padding: 20px;
    background-color: #fff;
    border-radius: 10px;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}

#title {
    font-size: 28px;
    font-weight: bold;
    text-align: center;
    margin-bottom: 20px;
}

.input-box {
    margin-bottom: 20px;
}

.input-box label {
    font-size: 18px;
    font-weight: bold;
    display: block;
    margin-bottom: 5px;
}

.input-box input,
.input-box select,
.input-box textarea {
    width: calc(100% - 20px);
    padding: 10px;
    font-size: 16px;
    border: 1px solid #ccc;
    border-radius: 4px;
    margin-bottom: 10px;
}

.input-box input:focus,
.input-box select:focus,
.input-box textarea:focus {
    outline: none;
    border-color: #007BFF;
}

.button {
    background-color: #007BFF;
    color: white;
    padding: 10px 20px;
    border: none;
    border-radius: 4px;
    cursor: pointer;
    font-size: 16px;
    width: 100%;
}

.button:hover {
    background-color: #0056b3;
}

.result {
    margin-top: 20px;
    text-align: center;
}

.result img {
    max-width: 100%;
    border-radius: 8px;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("""
    # Stable Diffusion 3 Medium Demo
    Generate images using Stable Diffusion 3. Try different prompts and settings to see the results.
    """, elem_id="title")

    with gr.Column(elem_id="col-container"):
        prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt here...", elem_classes="input-box")
        negative_prompt = gr.Textbox(label="Negative Prompt", lines=1, placeholder="Enter negative prompt (optional)", elem_classes="input-box")
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, elem_classes="input-box")
        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True, elem_classes="input-box")
        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=512, elem_classes="input-box")
        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=512, elem_classes="input-box")
        guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.5, value=7.5, elem_classes="input-box")
        num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=100, step=1, value=28, elem_classes="input-box")

        generate_button = gr.Button("Generate Image", elem_classes="button")
        result_image = gr.Image(elem_classes="result")

        generate_button.click(infer, inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps], outputs=[result_image, seed])

    gr.Examples(examples=examples, inputs=[prompt])

demo.launch()