Update demos/musicgen_app.py

demos/musicgen_app.py

@@ -30,8 +30,8 @@ torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
 torch.backends.cudnn.allow_tf32 = False
 torch.backends.cudnn.deterministic = False
 torch.backends.cudnn.benchmark = False
-# torch.backends.cuda.preferred_blas_library="cublas"
-# torch.backends.cuda.preferred_linalg_library="cusolver"
+torch.backends.cuda.preferred_blas_library="cublas"
+torch.backends.cuda.preferred_linalg_library="cusolver"
 torch.set_float32_matmul_precision("highest")
 
 class FileCleaner:
@@ -53,6 +53,103 @@ class FileCleaner:
                 
 file_cleaner = FileCleaner()
 
+
+# We'll load a model globally to be shared by our new UI tab
+# This simplifies things by avoiding the Predictor class for this specific feature.
+print("Loading model for segmented generation...")
+MODEL = MusicGen.get_pretrained("facebook/musicgen-melody", device='cuda')
+print("Model loaded.")
+
+@spaces.GPU(duration=90)
+def predict_segment(prompt, seed, current_segment, duration_per_segment):
+    """
+    This is the core prediction function for our new UI tab.
+    It calls the generate_segment method we created earlier.
+    """
+    print(f"Requesting to generate segment {current_segment} with seed {seed}.")
+    # Convert duration to token length for the model
+    # Note: The actual frame rate might differ, but this is a close approximation.
+    max_segment_len = int(duration_per_segment * 50) 
+
+    # If the seed is -1 or None, it means we should start with a random seed.
+    # The generate_segment function handles the logic of creating a new seed for segment 1.
+    seed_to_use = None if seed == -1 else int(seed)
+
+    # Call the method we added to the LMModel
+    # Note the path: MODEL.lm.generate_segment
+    full_codes, updated_seed = MODEL.lm.generate_segment(
+        segment=int(current_segment),
+        seed=seed_to_use,
+        prompt_text=prompt,
+        max_segment_len=max_segment_len,
+        # You can pass other generation params here if you add them to the UI
+        temp=1.0, 
+        top_k=250,
+        cfg_coef=3.0,
+    )
+
+    # Decode the *entire* audio generated so far to a waveform
+    audio_wav = MODEL.compression_model.decode(full_codes.to('cuda'))
+    audio_wav = audio_wav.detach().cpu()
+
+    # Save the waveform to a temporary file for the Gradio Audio component
+    with NamedTemporaryFile("wb", suffix=".wav", delete=False) as file:
+        audio_write(
+            file.name, audio_wav[0], MODEL.sample_rate, strategy="loudness",
+            loudness_headroom_db=16, loudness_compressor=True, add_suffix=False
+        )
+        file_cleaner.add(file.name)
+
+    # We must return values for all the outputs of the button's .click() event:
+    # 1. The audio file path for the gr.Audio component.
+    # 2. The updated seed for the gr.Number component.
+    # 3. The next segment number (current + 1) for the gr.Number state.
+    # 4. A new label for the button to guide the user.
+    return file.name, updated_seed, current_segment + 1, gr.update(value=f"Generate Segment {current_segment + 1}")
+
+
+def reset_segmented_ui():
+    """Resets the UI state for starting a new song."""
+    return None, -1, 1, gr.update(value="Generate Segment 1")
+
+
+def ui_segmented():
+    """Builds the Gradio UI for the new 'Long / Segmented Generation' tab."""
+    with gr.Blocks() as interface:
+        gr.Markdown("## Long / Segmented Generation\nGenerate music longer than the time limit by creating it in segments. The state is saved to a file after each segment.")
+        
+        with gr.Row():
+            prompt = gr.Text(label="Prompt", placeholder="e.g., An epic cinematic orchestral score")
+            # The Seed field is critical for linking segments. We use -1 to signify a random start.
+            seed = gr.Number(value=-1, label="Seed (-1 for random start)")
+        
+        duration = gr.Slider(label="Duration PER Segment (seconds)", minimum=10, maximum=90, value=30, step=5)
+        
+        # This hidden Number component will keep track of which segment we're on
+        current_segment = gr.Number(value=1, visible=False)
+        
+        with gr.Row():
+            run_button = gr.Button("Generate Segment 1")
+            reset_button = gr.Button("Reset & Start New Song")
+        
+        audio_output = gr.Audio(label="Full Generated Music So Far", type='filepath')
+
+        # --- Wire up the buttons ---
+        run_button.click(
+            fn=predict_segment,
+            inputs=[prompt, seed, current_segment, duration],
+            # The outputs update the audio player, the seed box, the hidden segment counter, and the button label
+            outputs=[audio_output, seed, current_segment, run_button]
+        )
+
+        reset_button.click(
+            fn=reset_segmented_ui,
+            inputs=[],
+            outputs=[audio_output, seed, current_segment, run_button]
+        )
+    return interface
+
+    
 def convert_wav_to_mp4(wav_path, output_path=None):
     """Converts a WAV file to a waveform MP4 video using ffmpeg."""
     if output_path is None:
@@ -489,9 +586,16 @@ if __name__ == '__main__':
     if args.share:
         launch_kwargs['share'] = args.share
     logging.basicConfig(level=logging.INFO, stream=sys.stderr)
-    # Added predictor shutdown
-    try:
-        ui_full(launch_kwargs)
-    finally:
-        if _predictor is not None:
-            _predictor.shutdown()
+
+    # Build the two interfaces
+    original_interface = ui_full({})
+    segmented_interface = ui_segmented()
+
+    # Combine them in a tabbed layout
+    tabbed_interface = gr.TabbedInterface(
+        [original_interface, segmented_interface],
+        ["Single Generation", "Long / Segmented Generation"]
+    )
+
+    # Launch the final combined app
+    tabbed_interface.queue().launch(**launch_kwargs)