99

app.py

@@ -187,11 +187,11 @@ except Exception as e:
 GRADIO_QUEUE_MAX_SIZE = int(os.environ.get("GRADIO_QUEUE_MAX_SIZE", "24"))
 GRADIO_DEFAULT_CONCURRENCY = int(os.environ.get("GRADIO_DEFAULT_CONCURRENCY", "1"))
 GPU_CONCURRENCY_LIMIT = int(os.environ.get("GRADIO_GPU_CONCURRENCY", "1"))
-STREAM_MIN_CHUNK_SEC = float(os.environ.get("STREAM_MIN_CHUNK_SEC", "29.76"))
+STREAM_MIN_CHUNK_SEC = float(os.environ.get("STREAM_MIN_CHUNK_SEC", "0"))
 
 
 class ModelManager:
-    def __init__(self, model_path: str):
+    def __init__(self, model_path: str, use_deepspeed_override: Optional[bool] = None):
         import torch
         from heartlib import HeartMuLaGenPipeline, HeartTranscriptorPipeline
 
@@ -200,7 +200,10 @@ class ModelManager:
         self.dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
         self._gen_pipes: Dict[Tuple[str, str, str], "HeartMuLaGenPipeline"] = {}
         self._transcribe_pipe: Optional["HeartTranscriptorPipeline"] = None
-        self.use_deepspeed = os.getenv("USE_DEEPSPEED_INFERENCE", "0").lower() in ("1", "true", "yes")
+        if use_deepspeed_override is None:
+            self.use_deepspeed = os.getenv("USE_DEEPSPEED_INFERENCE", "0").lower() in ("1", "true", "yes")
+        else:
+            self.use_deepspeed = use_deepspeed_override
         self.ds_inference_config = self._make_ds_inference_config()
         self._HeartMuLaGenPipeline = HeartMuLaGenPipeline
         self._HeartTranscriptorPipeline = HeartTranscriptorPipeline
@@ -270,16 +273,17 @@ class ModelManager:
         return self._transcribe_pipe
 
 
-model_manager: Optional[ModelManager] = None
+model_managers: Dict[str, ModelManager] = {}
 
 
-def get_model_manager() -> ModelManager:
-    global model_manager
-    if model_manager is None:
+def get_model_manager(use_acceleration: bool) -> ModelManager:
+    key = "accelerated" if use_acceleration else "original"
+    if key not in model_managers:
         os.makedirs(MODEL_PATH, exist_ok=True)
         download_models_if_needed(MODEL_PATH)
-        model_manager = ModelManager(MODEL_PATH)
-    return model_manager
+        use_deepspeed_override = None if use_acceleration else False
+        model_managers[key] = ModelManager(MODEL_PATH, use_deepspeed_override=use_deepspeed_override)
+    return model_managers[key]
 
 
 def update_tag_string(*args):
@@ -463,9 +467,11 @@ def download_transcriptor_if_needed(ckpt_dir):
     print("")
 
 
-def load_pipeline(model_path, version, codec_version, quant_mode):
+def load_pipeline(model_path, version, codec_version, quant_mode, use_acceleration: bool):
     """Load HeartMuLa pipeline (lazy)"""
-    manager = get_model_manager()
+    if not use_acceleration:
+        quant_mode = "none"
+    manager = get_model_manager(use_acceleration)
     print(f"Using model from {model_path} on {manager.device}...")
     return manager.get_gen_pipeline(version, codec_version, quant_mode)
 
@@ -473,7 +479,7 @@ def load_pipeline(model_path, version, codec_version, quant_mode):
 def load_transcriptor(model_path):
     """Load HeartTranscriptor pipeline"""
     download_transcriptor_if_needed(model_path)
-    manager = get_model_manager()
+    manager = get_model_manager(use_acceleration=True)
     return manager.get_transcriptor()
 
 
@@ -492,6 +498,7 @@ def generate(
     keep_model_loaded,
     offload_mode,
     backend,
+    use_acceleration,
 ):
     """Generate music"""
     import torch
