import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler
import gradio as gr

pipe = StableDiffusionPipeline.from_pretrained("MVRL/GeoSynth")
# scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base")
# pipe.scheduler = scheduler

def process(prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta):
    generator = torch.manual_seed(seed)
    output_images = pipe(prompt,
                         height=image_resolution,
                         width=image_resolution,
                         num_inference_steps=ddim_steps,
                         guidance_scale=scale,
                         negative_prompt=n_prompt,
                         num_images_per_prompt=num_samples,
                         eta=eta,
                         generator=generator,
                        ).images
    return output_images

block = gr.Blocks().queue()
with block:
    with gr.Row():
        gr.Markdown(
            """
            # GeoSynth: Contextually-Aware High-Resolution Satellite Image Synthesis
            Srikumar Sastry*, Subash Khanal, Aayush Dhakal, Nathan Jacobs (*Corresponding Author)<br>
            """
        )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            run_button = gr.Button(value="Run")
            with gr.Accordion("Advanced options", open=True):
                num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
                image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
                ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
                scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
                eta = gr.Number(label="eta (DDIM)", value=0.0)
                n_prompt = gr.Textbox(label="Negative Prompt",
                                      value='')
        with gr.Column():
            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery")
    ips = [prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta]
    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])

block.launch()