import gradio as gr
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"

pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token='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', torch_dtype=torch.float16, low_cpu_mem_usage=True)
pipe = pipe.to(device)

def inference(diffusion_prompt):
    samples = 4
    generator = torch.Generator(device=device)
    torch.cuda.empty_cache()
    with autocast("cuda"):
        images_list = pipe(
            [diffusion_prompt] * samples,
            height=256, width=384,
            num_inference_steps=50,
        )
        images = []
        for i, image in enumerate(images_list["sample"]):
            images.append(image)
    return images


with gr.Blocks() as demo:
    gr.Markdown("# Text to Image Generator")
    with gr.Row():
                prompt = gr.Textbox(
                    lines=1,
                    placeholder="Enter your prompt..",
                    interactive=True,
                    label="Prompt"
                )
                submit = gr.Button("Run")
    submit.click(fn=inference, inputs=[prompt], outputs=[images])
demo.launch()