From 45d41563dc5a97907accd3a440323d5e93477859 Mon Sep 17 00:00:00 2001
From: Yixiao Zhang <ldzhangyx@outlook.com>
Date: Wed, 28 Jun 2023 17:24:55 +0900
Subject: [PATCH] update accompaniment

---
 melodytalk/modules.py | 35 +++++++++++++++++++++++++++++++++++
 melodytalk/utils.py   | 13 ++++++++-----
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/melodytalk/modules.py b/melodytalk/modules.py
index 66fbf99..e74799f 100644
--- a/melodytalk/modules.py
+++ b/melodytalk/modules.py
@@ -49,6 +49,7 @@ def __init__(self, device):
     @prompts(
         name="Generate music from user input text with melody or track condition",
         description="useful if you want to generate, style transfer or remix music from a user input text with a given melody or track condition."
+                    "Unlike Accompaniment, this tool will also re-generate the given melody."
                     "like: remix the given melody with text description, or doing style transfer as text described with the given melody."
                     "The input to this tool should be a comma separated string of two, "
                     "representing the music_filename and the text description."
@@ -154,4 +155,38 @@ def inference(self, inputs):
         audio_write(updated_music_filename[:-4],
                     wav.cpu(), sr_1, strategy="loudness", loudness_compressor=True)
         print(f"\nProcessed TracksMixing, Output Music: {updated_music_filename}.")
+        return updated_music_filename
+
+
+class Accompaniment(object):
+    template_model = True
+    def __init__(self, device, Text2MusicWithMelody, ExtractTrack, SimpleTracksMixing):
+        print("Initializing Accompaniment")
+        self.device = device
+        self.Text2MusicWithMelody = Text2MusicWithMelody
+        self.ExtractTrack = ExtractTrack
+        self.SimpleTracksMixing = SimpleTracksMixing
+
+    @prompts(
+        name="Generate accompaniment music from user input text, keeping the given melody or track",
+        description="useful if you want to style transfer or remix music from a user input text with a given melody."
+                    "Unlike Text2MusicWithMelody, this tool will keep the given melody track instead of re-generate it."
+                    "Note that the user must assign a track (it must be one of `vocals`, `drums`, `bass`, `guitar`, `piano` or `other`) to keep."
+                    "like: keep the guitar track and remix the given music with text description, "
+                    "or generate accompaniment as text described with the given vocal track."
+                    "The input to this tool should be a comma separated string of three, "
+                    "representing the music_filename, track name, and the text description."
+    )
+
+    def inference(self, inputs):
+        music_filename, track_name, text = inputs.split(",")[0].strip(), inputs.split(",")[1].strip(), inputs.split(",")[2].strip()
+        print(f"Generating music from text with accompaniment condition, Input Text: {text}, Previous music: {music_filename}, Track: {track_name}.")
+        # separate the track
+        updated_main_track = self.ExtractTrack.inference(f"{music_filename}, {track_name}, extract")
+        # generate music
+        updated_new_music = self.Text2MusicWithMelody.inference(f"{text}, {updated_main_track}")
+        # remove the track in accompaniment
+        updated_accompaniment = self.ExtractTrack.inference(f"{updated_new_music}, {track_name}, remove")
+        # mix
+        updated_music_filename = self.SimpleTracksMixing.inference(f"{updated_main_track}, {updated_accompaniment}")
         return updated_music_filename
\ No newline at end of file
diff --git a/melodytalk/utils.py b/melodytalk/utils.py
index d4a1324..eb0c47b 100644
--- a/melodytalk/utils.py
+++ b/melodytalk/utils.py
@@ -4,6 +4,7 @@
 import numpy as np
 import os
 import openai
+import typing as tp
 
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
@@ -38,7 +39,7 @@ def decorator(func):
 
     return decorator
 
-def get_new_audio_name(org_audio_name, func_name="update"):
+def get_new_audio_name(org_audio_name: str, func_name: str ="update") -> str:
     head_tail = os.path.split(org_audio_name)
     head = head_tail[0]
     tail = head_tail[1]
@@ -86,7 +87,7 @@ def description_to_attributes(description: str) -> str:
     return response.choices[0].text
 
 
-def chord_generation(description: str, chord_num: int = 4) -> str:
+def chord_generation(description: str, chord_num: int = 4) -> tp.List:
     """ This function is a trick to generate chord sequence from the description.
 
     :param description:
@@ -97,7 +98,7 @@ def chord_generation(description: str, chord_num: int = 4) -> str:
     openai_prompt = f"""Please generate a chord sequence consists of according to the text description. Example:
 
     Q: Generate a chord sequence for sad pop song. 4 chords.
-    A:  Dm - Bb - F - C
+    A: Dm - Bb - F - C
     
     Q: {description}. {chord_num} chords.
     A:
@@ -106,7 +107,7 @@ def chord_generation(description: str, chord_num: int = 4) -> str:
     response = openai.Completion.create(
         model="text-davinci-003",
         prompt=openai_prompt,
-        temperature=0,
+        temperature=1,
         max_tokens=100,
         top_p=1,
         frequency_penalty=0.0,
@@ -114,4 +115,6 @@ def chord_generation(description: str, chord_num: int = 4) -> str:
         stop=["\n"]
     )
 
-    return response.choices[0].text
+    chord_list = [i.strip() for i in response.choices[0].text.split('-')]
+
+    return chord_list