import gradio as gr
from optimum.intel.openvino import OVStableDiffusionPipeline

# model AI
model_id = "segmind/tiny-sd"

pipe = OVStableDiffusionPipeline.from_pretrained(
    model_id,
    export=True
)

# fungsi generate
def generate_image(prompt, steps, width, height):

    image = pipe(
        prompt,
        negative_prompt="blurry, bad quality, low quality",
        num_inference_steps=int(steps),
        width=int(width),
        height=int(height)
    ).images[0]

    return image


# custom css aesthetic
custom_css = """
body {
    background: #0f1117;
}

.gradio-container {
    background: linear-gradient(
        135deg,
        #0f1117,
        #1c1f2b
    );
    color: white;
}

h1 {
    text-align: center;
    font-size: 42px !important;
    font-weight: bold;
    background: linear-gradient(to right, #8b5cf6, #ec4899);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

textarea {
    background: #1e2230 !important;
    color: white !important;
    border-radius: 15px !important;
}

button {
    background: linear-gradient(
        to right,
        #8b5cf6,
        #ec4899
    ) !important;

    color: white !important;
    border: none !important;
    border-radius: 14px !important;
    font-size: 18px !important;
    padding: 12px !important;
}

footer {
    display: none !important;
}
"""

# UI website
with gr.Blocks(
    css=custom_css,
    theme=gr.themes.Soft()
) as app:

    gr.Markdown("""
    # ✨ Pudel AI Generator
    
    Create beautiful AI images instantly
    """)

    with gr.Row():

        with gr.Column(scale=1):

            prompt = gr.Textbox(
                label="Prompt",
                placeholder="contoh: anime girl cyberpunk"
            )

            steps = gr.Slider(
                10,
                50,
                value=20,
                step=1,
                label="Steps"
            )

            width = gr.Slider(
                256,
                1024,
                value=512,
                step=64,
                label="Width"
            )

            height = gr.Slider(
                256,
                1024,
                value=512,
                step=64,
                label="Height"
            )

            generate_btn = gr.Button(
                "✨ Generate Image"
            )

        with gr.Column(scale=1):

            output = gr.Image(
                label="Hasil AI",
                type="filepath",
                height=500
            )

    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt,
            steps,
            width,
            height
        ],
        outputs=output
    )

app.launch()