no more duration params

app.py

@@ -18,159 +18,155 @@ def preprocess_audio(waveform):
     waveform_np = waveform.cpu().squeeze().numpy()
     return torch.from_numpy(waveform_np).unsqueeze(0).to(device)
 
-@spaces.GPU
-def simple_test(text):
-    return f"Hello {text}!"
-
-# @spaces.GPU(duration=10)
-# def generate_drum_sample():
-#     model = MusicGen.get_pretrained('pharoAIsanders420/micro-musicgen-jungle')
-#     model.set_generation_params(duration=10)
-#     wav = model.generate_unconditional(1).squeeze(0)
+@spaces.GPU  # Remove duration parameter
+def generate_drum_sample():
+    model = MusicGen.get_pretrained('pharoAIsanders420/micro-musicgen-jungle')
+    model.set_generation_params(duration=10)
+    wav = model.generate_unconditional(1).squeeze(0)
     
-#     filename_without_extension = f'jungle'
-#     filename_with_extension = f'{filename_without_extension}.wav'
+    filename_without_extension = f'jungle'
+    filename_with_extension = f'{filename_without_extension}.wav'
 
-#     audio_write(filename_without_extension, wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
+    audio_write(filename_without_extension, wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
 
-#     return filename_with_extension
+    return filename_with_extension
 
-# @spaces.GPU(duration=10)
-# def continue_drum_sample(existing_audio_path):
-#     if existing_audio_path is None:
-#         return None
+@spaces.GPU  # Remove duration parameter
+def continue_drum_sample(existing_audio_path):
+    if existing_audio_path is None:
+        return None
         
-#     existing_audio, sr = torchaudio.load(existing_audio_path)
-#     existing_audio = existing_audio.to(device)
+    existing_audio, sr = torchaudio.load(existing_audio_path)
+    existing_audio = existing_audio.to(device)
     
-#     prompt_duration = 2
-#     output_duration = 10
+    prompt_duration = 2
+    output_duration = 10
 
-#     num_samples = int(prompt_duration * sr)
-#     if existing_audio.shape[1] < num_samples:
-#         raise ValueError("The existing audio is too short for the specified prompt duration.")
+    num_samples = int(prompt_duration * sr)
+    if existing_audio.shape[1] < num_samples:
+        raise ValueError("The existing audio is too short for the specified prompt duration.")
 
-#     start_sample = existing_audio.shape[1] - num_samples
-#     prompt_waveform = existing_audio[..., start_sample:]
+    start_sample = existing_audio.shape[1] - num_samples
+    prompt_waveform = existing_audio[..., start_sample:]
 
-#     model = MusicGen.get_pretrained('pharoAIsanders420/micro-musicgen-jungle')
-#     model.set_generation_params(duration=output_duration)
+    model = MusicGen.get_pretrained('pharoAIsanders420/micro-musicgen-jungle')
+    model.set_generation_params(duration=output_duration)
 
-#     output = model.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
-#     output = output.to(device)
+    output = model.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
+    output = output.to(device)
 
-#     if output.dim() == 3:
-#         output = output.squeeze(0)
+    if output.dim() == 3:
+        output = output.squeeze(0)
 
-#     if output.dim() == 1:
-#         output = output.unsqueeze(0)
+    if output.dim() == 1:
+        output = output.unsqueeze(0)
 
-#     combined_audio = torch.cat((existing_audio, output), dim=1)
-#     combined_audio = combined_audio.cpu()
+    combined_audio = torch.cat((existing_audio, output), dim=1)
+    combined_audio = combined_audio.cpu()
 
-#     combined_file_path = f'./continued_jungle_{random.randint(1000, 9999)}.wav'
-#     torchaudio.save(combined_file_path, combined_audio, sr)
+    combined_file_path = f'./continued_jungle_{random.randint(1000, 9999)}.wav'
+    torchaudio.save(combined_file_path, combined_audio, sr)
 
-#     return combined_file_path
+    return combined_file_path
 
-# @spaces.GPU(duration=120)
-# def generate_music(wav_filename, prompt_duration, musicgen_model, output_duration):
-#     if wav_filename is None:
-#         return None
+@spaces.GPU  # Remove duration parameter
+def generate_music(wav_filename, prompt_duration, musicgen_model, output_duration):
+    if wav_filename is None:
+        return None
         
-#     song, sr = torchaudio.load(wav_filename)
-#     song = song.to(device)
-
-#     model_name = musicgen_model.split(" ")[0]
-#     model_continue = MusicGen.get_pretrained(model_name)
-
-#     model_continue.set_generation_params(
-#         use_sampling=True,
-#         top_k=250,
-#         top_p=0.0,
-#         temperature=1.0,
-#         duration=output_duration,
-#         cfg_coef=3
-#     )
-
-#     prompt_waveform = song[..., :int(prompt_duration * sr)]
-#     prompt_waveform = preprocess_audio(prompt_waveform)
+    song, sr = torchaudio.load(wav_filename)
+    song = song.to(device)
+
+    model_name = musicgen_model.split(" ")[0]
+    model_continue = MusicGen.get_pretrained(model_name)
+
+    model_continue.set_generation_params(
+        use_sampling=True,
+        top_k=250,
+        top_p=0.0,
+        temperature=1.0,
+        duration=output_duration,
+        cfg_coef=3
+    )
+
+    prompt_waveform = song[..., :int(prompt_duration * sr)]
+    prompt_waveform = preprocess_audio(prompt_waveform)
     
-#     output = model_continue.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
-#     output = output.cpu()
+    output = model_continue.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
+    output = output.cpu()
     
-#     if len(output.size()) > 2:
-#         output = output.squeeze()
+    if len(output.size()) > 2:
+        output = output.squeeze()
     
-#     filename_without_extension = f'continued_music'
-#     filename_with_extension = f'{filename_without_extension}.wav'
-#     audio_write(filename_without_extension, output, model_continue.sample_rate, strategy="loudness", loudness_compressor=True)
+    filename_without_extension = f'continued_music'
+    filename_with_extension = f'{filename_without_extension}.wav'
+    audio_write(filename_without_extension, output, model_continue.sample_rate, strategy="loudness", loudness_compressor=True)
 
-#     return filename_with_extension
+    return filename_with_extension
 
-# @spaces.GPU(duration=120)
-# def continue_music(input_audio_path, prompt_duration, musicgen_model, output_duration):
-#     if input_audio_path is None:
-#         return None
+@spaces.GPU  # Remove duration parameter
+def continue_music(input_audio_path, prompt_duration, musicgen_model, output_duration):
+    if input_audio_path is None:
+        return None
         
-#     song, sr = torchaudio.load(input_audio_path)
-#     song = song.to(device)
+    song, sr = torchaudio.load(input_audio_path)
+    song = song.to(device)
 
-#     model_continue = MusicGen.get_pretrained(musicgen_model.split(" ")[0])
-#     model_continue.set_generation_params(
-#         use_sampling=True,
-#         top_k=250,
-#         top_p=0.0,
-#         temperature=1.0,
-#         duration=output_duration,
-#         cfg_coef=3
-#     )
+    model_continue = MusicGen.get_pretrained(musicgen_model.split(" ")[0])
+    model_continue.set_generation_params(
+        use_sampling=True,
+        top_k=250,
+        top_p=0.0,
+        temperature=1.0,
+        duration=output_duration,
+        cfg_coef=3
+    )
 
-#     original_audio = AudioSegment.from_mp3(input_audio_path)
-#     current_audio = original_audio
+    original_audio = AudioSegment.from_mp3(input_audio_path)
+    current_audio = original_audio
 
-#     file_paths_for_cleanup = []
+    file_paths_for_cleanup = []
 
-#     for i in range(1):
-#         num_samples = int(prompt_duration * sr)
-#         if current_audio.duration_seconds * 1000 < prompt_duration * 1000:
-#             raise ValueError("The prompt_duration is longer than the current audio length.")
+    for i in range(1):
+        num_samples = int(prompt_duration * sr)
+        if current_audio.duration_seconds * 1000 < prompt_duration * 1000:
+            raise ValueError("The prompt_duration is longer than the current audio length.")
 
-#         start_time = current_audio.duration_seconds * 1000 - prompt_duration * 1000
-#         prompt_audio = current_audio[start_time:]
+        start_time = current_audio.duration_seconds * 1000 - prompt_duration * 1000
+        prompt_audio = current_audio[start_time:]
 
-#         prompt_bytes = prompt_audio.export(format="wav").read()
-#         prompt_waveform, _ = torchaudio.load(io.BytesIO(prompt_bytes))
-#         prompt_waveform = prompt_waveform.to(device)
+        prompt_bytes = prompt_audio.export(format="wav").read()
+        prompt_waveform, _ = torchaudio.load(io.BytesIO(prompt_bytes))
+        prompt_waveform = prompt_waveform.to(device)
 
-#         prompt_waveform = preprocess_audio(prompt_waveform)
+        prompt_waveform = preprocess_audio(prompt_waveform)
 
-#         output = model_continue.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
-#         output = output.cpu()
+        output = model_continue.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
+        output = output.cpu()
 
-#         if len(output.size()) > 2:
-#             output = output.squeeze()
+        if len(output.size()) > 2:
+            output = output.squeeze()
 
-#         filename_without_extension = f'continue_{i}'
-#         filename_with_extension = f'{filename_without_extension}.wav'
-#         correct_filename_extension = f'{filename_without_extension}.wav.wav'
+        filename_without_extension = f'continue_{i}'
+        filename_with_extension = f'{filename_without_extension}.wav'
+        correct_filename_extension = f'{filename_without_extension}.wav.wav'
 
-#         audio_write(filename_with_extension, output, model_continue.sample_rate, strategy="loudness", loudness_compressor=True)
-#         generated_audio_segment = AudioSegment.from_wav(correct_filename_extension)
+        audio_write(filename_with_extension, output, model_continue.sample_rate, strategy="loudness", loudness_compressor=True)
+        generated_audio_segment = AudioSegment.from_wav(correct_filename_extension)
 
-#         current_audio = current_audio[:start_time] + generated_audio_segment
+        current_audio = current_audio[:start_time] + generated_audio_segment
 
-#         file_paths_for_cleanup.append(correct_filename_extension)
+        file_paths_for_cleanup.append(correct_filename_extension)
 
-#     combined_audio_filename = f"combined_audio_{random.randint(1, 10000)}.mp3"
-#     current_audio.export(combined_audio_filename, format="mp3")
+    combined_audio_filename = f"combined_audio_{random.randint(1, 10000)}.mp3"
+    current_audio.export(combined_audio_filename, format="mp3")
 
-#     for file_path in file_paths_for_cleanup:
-#         os.remove(file_path)
+    for file_path in file_paths_for_cleanup:
+        os.remove(file_path)
 
-#     return combined_audio_filename
+    return combined_audio_filename
 
-# Define the expandable sections (keeping your existing content)
+# Define the expandable sections
 musicgen_micro_blurb = """
 ## musicgen_micro
 musicgen micro is an experimental series of models by aaron abebe. they are incredibly fast, and extra insane. this one does goated jungle drums. we're very excited about these.
@@ -213,7 +209,7 @@ thepatch/PhonkV2 was trained by MJ BERSABEph. there are multiple versions in the
 foureyednymph/musicgen-sza-sos-small was just trained by foureyednymph. We're all about to find out if it does continuations well.
 """
 
-# Create the Gradio interface with explicit types
+# Create the Gradio interface
 with gr.Blocks() as iface:
     gr.Markdown("# the-micro-slot-machine")
     gr.Markdown("two ai's jamming. warning: outputs will be very strange, likely stupid, and possibly rad.")
@@ -269,10 +265,10 @@ with gr.Blocks() as iface:
             continue_output_audio = gr.Audio(label="Continued Music Output", type="filepath")
 
     # Connecting the components
-    # generate_button.click(generate_drum_sample, outputs=[drum_audio])
-    # continue_drum_sample_button.click(continue_drum_sample, inputs=[drum_audio], outputs=[drum_audio])
-    # generate_music_button.click(generate_music, inputs=[drum_audio, prompt_duration, musicgen_model, output_duration], outputs=[output_audio])
-    # continue_button.click(continue_music, inputs=[output_audio, prompt_duration, musicgen_model, output_duration], outputs=continue_output_audio)
+    generate_button.click(generate_drum_sample, outputs=[drum_audio])
+    continue_drum_sample_button.click(continue_drum_sample, inputs=[drum_audio], outputs=[drum_audio])
+    generate_music_button.click(generate_music, inputs=[drum_audio, prompt_duration, musicgen_model, output_duration], outputs=[output_audio])
+    continue_button.click(continue_music, inputs=[output_audio, prompt_duration, musicgen_model, output_duration], outputs=continue_output_audio)
 
 if __name__ == "__main__":
     iface.launch()