optimize the prompt

app.py

@@ -2,7 +2,7 @@ import os
 
 import gradio as gr
 
-from gradio_components.image import generate_caption
+from gradio_components.image import generate_caption, improve_prompt
 from gradio_components.prediction import predict, transcribe
 
 theme = gr.themes.Glass(
@@ -83,7 +83,7 @@ def generate_prompt(difficulty, style):
         "Medum": "player who has 2-3 years experience",
         "Hard": "player who has more than 4 years experiences",
     }
-    prompt = "piano only music for a {} to pratice with the touch of {}".format(
+    prompt = "piano only music for a {} to practice with the touch of {}".format(
         _DIFFICULTY_MAPPIN[difficulty], style
     )
     return prompt
@@ -106,6 +106,16 @@ def toggle_melody_condition(melody_condition):
         )
 
 
+def toggle_custom_prompt(customize, difficulty, style):
+    if customize:
+        return gr.Textbox(label="Type your prompt", interactive=True, visible=True)
+    else:
+        prompt = generate_prompt(difficulty, style)
+        return gr.Textbox(
+            label="Generated Prompt", value=prompt, interactive=False, visible=True
+        )
+
+
 def show_caption(show_caption_condition, description, prompt):
     if show_caption_condition:
         return (
@@ -145,6 +155,17 @@ def show_caption(show_caption_condition, description, prompt):
         )
 
 
+def optimize_fn(prompt):
+    message_object, prompt = improve_prompt(prompt)
+    return prompt
+
+
+def display_prompt(prompt):
+    return gr.Textbox(
+        label="Generated Prompt", value=prompt, interactive=False, visible=True
+    )
+
+
 def post_submit(show_caption, model_path, image_input):
     _, description, prompt = generate_caption(image_input, model_path)
     return (
@@ -210,12 +231,36 @@ def UI():
                         )
                         if style == "Others":
                             style = gr.Textbox(label="Type your music genre")
-                        prompt = generate_prompt(difficulty.value, style.value)
                         customize = gr.Checkbox(
-                            label="Customize the prompt", interactive=True
+                            label="Customize the prompt", interactive=True, value=False
+                        )
+                        _init_prompt = generate_prompt(difficulty.value, style.value)
+                        prompt = gr.Textbox(
+                            label="",
+                            value=_init_prompt,
+                            interactive=False,
+                            visible=False,
+                        )
+                        customize.change(
+                            fn=toggle_custom_prompt,
+                            inputs=[customize, difficulty, style],
+                            outputs=prompt,
                         )
-                        if customize:
-                            prompt = gr.Textbox(label="Type your prompt")
+                        print(prompt)
+                        with gr.Column():
+                            optimize = gr.Button(
+                                "Optimize the prompt", interactive=True
+                            )
+                        with gr.Column():
+                            show_prompt = gr.Button("Show the prompt", interactive=True)
+                            prompt_text = gr.Textbox(
+                                "Optimized Prompt", interactive=False, visible=False
+                            )
+                        optimize.click(optimize_fn, inputs=[prompt], outputs=prompt)
+                        show_prompt.click(
+                            display_prompt, inputs=[prompt], outputs=prompt_text
+                        )
+
                 with gr.Column():
                     with gr.Row():
                         melody = gr.Audio(
@@ -385,12 +430,12 @@ def UI():
                 fn=post_submit,
                 inputs=[show_prompt, model_path, image_input],
                 outputs=[description, prompt, generate],
-                )
+            )
             show_prompt.change(
                 fn=show_caption,
                 inputs=[show_prompt, description, prompt],
                 outputs=[description, prompt, generate],
-                )
+            )
             transcribe_button.click(transcribe, inputs=[output_audio], outputs=d)
             generate.click(
                 fn=predict,

gradio_components/image.py

@@ -28,6 +28,36 @@ Try to make the prompt simple and concise with only 1-2 sentences
 Make sure the ouput is in JSON fomat, with two items `description` and `prompt`
 """
 
+PROMPT_IMPROVEMENT_GENERATE_PROMPT = """
+You are an export llm prompt enginner, you will be helping the user to improve their prompts. here are some examples of good prompts
+- "90s rock song with electric guitar and heavy drums"
+- "An 80s driving pop song with heavy drums and synth pads in the background"
+- "An energetic hip-hop music piece, with synth sounds and strong bass. There is a rhythmic hi-hat patten in the drums."
+- "A grand orchestral arrangement with thunderous percussion, epic brass fanfares, and soaring strings, creating a cinematic atmosphere fit for a heroic battle."
+- "Classic reggae track with an electronic guitar solo"
+
+You will be provided with a prompt and you need to improve it. Make sure the prompt is simple and concise with only 1-2 sentences. The output should be in JSON format, with one item `prompt`
+"""
+
+
+def improve_prompt(prompt):
+    message = client.messages.create(
+        model="claude-3-opus-20240229",
+        max_tokens=1024,
+        system=PROMPT_IMPROVEMENT_GENERATE_PROMPT,
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ],
+    )
+    message_object = json.loads(message.content[0].text)
+    prompt = message_object["prompt"]
+    return message_object, prompt
+
 
 def generate_caption(image_file, model_file, progress=gr.Progress()):
     if model_file == "facebook/audiogen-medium":

gradio_components/prediction.py

@@ -8,7 +8,7 @@ import gradio as gr
 import torch
 from audiocraft.data.audio import audio_write
 from audiocraft.data.audio_utils import convert_audio
-from audiocraft.models import MusicGen, AudioGen
+from audiocraft.models import AudioGen, MusicGen
 from basic_pitch import ICASSP_2022_MODEL_PATH
 from transformers import AutoModelForSeq2SeqLM
 
@@ -69,10 +69,7 @@ def _do_predictions(
         else:
             if model_file == "facebook/audiogen-medium":
                 # audio condition
-                outputs = model.generate(
-                    texts,
-                    progress=progress
-                )
+                outputs = model.generate(texts, progress=progress)
             else:
                 # text only
                 outputs = model.generate(texts, progress=progress)