Update app.py

app.py

@@ -1,22 +1,39 @@
-from transformers import pipeline, Wav2Vec2Processor, AutoModelForCTC
+import gradio as gr
 import torch
 from torch import autocast
 from diffusers import StableDiffusionPipeline
-import gradio as gr
 
 model_id = "CompVis/stable-diffusion-v1-4"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token='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', revision="fp16", torch_dtype=torch.float16)
-device = 'cuda'
-#has_cuda = torch.cuda.is_available()
-#device = torch.device('cpu' if not has_cuda else 'cuda')    
+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 convert(prompt):
-        with autocast("cuda"):
-                image = pipe(prompt)["sample"][0]
-        return image
+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
+
 
-gr.Interface(convert,
-             inputs = [gr.inputs.Textbox(label="Enter text")],    
-             outputs = [gr.outputs.Image(label="Generated Image")], 
-             title="Text to Image Generation").launch()
+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()