import gradio as gr
from diffusers import DiffusionPipeline

# import torch
# from diffusers import DDPMScheduler, UNet2DModel
# from PIL import Image
# import numpy as np

# pipeline = DiffusionPipeline.from_pretrained("google/ddpm-cat-256")
pipeline = DiffusionPipeline.from_pretrained("google/ddpm-celebahq-256")
# pipeline.to("cuda")


def erzeuge(prompt):
    return pipeline(prompt).images  # [0]


with gr.Blocks() as demo:
    with gr.Column(variant="panel"):
        with gr.Row(variant="compact"):
            text = gr.Textbox(
                label="Deine Beschreibung:",
                show_label=False,
                max_lines=1,
                placeholder="Bildbeschrei",
            )
            btn = gr.Button("erzeuge Bild")

        gallery = gr.Gallery(
            label="Erzeugtes Bild", show_label=False, elem_id="gallery"
        )

    btn.click(erzeuge, inputs=[text], outputs=[gallery])
    text.submit(erzeuge, inputs=[text], outputs=[gallery])

if __name__ == "__main__":
    demo.launch()