Add caption chaining feature, captioning chunks and chaining with openai if input is longer than 10 seconds

app.py

@@ -18,11 +18,17 @@ import torch
 import transformers
 import torchaudio
 
+from openai import OpenAI
+client = OpenAI()
+MODEL       = "gpt-4"
+SLEEP_BETWEEN_CALLS = 1.0
+
 from multi_token.model_utils import MultiTaskType
 from multi_token.training import ModelArguments
 from multi_token.inference import load_trained_lora_model
 from multi_token.data_tools import encode_chat
 
+CHUNK_LENGTH = 10
 
 @dataclass
 class ServeArguments(ModelArguments):
@@ -31,7 +37,6 @@ class ServeArguments(ModelArguments):
     temperature: float = field(default=0.01)
 
 
-# Load arguments and model
 logging.getLogger().setLevel(logging.INFO)
 
 parser = transformers.HfArgumentParser((ServeArguments,))
@@ -45,10 +50,82 @@ model, tokenizer = load_trained_lora_model(
     tasks_config=serve_args.tasks_config
 )
 
+def caption_audio(audio_file):
+    chunk_audio_files = split_audio(audio_file, CHUNK_LENGTH)
+    chunk_captions = []
+    for audio_chunk in chunk_audio_files:
+        chunk_captions.append(generate_caption(audio_chunk))
+
+    if len(chunk_captions) > 1:
+        audio_name = os.path.splitext(os.path.basename(audio_file))[0]
+        long_caption = summarize_song(audio_name, chunk_captions)
+
+        delete_files(chunk_audio_files)
+
+        return long_caption
+
+    else:
+        if len(chunk_captions) == 1:
+            return chunk_captions[0]
+        else:
+            return ""
+
+def summarize_song(song_name, chunks):
+    prompt = f"""
+You are a music critic. Given the following chronological 10‑second chunk descriptions of a single piece, write one flowing, detailed description of the entire song—its structure, instrumentation, and standout moments. Mention transition points in terms of time stamps. If the description of certain chunks does not seem to fit with those for the chunks before and after, treat those as bad descriptions with lower accuracy and do not incorporate the information. Retain concrete musical attributes such as key, chords, tempo.
+
+Chunks for “{song_name} ”:
+"""
+    for i, c in enumerate(chunks, 1):
+        prompt += f"\n {(i - 1)*0} to {i*10} seconds. {c.strip()}"
+    prompt += "\n\nFull song description:"
+
+    resp = client.chat.completions.create(model=MODEL,
+    messages=[
+        {"role": "system", "content": "You are an expert music writer."},
+        {"role": "user",   "content": prompt}
+    ],
+    temperature=0.0,
+    max_tokens=1000)
+    return resp.choices[0].message.content.strip()
+
+def delete_files(file_paths):
+    for path in file_paths:
+        try:
+            if os.path.isfile(path):
+                os.remove(path)
+                print(f"Deleted: {path}")
+            else:
+                print(f"Skipped (not a file or doesn't exist): {path}")
+        except Exception as e:
+            print(f"Error deleting {path}: {e}")
+
+def split_audio(input_path, chunk_length_seconds):
+
+    waveform, sample_rate = torchaudio.load(input_path)
+    num_channels, total_samples = waveform.shape
+    chunk_samples = int(chunk_length_seconds * sample_rate)
+
+    num_chunks = (total_samples + chunk_samples - 1) // chunk_samples
+
+    base, ext = os.path.splitext(input_path)
+    output_paths = []
+
+    if (num_chunks <= 1):
+        return [input_path]
+
+    for i in range(num_chunks):
+        start = i * chunk_samples
+        end = min((i + 1) * chunk_samples, total_samples)
+        chunk_waveform = waveform[:, start:end]
+        
+        output_file = f"{base}_{i+1:03d}{ext}"
+        torchaudio.save(output_file, chunk_waveform, sample_rate)
+        output_paths.append(output_file)
+
+    return output_paths
 
 def generate_caption(audio_file):
-    # waveform, sample_rate = torchaudio.load(audio_file)
-
     req_json = {
         "messages": [
             {"role": "user", "content": "Describe the music. <sound>"}
@@ -79,7 +156,7 @@ def generate_caption(audio_file):
 
 
 demo = gr.Interface(
-    fn=generate_caption,
+    fn=caption_audio,
     inputs=gr.Audio(type="filepath", label="Upload an audio file"),
     outputs=gr.Textbox(label="Generated Caption"),
     title="SonicVerse",
@@ -87,4 +164,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()