@@ -507,7 +514,7 @@ def generate(
         if backend == "exllama_v2":
             raise gr.Error("ExLlamaV2 backend is not implemented yet.")
 
-        pipe = load_pipeline(MODEL_PATH, version, codec_version, quant_mode)
+        pipe = load_pipeline(MODEL_PATH, version, codec_version, quant_mode, use_acceleration)
         output_path = os.path.join(DATA_DIR, f"gen_{uuid.uuid4().hex}.wav")
 
         with torch.no_grad():
@@ -540,6 +547,72 @@ def generate(
         raise gr.Error(f"Generation error: {str(e)}")
 
 
+def generate_original(
+    lyrics,
+    tags,
+    cfg_scale,
+    duration_sec,
+    temperature,
+    topk,
+    version,
+    codec_version,
+    quant_mode,
+    output_format,
+    keep_model_loaded,
+    offload_mode,
+    backend,
+):
+    return generate(
+        lyrics,
+        tags,
+        cfg_scale,
+        duration_sec,
+        temperature,
+        topk,
+        version,
+        codec_version,
+        quant_mode,
+        output_format,
+        keep_model_loaded,
+        offload_mode,
+        backend,
+        False,
+    )
+
+
+def generate_accelerated(
+    lyrics,
+    tags,
+    cfg_scale,
+    duration_sec,
+    temperature,
+    topk,
+    version,
+    codec_version,
+    quant_mode,
+    output_format,
+    keep_model_loaded,
+    offload_mode,
+    backend,
+):
+    return generate(
+        lyrics,
+        tags,
+        cfg_scale,
+        duration_sec,
+        temperature,
+        topk,
+        version,
+        codec_version,
+        quant_mode,
+        output_format,
+        keep_model_loaded,
+        offload_mode,
+        backend,
+        True,
+    )
+
+
 @_gpu_guard
 def transcribe_audio(audio_path, task, max_new_tokens, num_beams, temperature):
     """Transcribe or translate lyrics from audio"""
@@ -590,10 +663,11 @@ def generate_music_streaming(
     offload_mode,
     backend,
     chunk_frames,
+    use_acceleration,
 ) -> Iterator[Tuple[int, np.ndarray]]:
     if backend == "exllama_v2":
         raise gr.Error("ExLlamaV2 backend is not implemented yet.")
-    pipe = load_pipeline(MODEL_PATH, version, codec_version, quant_mode)
+    pipe = load_pipeline(MODEL_PATH, version, codec_version, quant_mode, use_acceleration)
     max_audio_length_ms = int(duration_sec * 1000)
     for chunk in pipe.stream(
         {"lyrics": lyrics, "tags": tags},
@@ -629,15 +703,16 @@ def stream_generate(
     backend,
     output_format,
     chunk_frames,
+    use_acceleration,
 ):
     try:
-        min_samples = max(1, int(STREAM_MIN_CHUNK_SEC * 48000))
+        min_samples = max(0, int(STREAM_MIN_CHUNK_SEC * 48000))
         buffer = []
         buffered_samples = 0
         last_yield_samples = 0
         print(
             "stream start:",
-            f"min_samples={min_samples}",
+            f"min_chunk_sec={STREAM_MIN_CHUNK_SEC}",
             f"duration_sec={duration_sec}",
             f"chunk_frames={chunk_frames}",
         )
@@ -656,23 +731,22 @@ def stream_generate(
             offload_mode=offload_mode,
             backend=backend,
             chunk_frames=chunk_frames,
+            use_acceleration=use_acceleration,
         ):
             chunk_np = chunk_np.astype("float32", copy=False)
+            if min_samples <= 0:
+                print(f"stream yield: samples={chunk_np.shape[0]}")
+                yield sr, chunk_np
+                continue
             buffer.append(chunk_np)
             buffered_samples += chunk_np.shape[0]
-            print(
-                "stream buffer:",
-                f"chunk={chunk_np.shape[0]}",
-                f"buffered={buffered_samples}",
-            )
             if buffered_samples - last_yield_samples < min_samples:
                 continue
             full_audio = np.concatenate(buffer)
             last_yield_samples = buffered_samples
             print(f"stream yield: samples={full_audio.shape[0]}")
             yield sr, full_audio
-
-        if buffer:
+        if min_samples > 0 and buffer:
             full_audio = np.concatenate(buffer)
             print(f"stream final yield: samples={full_audio.shape[0]}")
             yield 48000, full_audio
@@ -680,6 +754,41 @@ def stream_generate(
         raise gr.Error(f"Streaming error: {str(e)}")
 
 
+def stream_generate_accelerated(
+    lyrics,
+    tags,
+    cfg_scale,
+    duration_sec,
+    temperature,
+    topk,
+    version,
+    codec_version,
+    quant_mode,
+    keep_model_loaded,
+    offload_mode,
+    backend,
+    output_format,
+    chunk_frames,
+):
+    return stream_generate(
+        lyrics,
+        tags,
+        cfg_scale,
+        duration_sec,
+        temperature,
+        topk,
+        version,
+        codec_version,
+        quant_mode,
+        keep_model_loaded,
+        offload_mode,
+        backend,
+        output_format,
+        chunk_frames,
+        True,
+    )
+
+
 def generate_lyrics(theme, tags, language, api_choice, api_key_input, custom_base_url, custom_model, progress=gr.Progress()):
     """Generate lyrics using selected LLM API"""
 
@@ -946,9 +1055,25 @@ def create_ui():
                                 5, 100, value=20, step=1, label="Streaming Chunk Frames"
                             )
 
-                        generate_btn = gr.Button("Generate Music", variant="primary", size="lg")
-                        stream_btn = gr.Button("Generate Music (Streaming)", variant="primary", size="lg")
-                        cancel_stream_btn = gr.Button("Cancel Streaming", variant="secondary", size="lg")
+                        gr.Markdown("### 🚀 Generation")
+
+                        generation_mode = gr.Radio(
+                            choices=["Original (No Acceleration)", "Accelerated"],
+                            value="Original (No Acceleration)",
+                            label="Generation Mode",
+                        )
+
+                        speed_submode = gr.Radio(
+                            choices=["Standard", "Streaming"],
+                            value="Standard",
+                            label="Accelerated Options",
+                            visible=False,
+                        )
+
+                        btn_original = gr.Button("🎼 Generate Music (Original)", variant="primary", size="lg", visible=True)
+                        btn_accel = gr.Button("🎼 Generate Music (Accelerated)", variant="primary", size="lg", visible=False)
+                        btn_stream = gr.Button("🎼 Generate Music (Streaming)", variant="primary", size="lg", visible=False)
+                        cancel_stream_btn = gr.Button("Cancel Streaming", variant="secondary", size="lg", visible=False)
                         cancel_state = gr.State()
 
                     with gr.Column():
@@ -1015,8 +1140,59 @@ Every day the fire burns
                     outputs=[lyrics]
                 )
 
-                generate_btn.click(
-                    fn=generate,
+                def update_visibility(gen_mode, spd_mode):
+                    if gen_mode == "Original (No Acceleration)":
+                        return (
+                            gr.update(visible=False),  # speed_submode
+                            gr.update(visible=True),  # btn_original
+                            gr.update(visible=False),  # btn_accel
+                            gr.update(visible=False),  # btn_stream
+                            gr.update(visible=False),  # cancel_stream_btn
+                        )
+                    show_stream = spd_mode == "Streaming"
+                    return (
+                        gr.update(visible=True),  # speed_submode
+                        gr.update(visible=False),  # btn_original
+                        gr.update(visible=not show_stream),  # btn_accel
+                        gr.update(visible=show_stream),  # btn_stream
+                        gr.update(visible=show_stream),  # cancel_stream_btn
+                    )
+
+                generation_mode.change(
+                    fn=update_visibility,
+                    inputs=[generation_mode, speed_submode],
+                    outputs=[speed_submode, btn_original, btn_accel, btn_stream, cancel_stream_btn],
+                )
+                speed_submode.change(
+                    fn=update_visibility,
+                    inputs=[generation_mode, speed_submode],
+                    outputs=[speed_submode, btn_original, btn_accel, btn_stream, cancel_stream_btn],
+                )
+
+                btn_original.click(
+                    fn=generate_original,
+                    inputs=[
+                        lyrics,
+                        tags,
+                        cfg_scale,
+                        duration,
+                        temperature,
+                        topk,
+                        version,
+                        codec_version,
+                        quant_mode,
+                        output_format,
+                        keep_model_loaded,
+                        offload_mode,
+                        backend,
+                    ],
+                    outputs=[output_audio_file],
+                    concurrency_id="gpu_queue",
+                    concurrency_limit=GPU_CONCURRENCY_LIMIT,
+                )
+
+                btn_accel.click(
+                    fn=generate_accelerated,
                     inputs=[
                         lyrics,
                         tags,
@@ -1037,8 +1213,8 @@ Every day the fire burns
                     concurrency_limit=GPU_CONCURRENCY_LIMIT,
                 )
 
-                stream_event = stream_btn.click(
-                    fn=stream_generate,
+                stream_event = btn_stream.click(
+                    fn=stream_generate_accelerated,
                     inputs=[
                         lyrics,
                         tags,