Adding Image Generation

app.py

@@ -1,5 +1,6 @@
 from transformers import pipeline
 import gradio as gr
+from diffusers import DiffusionPipeline
 
 # 1. text summarizer
 summarizer = pipeline("summarization", model = "facebook/bart-large-cnn")
@@ -15,10 +16,14 @@ def get_ner(text):
 
 # 3. Image Captioning
 caption_model = pipeline("image-to-text", model = "Salesforce/blip-image-captioning-base")
-# processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 def get_caption(img):
     output = caption_model(img)
     return output[0]["generated_text"]
+
+# 4. Image Generation
+img_model = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+def get_img(prompt):
+    return img_model(prompt).images[0]
     
 
 demo = gr.Blocks()
@@ -31,15 +36,18 @@ with demo:
         sum_btn.click(get_summary, sum_input, sum_output)
     with gr.Tab("Named Entity Recognition"):
         ner_input = [gr.Textbox(label="Text to find Entities", placeholder = "Enter text...", lines = 4)]
-        # ner_output = gr.Textbox()
         ner_output = [gr.HighlightedText(label="Text with entities")]
         ner_btn = gr.Button("Generate entities")
-        # allow_flagging = "never"
         ner_btn.click(get_ner, ner_input, ner_output)
     with gr.Tab("Image Captioning"):
         cap_input = [gr.Image(label="Upload Image", type="pil")]
         cap_btn = gr.Button("Generate Caption")
         cap_output = [gr.Textbox(label="Caption")]
         cap_btn.click(get_caption, cap_input, cap_output)
+    with gr.Tab("Image Generation"):
+        img_input = [gr.Textbox(label="Your Text")]
+        img_btn = gr.Button("Generate Image")
+        img_output = [gr.Image(label="Generated Image")]
+        img_btn.click(get_img, img_input, img_output)
 
 demo.launch()