import gradio as gr
from diffusers import DDPMPipeline
import torch

# Load model
pipe = DDPMPipeline.from_pretrained("Docty/pipecorrode", torch_dtype=torch.float16)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Generation function
def generate_images(num_images: int = 1, steps: int = 50):
    output = pipe(num_inference_steps=steps, batch_size=num_images)
    return output.images

# Gradio Interface
gr.Interface(
    fn=generate_images,
    inputs=[
        gr.Slider(1, 8, step=1, label="Number of Images"),
        gr.Slider(10, 100, step=10, label="Sampling Steps"),
    ],
    outputs=gr.Gallery(label="Generated Images"),
    title="Unconditional Diffusion Generator"
).launch()