add multiprocessing

app_dialogue.py

@@ -10,6 +10,7 @@ from typing import List, Optional, Tuple
 from urllib.parse import urlparse
 from PIL import Image, ImageDraw, ImageFont
 
+import concurrent.futures
 import random
 import gradio as gr
 import PIL
@@ -777,6 +778,28 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
     with gr.Row():
         chatbot.render()
 
+    def generate_meme(
+        i,
+        client,
+        query,
+        image,
+        font_meme_text,
+        all_caps_meme_text,
+        text_at_the_top,
+        generation_args,
+    ):
+        text = client.generate(prompt=query, **generation_args).generated_text
+        if image is not None and text != "":
+            meme_image = make_meme_image(
+                image=image,
+                text=text,
+                font_meme_text=font_meme_text,
+                all_caps_meme_text=all_caps_meme_text,
+                text_at_the_top=text_at_the_top,
+            )
+            meme_image = pil_to_temp_file(meme_image)
+            return meme_image
+
     def model_inference(
         model_selector,
         system_prompt,
@@ -849,21 +872,27 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
 
         query = prompt_list_to_tgi_input(formated_prompt_list)
         all_meme_images = []
-        for i in range(4):
-            text = client.generate(prompt=query, **generation_args).generated_text
-            if image is not None and text != "":
-                meme_image = make_meme_image(
-                    image=image,
-                    text=text,
-                    font_meme_text=font_meme_text,
-                    all_caps_meme_text=all_caps_meme_text,
-                    text_at_the_top=text_at_the_top,
+        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+            futures = [
+                executor.submit(
+                    generate_meme,
+                    i,
+                    client,
+                    query,
+                    image,
+                    font_meme_text,
+                    all_caps_meme_text,
+                    text_at_the_top,
+                    generation_args,
                 )
-                meme_image = pil_to_temp_file(meme_image)
-                all_meme_images.append(meme_image)
-                yield user_prompt_str, all_meme_images, chat_history
-            if i == 3:
-                return
+                for i in range(4)
+            ]
+
+            for future in concurrent.futures.as_completed(futures):
+                meme_image = future.result()
+                if meme_image:
+                    all_meme_images.append(meme_image)
+        return user_prompt_str, all_meme_images, chat_history
 
     gr.on(
         triggers=[textbox.submit, imagebox.upload, submit_btn.click],