Update app.py

app.py

@@ -1,16 +1,31 @@
 import gradio as gr
 from transformers import pipeline
+from diffusers import StableDiffusionPipeline
 
 playground = gr.Blocks()
 
 image_pipe = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 summary_pipe = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
 ner_pipe = pipeline("ner", model="dslim/bert-base-NER")
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
 
 def launch_image_pipe(input):
     out = image_pipe(input)
     return out[0]['generated_text']
 
+def base64_to_pil(img_base64):
+    base64_decoded = base64.b64decode(img_base64)
+    byte_stream = io.BytesIO(base64_decoded)
+    pil_image = Image.open(byte_stream)
+    return pil_image
+
+def image_generate(prompt):
+    # output = get_completion(prompt)
+    image = pipe(prompt).images[0]
+    result_image = base64_to_pil(image)
+    return result_image
+    
 def translate(input_text, source, target):
     try:
       model = f"Helsinki-NLP/opus-mt-{source}-{target}"
@@ -100,12 +115,17 @@ with playground:
                 with gr.Column():
                     img = gr.Image(type='pil')
                 with gr.Column():
-                    generated_textbox = gr.Textbox(lines=2, placeholder="", label="Generated Text")
-                    
-            image_pipeline_button.click(launch_image_pipe,
-                                        inputs=[img],
-                                        outputs=[generated_textbox])
-            
+                    with gr.Row():
+                        with gr.Column(scale=4):
+                            generated_textbox = gr.Textbox(lines=2, placeholder="", label="Generated Text")
+                        with gr.Column(scale=1):
+                            image_generation_button = gr.Button(value="GEN")
+                    with gr.Row():
+                        generated_image = gr.Image(label="Generated Image"))
+                        
+            image_pipeline_button.click(launch_image_pipe, inputs=[img], outputs=[generated_textbox])
+            image_generation_button.click(image_generate, inputs=[generated_textbox], outputs=[generated_image])
+        
         with gr.TabItem("Text"):
             with gr.Row():
                 with gr.Column(scale=4):