import gradio as gr
from diffusers import DDPMPipeline
import torch 

model = 'alexktrs/CumulusCloudsGenerator'

if torch.cuda.is_available():
    device='cuda'
else:
    device='cpu'

generator = DDPMPipeline.from_pretrained(model)
generator.to(device)

def generate(num_images, num_inference_steps):
    images=[]
    print(num_images)
    if num_images==None:
        num_images=1
    num_images=int(num_images)

    for i in range(num_images):
        image = generator(num_inference_steps=num_inference_steps).images[0]
        images.append(image)
    return images

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
                
                # Generate Cumulus Clouds

                """)
    
    gallery=gr.Gallery(type="pil")

    with gr.Row():
        slider=gr.Slider(label='Inference Steps', minimum=1, maximum=100, step=1, value=20)
        n=gr.Number(label='Number of Generated Images', minimum=1, maximum=4, value=2)

    btn = gr.Button("Generate Clouds")
    btn.click(fn=generate, inputs=[n, slider], outputs=gallery)

demo.launch()