import gradio as gr
from diffusers import StableDiffusionXLPipeline
import torch

# Lade das Modell
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    use_safetensors=True,
)
pipe.to("cuda")

# Bildgenerierungsfunktion
def generate(prompt, negative_prompt, width, height, steps, guidance, seed, progress=gr.Progress()):
    generator = torch.manual_seed(seed) if seed != -1 else None
    with progress.tqdm(total=steps) as pbar:
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            num_inference_steps=steps,
            guidance_scale=guidance,
            generator=generator,
            callback=lambda i, t, s: pbar.update(1)
        ).images[0]
    return image

# UI
with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion XL Generator")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="z. B. futuristic city at sunset", lines=2)
        negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="blurry, low quality, distorted", lines=1)

    with gr.Row():
        width = gr.Slider(label="Width", minimum=512, maximum=1024, step=64, value=768)
        height = gr.Slider(label="Height", minimum=512, maximum=1024, step=64, value=768)

    with gr.Row():
        steps = gr.Slider(label="Steps", minimum=10, maximum=100, step=1, value=30)
        guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=7.5)
        seed = gr.Slider(label="Seed (set -1 for random)", minimum=-1, maximum=999999, step=1, value=-1)

    run_button = gr.Button("Generate")
    output = gr.Image(label="Output")

    run_button.click(
        fn=generate,
        inputs=[prompt, negative_prompt, width, height, steps, guidance, seed],
        outputs=output,
    )

demo.launch()