import gradio as gr
from transformers import pipeline
import torch
from diffusers import DiffusionPipeline
    
def get_completion(prompt,params):
    # return pipeline(prompt=prompt, height=params['height'], width=params['width'], num_inference_steps=int(params['num_inference_steps']), guidance_scale=params['guidance_scale'])['sample'][0]
    return pipeline(prompt=prompt, height=params['height'], width=params['width'], num_inference_steps=int(params['num_inference_steps']), guidance_scale=params['guidance_scale']).images[0]
def generate(prompt,steps,guidance,width,height):
    params = {
        "num_inference_steps": steps,
        "guidance_scale": guidance,
        "width": width,
        "height": height
    }
    output = get_completion(prompt,params)
    return output
    
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")

with gr.Blocks() as demo:
    gr.Markdown("# Image Generation Demo & Test App by Srinivas")
    gr.Markdown("## Generates an Image based on Your Promt inputted and Optional parameters selected")
    with gr.Row():
        with gr.Column(scale=4):
            prompt = gr.Textbox(label="Your Prompt") #Give prompt some real estate
        with gr.Column(scale=1, min_width=50):
            btn = gr.Button("Submit") #Submit button side by side!
    with gr.Accordion("Advanced options", open=False): #Let's hide the advanced options!
            # negative_prompt = gr.Textbox(label="Negative prompt")
            with gr.Row():
                with gr.Column():
                    steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=25,
                      info="In many steps will the denoiser denoise the image?")
                    guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=7.0,
                      info="Controls how much the text prompt influences the result")
                with gr.Column():
                    width = gr.Slider(label="Width", minimum=64, maximum=512, step=64, value=512)
                    height = gr.Slider(label="Height", minimum=64, maximum=512, step=64, value=512)
    output = gr.Image(label="Result") #Move the output up too

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

demo.launch()