Update app.py

app.py

@@ -113,26 +113,26 @@ instruction = f"""
 """
 
 @spaces.GPU(enable_queue=True, duration=60)
-def infer(image_in):
-    gr.Info("Getting image caption with Kosmos2...")
-    user_prompt = get_caption(image_in)
-    
+def get_musical_prompt(user_prompt):
     prompt = f"{instruction.strip()}\n{user_prompt}</s>"    
-    #print(f"PROMPT: {prompt}")
-    
-    gr.Info("Building a musical prompt according to the image caption ...")
     outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-    
-
     pattern = r'\<\|system\|\>(.*?)\<\|assistant\|\>'
     cleaned_text = re.sub(pattern, '', outputs[0]["generated_text"], flags=re.DOTALL)
     
     print(f"SUGGESTED Musical prompt: {cleaned_text}")
+    return cleaned_text.lstrip("\n")
+
+def infer(image_in):
+    gr.Info("Getting image caption with Kosmos2...")
+    user_prompt = get_caption(image_in)
+    
+    gr.Info("Building a musical prompt according to the image caption ...")
+    musical_prompt = get_musical_prompt(user_prompt)
 
     gr.Info("Now calling MAGNet for music ...")
-    music_o = get_magnet(cleaned_text)
+    music_o = get_magnet(musical_prompt)
     
-    return cleaned_text.lstrip("\n"), music_o
+    return cleaned_text, music_o
 
 demo_title = "Image to Music V2"
 description = "Get music from a picture"