use standard generate

app_dialogue.py

@@ -355,6 +355,7 @@ def resize_with_ratio(image: PIL.Image.Image, fixed_width: int) -> PIL.Image.Ima
 
     return resized_img
 
+
 def make_new_lines(draw, image, font, text_is_too_long, lines, num_lines, num_loops):
     max_len_increment = 0
     while text_is_too_long and max_len_increment < 10:
@@ -389,6 +390,7 @@ def make_new_lines(draw, image, font, text_is_too_long, lines, num_lines, num_lo
         lines = new_lines
     return lines, text_width, text_is_too_long
 
+
 def test_font_size(
     draw,
     image,
@@ -546,9 +548,11 @@ def format_user_prompt_with_im_history_and_system_conditioning(
 
     return resulting_list, current_user_prompt_list
 
+
 def expand_layout():
     return gr.Column(scale=2), gr.Gallery(height=682)
 
+
 textbox = gr.Textbox(
     placeholder="Upload an image and ask the AI to create a meme!",
     show_label=False,
@@ -565,17 +569,17 @@ chatbot = gr.Chatbot(
     visible=False,
     avatar_images=[None, BOT_AVATAR],
 )
-css='''
+css = """
 .gradio-container{max-width: 1000px!important}
 h1{display: flex;align-items: center;justify-content: center;gap: .25em}
 *{transition: width 0.5s ease, flex-grow 0.5s ease}
-'''
+"""
 with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as demo:
     with gr.Row(scale=0.5):
         gr.HTML(
             """<h1 align="center">AI Meme Generator <span style="font-size: 13px;">powered by <a href="https://huggingface.co/blog/idefics">IDEFICS</a></span><img width=40 height=40 src="https://cdn-uploads.huggingface.co/production/uploads/624bebf604abc7ebb01789af/v770xGti5vH1SYLBgyOO_.png" /></h1>"""
         )
-        
+
     with gr.Row(elem_id="model_selector_row"):
         model_selector = gr.Dropdown(
             choices=MODELS,
@@ -587,7 +591,7 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
             visible=False,
         )
     with gr.Row(equal_height=True):
-        #scale=2 when expanded
+        # scale=2 when expanded
         with gr.Column(scale=4, min_width=250) as upload_area:
             imagebox = gr.Image(
                 type="filepath", label="Image to meme", height=272, visible=True
@@ -596,13 +600,20 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
                 with gr.Row():
                     textbox.render()
                 with gr.Row():
-                    submit_btn = gr.Button(value="▶️ Submit", visible=True, min_width=120)
+                    submit_btn = gr.Button(
+                        value="▶️ Submit", visible=True, min_width=120
+                    )
                     clear_btn = gr.ClearButton(
                         [textbox, imagebox, chatbot], value="🧹 Clear", min_width=120
                     )
-                    regenerate_btn = gr.Button(value="🔄 Regenerate", visible=True, min_width=120)
+                    regenerate_btn = gr.Button(
+                        value="🔄 Regenerate", visible=True, min_width=120
+                    )
                     upload_btn = gr.UploadButton(
-                        "📁 Upload image", file_types=["image"], visible=False, min_width=120
+                        "📁 Upload image",
+                        file_types=["image"],
+                        visible=False,
+                        min_width=120,
                     )
             with gr.Accordion(
                 "Advanced settings", open=False, visible=True
@@ -717,8 +728,9 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
                 show_download_button=True,
                 show_share_button=True,
                 columns=[2],
-                object_fit="contain",height=428
-            ) #height 600 when expanded
+                object_fit="contain",
+                height=428,
+            )  # height 600 when expanded
     with gr.Row(equal_height=True):
         with gr.Box(elem_id="gallery_box"):
             gallery_type_choice = gr.Radio(
@@ -799,7 +811,7 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
         client = Client(
             base_url=client_endpoint,
             headers={"x-use-cache": "0", "Authorization": f"Bearer {API_TOKEN}"},
-            timeout=30
+            timeout=30,
         )
 
         # Common parameters to all decoding strategies
@@ -838,61 +850,29 @@ 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):
-            stream = client.generate_stream(prompt=query, **generation_args)
-
-            acc_text = ""
-            full_text = ""
-            for idx, response in enumerate(stream):
-                text_token = response.token.text
-
-                if response.details:
-                    # That's the exit condition
-                    if image is not None and full_text != "":
-                        meme_image = make_meme_image(
-                            image=image,
-                            text=full_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)
-                        all_meme_images.append(meme_image)
-                        yield user_prompt_str, all_meme_images, chat_history
-                    if i == 3:
-                        return
-
-                if text_token in STOP_SUSPECT_LIST:
-                    acc_text += text_token
-                    continue
-
-                if idx == 0 and text_token.startswith(" "):
-                    text_token = text_token.lstrip()
-
-                acc_text += text_token
-                # Commented to not have a chatbot history that could confuse user
-
-                # last_turn = chat_history.pop(-1)
-                # last_turn[-1] += acc_text
-                # if last_turn[-1].endswith("\nUser"):
-                #     # Safeguard: sometimes (rarely), the model won't generate the token `<end_of_utterance>` and will go directly to generating `\nUser:`
-                #     # It will thus stop the generation on `\nUser:`. But when it exits, it will have already generated `\nUser`
-                #     # This post-processing ensures that we don't have an additional `\nUser` wandering around.
-                #     last_turn[-1] = last_turn[-1][:-5]
-                # chat_history.append(last_turn)
-                # yield "", None, chat_history
-                full_text += acc_text
-                acc_text = ""
+            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)
+                all_meme_images.append(meme_image)
+                yield user_prompt_str, all_meme_images, chat_history
+            if i == 3:
+                return
+
     gr.on(
-        triggers=[
-            textbox.submit,
-            imagebox.upload,
-            submit_btn.click
-        ],
+        triggers=[textbox.submit, imagebox.upload, submit_btn.click],
         fn=expand_layout,
         outputs=[upload_area, generated_memes_gallery],
-        queue=False
+        queue=False,
+    ).then(
+        fn=lambda: "", inputs=[], outputs=[generated_memes_gallery], queue=False
     ).then(
-        fn=lambda: "", inputs=[], outputs=[generated_memes_gallery], queue=False).then(
         fn=model_inference,
         inputs=[
             model_selector,
@@ -926,7 +906,7 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
         fn=remove_last_turn,
         inputs=chatbot,
         outputs=[chatbot, textbox, generated_memes_gallery],
-        queue=False
+        queue=False,
     ).then(
         fn=model_inference,
         inputs=[
@@ -953,13 +933,22 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
 
     upload_btn.upload(add_file, [upload_btn], [imagebox, upload_btn], queue=False)
     submit_btn.click(
-        lambda: gr.update(label="📁 Upload image", interactive=True), [], upload_btn, queue=False
+        lambda: gr.update(label="📁 Upload image", interactive=True),
+        [],
+        upload_btn,
+        queue=False,
     )
     textbox.submit(
-        lambda: gr.update(label="📁 Upload image", interactive=True), [], upload_btn, queue=False
+        lambda: gr.update(label="📁 Upload image", interactive=True),
+        [],
+        upload_btn,
+        queue=False,
     )
     clear_btn.click(
-        lambda: gr.update(label="📁 Upload image", interactive=True), [], upload_btn, queue=False
+        lambda: gr.update(label="📁 Upload image", interactive=True),
+        [],
+        upload_btn,
+        queue=False,
     )
     gallery_type_choice.change(
         fn=choose_gallery,
@@ -971,11 +960,9 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
         fn=add_file_gallery,
         inputs=[template_gallery],
         outputs=[textbox, imagebox, generated_memes_gallery],
-        queue=False
+        queue=False,
     ).success(
-        fn=expand_layout,
-        outputs=[upload_area, generated_memes_gallery],
-        queue=False
+        fn=expand_layout, outputs=[upload_area, generated_memes_gallery], queue=False
     ).success(
         fn=model_inference,
         inputs=[
@@ -1000,7 +987,10 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
         ],
     )
     demo.load(
-        fn=choose_gallery, inputs=[gallery_type_choice], outputs=[template_gallery], queue=False
+        fn=choose_gallery,
+        inputs=[gallery_type_choice],
+        outputs=[template_gallery],
+        queue=False,
     )
 demo.queue(concurrency_count=4, max_size=40)
-demo.launch(max_threads=400)
+demo.launch(max_threads=400)