update reqs + small refactor

app_dialogue.py

@@ -546,6 +546,125 @@ def expand_layout():
     return gr.Column(scale=2), gr.Gallery(height=682)
 
 
+def generate_meme(
+    client,
+    query,
+    image,
+    font_meme_text,
+    all_caps_meme_text,
+    text_at_the_top,
+    generation_args,
+):
+    try:
+        text = client.generate(prompt=query, **generation_args).generated_text
+    except Exception as e:
+        logger.error(f"Error {e} while generating meme text")
+        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,
+        )
+        return meme_image
+    else:
+        return None
+
+
+def model_inference(
+    model_selector,
+    system_prompt,
+    user_prompt_str,
+    chat_history,
+    image,
+    decoding_strategy,
+    temperature,
+    max_new_tokens,
+    repetition_penalty,
+    top_p,
+    all_caps_meme_text,
+    text_at_the_top,
+    font_meme_text,
+):
+    chat_history = []
+    if user_prompt_str.strip() == "" and image is None:
+        return "", None, chat_history
+
+    system_prompt = ast.literal_eval(system_prompt)
+    (
+        formated_prompt_list,
+        user_prompt_list,
+    ) = format_user_prompt_with_im_history_and_system_conditioning(
+        system_prompt=system_prompt,
+        current_user_prompt_str=user_prompt_str.strip(),
+        current_image=image,
+        history=chat_history,
+    )
+
+    client_endpoint = API_PATHS[model_selector]
+    client = Client(
+        base_url=client_endpoint,
+        headers={"x-use-cache": "0", "Authorization": f"Bearer {API_TOKEN}"},
+        timeout=45,
+    )
+
+    # Common parameters to all decoding strategies
+    # This documentation is useful to read: https://huggingface.co/docs/transformers/main/en/generation_strategies
+    generation_args = {
+        "max_new_tokens": max_new_tokens,
+        "repetition_penalty": repetition_penalty,
+        "stop_sequences": EOS_STRINGS,
+    }
+
+    assert decoding_strategy in [
+        "Greedy",
+        "Top P Sampling",
+    ]
+    if decoding_strategy == "Greedy":
+        generation_args["do_sample"] = False
+    elif decoding_strategy == "Top P Sampling":
+        generation_args["temperature"] = temperature
+        generation_args["do_sample"] = True
+        generation_args["top_p"] = top_p
+
+    chat_history.append([prompt_list_to_markdown(user_prompt_list), ""])
+
+    query = prompt_list_to_tgi_input(formated_prompt_list)
+    all_meme_images = []
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+        futures = [
+            executor.submit(
+                generate_meme,
+                client,
+                query,
+                image,
+                font_meme_text,
+                all_caps_meme_text,
+                text_at_the_top,
+                generation_args,
+            )
+            for i in range(4)
+        ]
+
+        for future in concurrent.futures.as_completed(futures):
+            meme_image = future.result(timeout=45)
+            if meme_image:
+                all_meme_images.append(meme_image)
+    return user_prompt_str, all_meme_images, chat_history
+
+
+def remove_last_turn(chat_history):
+    if len(chat_history) == 0:
+        return chat_history, "", ""
+    last_interaction = chat_history[-1]
+    chat_history = chat_history[:-1]
+    chat_update = chat_history
+    text_update = last_interaction[0]
+    return chat_update, text_update, ""
+
+
 textbox = gr.Textbox(
     placeholder="Upload an image and ask the AI to create a meme!",
     show_label=False,
@@ -764,115 +883,6 @@ 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,
-    ):
-        try:
-            text = client.generate(prompt=query, **generation_args).generated_text
-        except Exception as e:
-            logger.error(f"Error {e} while generating meme text")
-            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,
-            )
-            return meme_image
-        else:
-            return None
-
-    def model_inference(
-        model_selector,
-        system_prompt,
-        user_prompt_str,
-        chat_history,
-        image,
-        decoding_strategy,
-        temperature,
-        max_new_tokens,
-        repetition_penalty,
-        top_p,
-        all_caps_meme_text,
-        text_at_the_top,
-        font_meme_text,
-    ):
-        chat_history = []
-        if user_prompt_str.strip() == "" and image is None:
-            return "", None, chat_history
-
-        system_prompt = ast.literal_eval(system_prompt)
-        (
-            formated_prompt_list,
-            user_prompt_list,
-        ) = format_user_prompt_with_im_history_and_system_conditioning(
-            system_prompt=system_prompt,
-            current_user_prompt_str=user_prompt_str.strip(),
-            current_image=image,
-            history=chat_history,
-        )
-
-        client_endpoint = API_PATHS[model_selector]
-        client = Client(
-            base_url=client_endpoint,
-            headers={"x-use-cache": "0", "Authorization": f"Bearer {API_TOKEN}"},
-            timeout=45,
-        )
-
-        # Common parameters to all decoding strategies
-        # This documentation is useful to read: https://huggingface.co/docs/transformers/main/en/generation_strategies
-        generation_args = {
-            "max_new_tokens": max_new_tokens,
-            "repetition_penalty": repetition_penalty,
-            "stop_sequences": EOS_STRINGS,
-        }
-
-        assert decoding_strategy in [
-            "Greedy",
-            "Top P Sampling",
-        ]
-        if decoding_strategy == "Greedy":
-            generation_args["do_sample"] = False
-        elif decoding_strategy == "Top P Sampling":
-            generation_args["temperature"] = temperature
-            generation_args["do_sample"] = True
-            generation_args["top_p"] = top_p
-
-        chat_history.append([prompt_list_to_markdown(user_prompt_list), ""])
-
-        query = prompt_list_to_tgi_input(formated_prompt_list)
-        all_meme_images = []
-        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,
-                )
-                for i in range(4)
-            ]
-
-            for future in concurrent.futures.as_completed(futures):
-                meme_image = future.result(timeout=45)
-                if meme_image:
-                    all_meme_images.append(meme_image)
-        return user_prompt_str, all_meme_images, chat_history
-
     gr.on(
         triggers=[
             textbox.submit,
@@ -906,15 +916,6 @@ with gr.Blocks(title="AI Meme Generator", theme=gr.themes.Base(), css=css) as de
         outputs=[textbox, generated_memes_gallery, chatbot],
     )
 
-    def remove_last_turn(chat_history):
-        if len(chat_history) == 0:
-            return chat_history, "", ""
-        last_interaction = chat_history[-1]
-        chat_history = chat_history[:-1]
-        chat_update = chat_history
-        text_update = last_interaction[0]
-        return chat_update, text_update, ""
-
     regenerate_btn.click(
         fn=remove_last_turn,
         inputs=chatbot,

requirements.txt

@@ -9,10 +9,8 @@ opencv-python
 numpy
 accelerate
 joblib
-deepspeed
 parameterized
 einops
 pynvml
 sentencepiece
-text_generation
-https://gradio-builds.s3.amazonaws.com/2060bfe3e7eb57fb9b5c8695ebfc900469263d1f/gradio-3.46.0-py3-none-any.whl
+text_generation