import gradio as gr
import numpy as np
import torch
from PIL import Image
from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(model_id).to('cpu')

def infer(prompt, negative, steps, scale, seed):
    generator = torch.Generator(device='cpu').manual_seed(seed)
    img = pipe(
            prompt,
            height=512, 
            width=512,
            num_inference_steps=steps,
            guidance_scale=scale,
            negative_prompt = negative,
            generator=generator,
        ).images
    return img

block = gr.Blocks()

with block:
  with gr.Group():
    with gr.Box():
      with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
            with gr.Column():
                text = gr.Textbox(
                    label="Enter your prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    ).style(
                        border=(True, False, True, True),
                        rounded=(True, False, False, True),
                        container=False,
                    )

                negative = gr.Textbox(
                    label="Enter your negative prompt",
                    show_label=False,
                    placeholder="Enter a negative prompt",
                    elem_id="negative-prompt-text-input",
                    ).style(
                        border=(True, False, True, True),
                        rounded=(True, False, False, True),container=False,
                    )
                
            btn = gr.Button("Generate image").style(
                        margin=False,
                        rounded=(False, True, True, False),
                    )  
    gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(columns=(1, 2), height="auto")
        
    with gr.Row(elem_id="advanced-options"):
          samples = gr.Slider(label="Images", minimum=1, maximum=1, value=1, step=1, interactive=False)
          steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=12, step=1, interactive=True)
          scale = gr.Slider(label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1, interactive=True)
          seed = gr.Slider(label="Random seed",minimum=0,maximum=2147483647,step=1,randomize=True,interactive=True)

    btn.click(infer, inputs=[text, negative, steps, scale, seed], outputs=[gallery])

block.launch(show_api=False)