22

app.py

@@ -17,7 +17,6 @@ import uuid
 import subprocess
 import tempfile
 import wave
-import time
 from pathlib import Path
 from typing import Dict, Any, Iterator, Optional, Tuple, TYPE_CHECKING
 
@@ -29,7 +28,15 @@ try:
 except Exception:
     spaces = None
 
-GPU_MAX_DURATION = int(os.environ.get("GPU_MAX_DURATION", "100"))
+IS_SPACE = spaces is not None and os.environ.get("SPACE_ID")
+
+_gpu_duration_env = os.environ.get("GPU_MAX_DURATION")
+if _gpu_duration_env is None:
+    aoti_env = os.environ.get("ENABLE_AOTI")
+    enable_aoti_default = aoti_env is None or aoti_env.strip().lower() in ("1", "true", "yes", "y", "on")
+    GPU_MAX_DURATION = 600 if IS_SPACE and enable_aoti_default else 100
+else:
+    GPU_MAX_DURATION = int(_gpu_duration_env)
 
 
 def _env_bool(name: str) -> Optional[bool]:
@@ -41,7 +48,7 @@ def _env_bool(name: str) -> Optional[bool]:
 
 _default_keep_model_loaded_env = _env_bool("KEEP_MODEL_LOADED_DEFAULT")
 if _default_keep_model_loaded_env is None:
-    DEFAULT_KEEP_MODEL_LOADED = not (spaces is not None and os.environ.get("SPACE_ID"))
+    DEFAULT_KEEP_MODEL_LOADED = not IS_SPACE
 else:
     DEFAULT_KEEP_MODEL_LOADED = _default_keep_model_loaded_env
 
@@ -59,9 +66,11 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
 
 # HF Spaces: disable SSR to avoid Node proxy/port issues.
 os.environ.setdefault("GRADIO_SSR_MODE", "0")
+# Mitigate CUDA memory fragmentation on small GPUs.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 if TYPE_CHECKING:
-    from heartlib import HeartMuLaGenPipeline, HeartTranscriptorPipeline
+    from heartlib import HeartMuLaGenPipeline
 import gradio as gr
 import re
 from google import genai
@@ -209,26 +218,100 @@ 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", "0"))
+DEFAULT_DURATION_SEC = int(os.environ.get("DEFAULT_DURATION_SEC", "60" if IS_SPACE else "180"))
+DEFAULT_QUANT_MODE = os.environ.get("DEFAULT_QUANT_MODE", "4bit" if IS_SPACE else "none")
+DEFAULT_OFFLOAD_MODE = os.environ.get("DEFAULT_OFFLOAD_MODE", "aggressive" if IS_SPACE else "auto")
+DEFAULT_GENERATION_MODE = os.environ.get(
+    "DEFAULT_GENERATION_MODE",
+    "Accelerated" if IS_SPACE else "Original (No Acceleration)",
+)
+DEFAULT_SPEED_SUBMODE = os.environ.get("DEFAULT_SPEED_SUBMODE", "Standard")
+DEFAULT_PRESET = os.environ.get("DEFAULT_PRESET", "ZeroGPU Safe" if IS_SPACE else "Balanced")
+MODEL_DETOKENIZE_INTERVAL_SEC = float(os.environ.get("MODEL_DETOKENIZE_INTERVAL_SEC", "29.76"))
+AUDIO_SAMPLE_RATE = int(os.environ.get("AUDIO_SAMPLE_RATE", "48000"))
+FRAME_MS = 80.0
+BLOCK_FRAMES = max(1, int(round((MODEL_DETOKENIZE_INTERVAL_SEC * 1000.0) / FRAME_MS)))
+BLOCK_SAMPLES = max(1, int(round(MODEL_DETOKENIZE_INTERVAL_SEC * AUDIO_SAMPLE_RATE)))
+MIN_BUFFER_BLOCKS = int(os.environ.get("MIN_BUFFER_BLOCKS", "1"))
+PREFETCH_BLOCKS = int(os.environ.get("PREFETCH_BLOCKS", "2"))
+MAX_QUEUE_BLOCKS = int(os.environ.get("MAX_QUEUE_BLOCKS", "4"))
+
+PRESET_CONFIGS = {
+    "ZeroGPU Safe": {
+        "duration": 60,
+        "quant_mode": "4bit",
+        "offload_mode": "aggressive",
+        "keep_model_loaded": False,
+        "temperature": 1.0,
+        "topk": 50,
+        "cfg_scale": 1.5,
+        "generation_mode": "Accelerated",
+        "speed_submode": "Standard",
+    },
+    "ZeroGPU AOTI FP8": {
+        "duration": 60,
+        "quant_mode": "fp8",
+        "offload_mode": "aggressive",
+        "keep_model_loaded": False,
+        "temperature": 1.0,
+        "topk": 50,
+        "cfg_scale": 1.5,
+        "generation_mode": "Accelerated",
+        "speed_submode": "Standard",
+    },
+    "Balanced": {
+        "duration": 120,
+        "quant_mode": "4bit" if IS_SPACE else "none",
+        "offload_mode": "auto",
+        "keep_model_loaded": DEFAULT_KEEP_MODEL_LOADED,
+        "temperature": 1.0,
+        "topk": 50,
+        "cfg_scale": 1.5,
+        "generation_mode": "Accelerated",
+        "speed_submode": "Standard",
+    },
+    "Quality": {
+        "duration": 180,
+        "quant_mode": "none",
+        "offload_mode": "auto",
+        "keep_model_loaded": DEFAULT_KEEP_MODEL_LOADED,
+        "temperature": 0.9,
+        "topk": 50,
+        "cfg_scale": 2.0,
+        "generation_mode": "Original (No Acceleration)",
+        "speed_submode": "Standard",
+    },
+}
+if DEFAULT_PRESET not in PRESET_CONFIGS:
+    DEFAULT_PRESET = "Balanced"
+_default_preset_config = PRESET_CONFIGS[DEFAULT_PRESET]
+if "DEFAULT_DURATION_SEC" not in os.environ:
+    DEFAULT_DURATION_SEC = _default_preset_config["duration"]
+if "DEFAULT_QUANT_MODE" not in os.environ:
+    DEFAULT_QUANT_MODE = _default_preset_config["quant_mode"]
+if "DEFAULT_OFFLOAD_MODE" not in os.environ:
+    DEFAULT_OFFLOAD_MODE = _default_preset_config["offload_mode"]
+if "DEFAULT_GENERATION_MODE" not in os.environ:
+    DEFAULT_GENERATION_MODE = _default_preset_config["generation_mode"]
+if "DEFAULT_SPEED_SUBMODE" not in os.environ:
+    DEFAULT_SPEED_SUBMODE = _default_preset_config["speed_submode"]
 
 
 class ModelManager:
     def __init__(self, model_path: str, use_deepspeed_override: Optional[bool] = None):
         import torch
-        from heartlib import HeartMuLaGenPipeline, HeartTranscriptorPipeline
+        from heartlib import HeartMuLaGenPipeline
 
         self.model_path = model_path
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         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
         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
 
     def _make_ds_inference_config(self) -> Dict[str, Any]:
         if not self.use_deepspeed:
@@ -249,6 +332,8 @@ class ModelManager:
             return None
         if self.device.type != "cuda":
             raise gr.Error("Quantization requires CUDA.")
+        if quant_mode == "fp8":
+            return None
         if quant_mode == "4bit":
             quant_type = "nf4"
             try:
@@ -278,22 +363,13 @@ class ModelManager:
                 version=version,
                 codec_version=codec_version,
                 bnb_config=bnb_config,
+                torchao_quantize=quant_mode == "fp8",
                 lazy_load=True,
                 use_deepspeed=self.use_deepspeed,
                 ds_inference_config=self.ds_inference_config,
             )
         return self._gen_pipes[key]
 
-    def get_transcriptor(self) -> "HeartTranscriptorPipeline":
-        import torch
-        if self._transcribe_pipe is None:
-            self._transcribe_pipe = self._HeartTranscriptorPipeline.from_pretrained(
-                self.model_path,
-                device=self.device,
-                dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
-            )
-        return self._transcribe_pipe
-
 
 model_managers: Dict[str, ModelManager] = {}
 
@@ -459,50 +535,18 @@ def download_models_if_needed(ckpt_dir):
     print("")
 
 
-def check_transcriptor_exists(ckpt_dir):
-    """Check if HeartTranscriptor checkpoint exists"""
-    return os.path.exists(os.path.join(ckpt_dir, "HeartTranscriptor-oss"))
-
-
-def download_transcriptor_if_needed(ckpt_dir):
-    """Download HeartTranscriptor from ModelScope if not present"""
-    if check_transcriptor_exists(ckpt_dir):
-        print("=" * 50)
-        print(f"✓ HeartTranscriptor found in {ckpt_dir}")
-        print("✓ Skipping HeartTranscriptor download")
-        print("=" * 50)
-        return
-
-    print("=" * 50)
-    print("⬇  Downloading HeartTranscriptor from ModelScope")
-    print("=" * 50)
-    print("")
-
-    from modelscope import snapshot_download
-
-    snapshot_download(
-        'HeartMuLa/HeartTranscriptor-oss',
-        local_dir=os.path.join(ckpt_dir, 'HeartTranscriptor-oss')
-    )
-
-    print("✓ HeartTranscriptor download completed")
-    print("")
-
-
 def load_pipeline(model_path, version, codec_version, quant_mode, use_acceleration: bool):
     """Load HeartMuLa pipeline (lazy)"""
-    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)
 
 
-def load_transcriptor(model_path):
-    """Load HeartTranscriptor pipeline"""
-    download_transcriptor_if_needed(model_path)
-    manager = get_model_manager(use_acceleration=True)
-    return manager.get_transcriptor()
+def _validate_generation_inputs(lyrics: str, tags: str) -> None:
+    if not lyrics.strip():
+        raise gr.Error("Please enter lyrics")
+    if not tags.strip():
+        raise gr.Error("Please enter tags")
 
 
 def generate(
@@ -523,11 +567,7 @@ def generate(
 ):
     """Generate music"""
     import torch
-    if not lyrics.strip():
-        raise gr.Error("Please enter lyrics")
-
-    if not tags.strip():
-        raise gr.Error("Please enter tags")
+    _validate_generation_inputs(lyrics, tags)
 
     max_audio_length_ms = int(duration_sec * 1000)
 
@@ -538,7 +578,7 @@ def generate(
         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():
+        with torch.inference_mode():
             pipe(
                 {
                     "lyrics": lyrics,
@@ -566,6 +606,9 @@ def generate(
 
     except Exception as e:
         raise gr.Error(f"Generation error: {str(e)}")
+    finally:
+        if not keep_model_loaded and torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 
 @_gpu_guard
@@ -636,37 +679,6 @@ def generate_accelerated(
     )
 
 
-@_gpu_guard
-def transcribe_audio(audio_path, task, max_new_tokens, num_beams, temperature):
-    """Transcribe or translate lyrics from audio"""
-    import torch
-    if not audio_path:
-        raise gr.Error("Please upload an audio file")
-
-    pipe = load_transcriptor(MODEL_PATH)
-
-    try:
-        with torch.no_grad():
-            result = pipe(
-                audio_path,
-                **{
-                    "max_new_tokens": int(max_new_tokens),
-                    "num_beams": int(num_beams),
-                    "task": task,
-                    "condition_on_prev_tokens": False,
-                    "compression_ratio_threshold": 1.8,
-                    "temperature": float(temperature),
-                    "logprob_threshold": -1.0,
-                    "no_speech_threshold": 0.4,
-                },
-            )
-        if isinstance(result, dict):
-            return result.get("text", "")
-        return str(result)
-    except Exception as e:
-        raise gr.Error(f"Transcription error: {str(e)}")
-
-
 def _normalize_stream_chunk(chunk: np.ndarray) -> np.ndarray:
     chunk = np.nan_to_num(chunk, nan=0.0, posinf=0.0, neginf=0.0)
     return np.clip(chunk, -1.0, 1.0)
@@ -685,9 +697,9 @@ def generate_music_streaming(
     keep_model_loaded,
     offload_mode,
     backend,
-    chunk_frames,
     use_acceleration,
 ) -> Iterator[Tuple[int, np.ndarray]]:
+    _validate_generation_inputs(lyrics, tags)
     if backend == "exllama_v2":
         raise gr.Error("ExLlamaV2 backend is not implemented yet.")
     pipe = load_pipeline(MODEL_PATH, version, codec_version, quant_mode, use_acceleration)
@@ -698,7 +710,7 @@ def generate_music_streaming(
         temperature=temperature,
         topk=topk,
         cfg_scale=cfg_scale,
-        chunk_frames=int(chunk_frames),
+        chunk_frames=BLOCK_FRAMES,
         keep_model_loaded=keep_model_loaded,
         offload_mode=offload_mode,
     ):
@@ -706,8 +718,8 @@ def generate_music_streaming(
             chunk = chunk.squeeze(0)
         chunk_np = chunk.cpu().numpy()
         chunk_np = _normalize_stream_chunk(chunk_np)
-        print(f"stream chunk: samples={chunk_np.shape[0]} sr=48000")
-        yield 48000, chunk_np
+        print(f"stream chunk: samples={chunk_np.shape[0]} sr={AUDIO_SAMPLE_RATE}")
+        yield AUDIO_SAMPLE_RATE, chunk_np
 
 
 def stream_generate(
@@ -723,22 +735,14 @@ def stream_generate(
     keep_model_loaded,
     offload_mode,
     backend,
-    output_format,
-    chunk_frames,
     use_acceleration,
 ):
     try:
-        min_samples = max(0, int(STREAM_MIN_CHUNK_SEC * 48000))
-        buffer = []
-        buffered_samples = 0
-        last_yield_samples = 0
-        print(
-            "stream start:",
-            f"min_chunk_sec={STREAM_MIN_CHUNK_SEC}",
-            f"duration_sec={duration_sec}",
-            f"chunk_frames={chunk_frames}",
-        )
-
+        _validate_generation_inputs(lyrics, tags)
+        start_threshold = max(MIN_BUFFER_BLOCKS, PREFETCH_BLOCKS)
+        queue = []
+        started = False
+        print("block stream start:", f"block_sec={MODEL_DETOKENIZE_INTERVAL_SEC}", f"duration_sec={duration_sec}")
         for sr, chunk_np in generate_music_streaming(
             lyrics=lyrics,
             tags=tags,
@@ -752,26 +756,19 @@ def stream_generate(
             keep_model_loaded=keep_model_loaded,
             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]
-            if buffered_samples - last_yield_samples < min_samples:
+            queue.append(chunk_np)
+            if not started and len(queue) < start_threshold and len(queue) < MAX_QUEUE_BLOCKS:
                 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 min_samples > 0 and buffer:
-            full_audio = np.concatenate(buffer)
-            print(f"stream final yield: samples={full_audio.shape[0]}")
-            yield 48000, full_audio
+            if not started:
+                started = True
+                print(f"block stream start playback: buffered_blocks={len(queue)}")
+            while queue:
+                block = queue.pop(0)
+                print(f"block stream yield: samples={block.shape[0]}")
+                yield sr, block
     except Exception as e:
         raise gr.Error(f"Streaming error: {str(e)}")
 
@@ -790,8 +787,6 @@ def stream_generate_accelerated(
     keep_model_loaded,
     offload_mode,
     backend,
-    output_format,
-    chunk_frames,
 ):
     return stream_generate(
         lyrics,
@@ -806,8 +801,6 @@ def stream_generate_accelerated(
         keep_model_loaded,
         offload_mode,
         backend,
-        output_format,
-        chunk_frames,
         True,
     )
 
@@ -977,10 +970,15 @@ Output ONLY the lyrics with structure tags, no explanations.
 
 def create_ui():
     """Create Gradio UI"""
+    show_stream_default = (
+        DEFAULT_GENERATION_MODE != "Original (No Acceleration)"
+        and DEFAULT_SPEED_SUBMODE == "Streaming"
+    )
 
     with gr.Blocks(title="HeartMuLa Music Generation") as demo:
         gr.Markdown("# HeartMuLa Music Generation")
         gr.Markdown("Generate music from lyrics and style tags")
+        gr.Markdown("Tip: start with the **ZeroGPU Safe** preset for reliable generation on small GPUs.")
 
         with gr.Tabs():
             # Tab 1: Music Generation
@@ -1030,10 +1028,19 @@ def create_ui():
                                     t7 = gr.CheckboxGroup(choices=TAG_DATA["Topic"], label="Topic")
                                     tag_checkboxes.append(t7)
 
+                        gr.Markdown("### Quick Presets")
+                        with gr.Row():
+                            preset_selector = gr.Dropdown(
+                                choices=list(PRESET_CONFIGS.keys()),
+                                value=DEFAULT_PRESET,
+                                label="Preset"
+                            )
+                            apply_preset_btn = gr.Button("Apply Preset", size="sm")
+
                         # Generation parameters
                         with gr.Row():
                             cfg_scale = gr.Slider(0.0, 3.0, value=1.5, step=0.1, label="CFG Scale")
-                            duration = gr.Slider(10, 300, value=180, step=10, label="Duration (sec)")
+                            duration = gr.Slider(10, 300, value=DEFAULT_DURATION_SEC, step=10, label="Duration (sec)")
 
                         with gr.Row():
                             temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature")
@@ -1056,9 +1063,14 @@ def create_ui():
                                 label="Codec Version"
                             )
                             quant_mode = gr.Dropdown(
-                                choices=[("None", "none"), ("4-bit (NF4/FP4)", "4bit"), ("8-bit", "8bit")],
-                                value="none",
-                                label="Quantization"
+                                choices=[
+                                    ("None", "none"),
+                                    ("4-bit (NF4/FP4)", "4bit"),
+                                    ("8-bit", "8bit"),
+                                    ("FP8 (TorchAO)", "fp8"),
+                                ],
+                                value=DEFAULT_QUANT_MODE,
+                                label="Quantization (ZeroGPU recommended: 4-bit)"
                             )
                             keep_model_loaded = gr.Checkbox(
                                 value=DEFAULT_KEEP_MODEL_LOADED,
@@ -1066,43 +1078,76 @@ def create_ui():
                             )
                             offload_mode = gr.Dropdown(
                                 choices=["auto", "aggressive"],
-                                value="auto",
+                                value=DEFAULT_OFFLOAD_MODE,
                                 label="Offload Mode"
                             )
                             output_format = gr.Radio(
                                 choices=[("WAV", "wav"), ("MP3", "mp3")],
                                 value="wav",
-                                label="Output Format"
+                                label="Output Format",
+                                visible=not show_stream_default,
                             )
-                            chunk_frames = gr.Slider(
-                                5, 100, value=20, step=1, label="Streaming Chunk Frames"
+                            gr.Markdown(
+                                f"Streaming is block-based: {MODEL_DETOKENIZE_INTERVAL_SEC:.2f}s per block "
+                                f"({BLOCK_FRAMES} frames, {BLOCK_SAMPLES} samples)."
                             )
 
                         gr.Markdown("### 🚀 Generation")
 
                         generation_mode = gr.Radio(
                             choices=["Original (No Acceleration)", "Accelerated"],
-                            value="Original (No Acceleration)",
+                            value=DEFAULT_GENERATION_MODE,
                             label="Generation Mode",
                         )
 
                         speed_submode = gr.Radio(
-                            choices=["Standard", "Streaming"],
-                            value="Standard",
+                            choices=[
+                                ("Standard", "Standard"),
+                                ("Block Streaming (Preview)", "Streaming"),
+                            ],
+                            value=DEFAULT_SPEED_SUBMODE,
                             label="Accelerated Options",
-                            visible=False,
+                            visible=DEFAULT_GENERATION_MODE != "Original (No Acceleration)",
                         )
 
-                        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)
+                        btn_original = gr.Button(
+                            "🎼 Generate Music (Original)",
+                            variant="primary",
+                            size="lg",
+                            visible=DEFAULT_GENERATION_MODE == "Original (No Acceleration)",
+                        )
+                        btn_accel = gr.Button(
+                            "🎼 Generate Music (Accelerated)",
+                            variant="primary",
+                            size="lg",
+                            visible=DEFAULT_GENERATION_MODE != "Original (No Acceleration)"
+                            and not show_stream_default,
+                        )
+                        btn_stream = gr.Button(
+                            "🎼 Generate Music (Block Streaming)",
+                            variant="primary",
+                            size="lg",
+                            visible=show_stream_default,
+                        )
+                        cancel_stream_btn = gr.Button(
+                            "Cancel Streaming",
+                            variant="secondary",
+                            size="lg",
+                            visible=show_stream_default,
+                        )
                         cancel_state = gr.State()
 
                     with gr.Column():
                         # Notice section
                         with gr.Accordion("Usage Notice", open=True):
                             gr.Markdown("""
+### ZeroGPU Tips
+- Prefer **4-bit quantization** and **aggressive offload**
+- Keep durations short (<= 60s) for faster first results
+- **Streaming is block-based**, not token streaming: each block is ~29.76s
+- Use Streaming for preview; Standard mode outputs a full downloadable file
+- Try **FP8 + AOTI** (preset: ZeroGPU AOTI FP8) for the fastest H200 path
+
 ### Lyrics Format Requirements
 
 **Automatic Processing:**
@@ -1142,14 +1187,16 @@ Every day the fire burns
                             label="Generated Music (Full)",
                             type="filepath",
                             interactive=False,
+                            visible=not show_stream_default,
                         )
                         stream_audio = gr.Audio(
-                            label="Streaming Audio (Preview)",
+                            label="Block Streaming Audio (Preview)",
                             streaming=True,
                             autoplay=True,
                             type="numpy",
                             format="wav",
                             interactive=False,
+                            visible=show_stream_default,
                         )
 
                 # Event handlers for tag selection
@@ -1171,6 +1218,9 @@ Every day the fire burns
                             gr.update(visible=False),  # btn_accel
                             gr.update(visible=False),  # btn_stream
                             gr.update(visible=False),  # cancel_stream_btn
+                            gr.update(visible=True),  # output_format
+                            gr.update(visible=True),  # output_audio_file
+                            gr.update(visible=False),  # stream_audio
                         )
                     show_stream = spd_mode == "Streaming"
                     return (
@@ -1179,17 +1229,83 @@ Every day the fire burns
                         gr.update(visible=not show_stream),  # btn_accel
                         gr.update(visible=show_stream),  # btn_stream
                         gr.update(visible=show_stream),  # cancel_stream_btn
+                        gr.update(visible=not show_stream),  # output_format
+                        gr.update(visible=not show_stream),  # output_audio_file
+                        gr.update(visible=show_stream),  # stream_audio
+                    )
+
+                def apply_preset(preset_name):
+                    preset = PRESET_CONFIGS.get(preset_name, PRESET_CONFIGS[DEFAULT_PRESET])
+                    return (
+                        gr.update(value=preset["cfg_scale"]),
+                        gr.update(value=preset["duration"]),
+                        gr.update(value=preset["temperature"]),
+                        gr.update(value=preset["topk"]),
+                        gr.update(value=preset["quant_mode"]),
+                        gr.update(value=preset["keep_model_loaded"]),
+                        gr.update(value=preset["offload_mode"]),
+                        gr.update(value=preset["generation_mode"]),
+                        gr.update(value=preset["speed_submode"]),
                     )
 
                 generation_mode.change(
                     fn=update_visibility,
                     inputs=[generation_mode, speed_submode],
-                    outputs=[speed_submode, btn_original, btn_accel, btn_stream, cancel_stream_btn],
+                    outputs=[
+                        speed_submode,
+                        btn_original,
+                        btn_accel,
+                        btn_stream,
+                        cancel_stream_btn,
+                        output_format,
+                        output_audio_file,
+                        stream_audio,
+                    ],
                 )
                 speed_submode.change(
                     fn=update_visibility,
                     inputs=[generation_mode, speed_submode],
-                    outputs=[speed_submode, btn_original, btn_accel, btn_stream, cancel_stream_btn],
+                    outputs=[
+                        speed_submode,
+                        btn_original,
+                        btn_accel,
+                        btn_stream,
+                        cancel_stream_btn,
+                        output_format,
+                        output_audio_file,
+                        stream_audio,
+                    ],
+                )
+
+                preset_event = apply_preset_btn.click(
+                    fn=apply_preset,
+                    inputs=[preset_selector],
+                    outputs=[
+                        cfg_scale,
+                        duration,
+                        temperature,
+                        topk,
+                        quant_mode,
+                        keep_model_loaded,
+                        offload_mode,
+                        generation_mode,
+                        speed_submode,
+                    ],
+                )
+
+                preset_event.then(
+                    fn=update_visibility,
+                    inputs=[generation_mode, speed_submode],
+                    outputs=[
+                        speed_submode,
+                        btn_original,
+                        btn_accel,
+                        btn_stream,
+                        cancel_stream_btn,
+                        output_format,
+                        output_audio_file,
+                        stream_audio,
+                    ],
                 )
 
                 btn_original.click(
@@ -1251,8 +1367,6 @@ Every day the fire burns
                         keep_model_loaded,
                         offload_mode,
                         backend,
-                        output_format,
-                        chunk_frames,
                     ],
                     outputs=[stream_audio],
                     concurrency_id="gpu_queue",
@@ -1397,81 +1511,8 @@ Every day the fire burns
                     outputs=[lyrics]
                 )
 
-            # Tab 3: Lyrics Transcription
-            with gr.Tab("Lyrics Transcription"):
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("### Transcribe or Translate Lyrics from Audio")
-
-                        audio_input = gr.Audio(
-                            label="Audio Input",
-                            type="filepath"
-                        )
-
-                        task_select = gr.Radio(
-                            choices=[
-                                ("Transcribe (Original Language)", "transcribe"),
-                                ("Translate to English", "translate")
-                            ],
-                            value="transcribe",
-                            label="Task"
-                        )
-
-                        max_new_tokens = gr.Slider(
-                            64, 512, value=256, step=16, label="Max New Tokens"
-                        )
-                        num_beams = gr.Slider(
-                            1, 5, value=2, step=1, label="Beam Search"
-                        )
-                        transcribe_temperature = gr.Slider(
-                            0.0, 0.8, value=0.2, step=0.1, label="Temperature"
-                        )
-
-                        transcribe_btn = gr.Button(
-                            "Run Transcription",
-                            variant="primary",
-                            size="lg"
-                        )
-
-                        use_generated_audio = gr.Button(
-                            "Use Generated Music",
-                            size="sm"
-                        )
-
-                    with gr.Column():
-                        with gr.Accordion("Notes", open=True):
-                            gr.Markdown("""
-### Notes
-- Best results come from **vocals-only** stems.
-- If you pass full mixes, consider source separation first.
-- The HeartTranscriptor model will auto-download on first use.
-                            """)
-
-                        transcription_output = gr.Textbox(
-                            label="Transcription Result",
-                            lines=18,
-                            placeholder="Transcribed lyrics will appear here...",
-                            interactive=False
-                        )
-
-                transcribe_btn.click(
-                    fn=transcribe_audio,
-                    inputs=[audio_input, task_select, max_new_tokens, num_beams, transcribe_temperature],
-                    outputs=[transcription_output],
-                    concurrency_id="gpu_queue",
-                    concurrency_limit=GPU_CONCURRENCY_LIMIT,
-                )
-
-                use_generated_audio.click(
-                    fn=lambda x: x,
-                    inputs=[output_audio_file],
-                    outputs=[audio_input]
-                )
-
     return demo
 
-
-demo = create_ui()
 demo.queue(max_size=GRADIO_QUEUE_MAX_SIZE, default_concurrency_limit=GRADIO_DEFAULT_CONCURRENCY)
 
 if __name__ == "__main__":

src/heartlib/pipelines/music_generation.py

@@ -1,17 +1,30 @@
 import gc
+import inspect
 import json
 import os
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Iterator
+from typing import Any, Dict, Optional, Iterator, Tuple
 
 import torch
 import torchaudio
 from tokenizers import Tokenizer
 from transformers import BitsAndBytesConfig
 
-from ..heartmula.modeling_heartmula import HeartMuLa
+from ..heartmula.modeling_heartmula import HeartMuLa, _index_causal_mask
 from ..heartcodec.modeling_heartcodec import HeartCodec
 
+try:
+    import spaces  # type: ignore
+except Exception:  # pragma: no cover - optional dependency
+    spaces = None
+
+
+def _env_bool(name: str) -> Optional[bool]:
+    value = os.environ.get(name)
+    if value is None:
+        return None
+    return value.strip().lower() in ("1", "true", "yes", "y", "on")
+
 
 @dataclass
 class HeartMuLaGenConfig:
@@ -44,6 +57,9 @@ class HeartMuLaGenPipeline:
         lazy_load: bool = False,
         use_deepspeed: bool = False,
         ds_inference_config: Optional[Dict[str, Any]] = None,
+        torchao_quantize: bool = False,
+        aoti_enabled: Optional[bool] = None,
+        aoti_dynamic_shapes: Optional[bool] = None,
     ):
         self.model = model
         self.audio_codec = audio_codec
@@ -62,6 +78,17 @@ class HeartMuLaGenPipeline:
         self._ds_inference_applied = False
         self._parallel_number = num_quantizers + 1 if num_quantizers else 9
         self._muq_dim = model.config.muq_dim if model else None
+        self._torchao_quantize = torchao_quantize
+        space_default = spaces is not None and os.environ.get("SPACE_ID") is not None
+        env_aoti = _env_bool("ENABLE_AOTI")
+        env_dynamic = _env_bool("AOTI_DYNAMIC_SHAPES")
+        self._aoti_enabled = env_aoti if aoti_enabled is None else aoti_enabled
+        if self._aoti_enabled is None:
+            self._aoti_enabled = space_default
+        self._aoti_dynamic_shapes = env_dynamic if aoti_dynamic_shapes is None else aoti_dynamic_shapes
+        if self._aoti_dynamic_shapes is None:
+            self._aoti_dynamic_shapes = space_default
+        self._aoti_compiled = False
 
     def _init_deepspeed_inference(self):
         if not self.use_deepspeed or self._ds_inference_applied:
@@ -92,6 +119,9 @@ class HeartMuLaGenPipeline:
             )
         if str(next(self.model.parameters()).device) != str(self.device):
             self.model.to(self.device)
+        if self._torchao_quantize and not self._aoti_enabled:
+            self._maybe_fp8_quantize(self.model.backbone)
+            self._maybe_fp8_quantize(self.model.decoder)
         self.model.eval()
         self._init_deepspeed_inference()
         self._muq_dim = self.model.config.muq_dim
@@ -120,6 +150,157 @@ class HeartMuLaGenPipeline:
             return torch.autocast(device_type="cuda", dtype=self.dtype)
         return torch.inference_mode()
 
+    def _supports_aoti(self) -> Tuple[bool, str]:
+        if not self._aoti_enabled:
+            return False, "AOTI disabled"
+        if spaces is None:
+            return False, "spaces package not available"
+        if not torch.cuda.is_available():
+            return False, "CUDA unavailable"
+        if not hasattr(torch, "export"):
+            return False, "torch.export unavailable"
+        if self.use_deepspeed:
+            return False, "DeepSpeed enabled"
+        if self.bnb_config is not None:
+            return False, "bitsandbytes quantization active"
+        return True, ""
+
+    def _maybe_fp8_quantize(self, module: torch.nn.Module) -> bool:
+        if not self._torchao_quantize:
+            return False
+        if not torch.cuda.is_available():
+            print("AOTI FP8: CUDA unavailable, skipping torchao quantization.")
+            return False
+        try:
+            major, _ = torch.cuda.get_device_capability()
+        except Exception:
+            major = 0
+        if major < 9:
+            print("AOTI FP8: requires CUDA compute capability >= 9.0, skipping.")
+            return False
+        try:
+            from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
+        except Exception:
+            print("AOTI FP8: torchao not available, skipping.")
+            return False
+        quantize_(module, Float8DynamicActivationFloat8WeightConfig())
+        return True
+
+    def _aoti_compile_module(
+        self,
+        module: torch.nn.Module,
+        example_kwargs: Dict[str, Any],
+        dynamic_shapes: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        if spaces is None:
+            return
+        with spaces.aoti_capture(module) as call:
+            module(**example_kwargs)
+        exported = torch.export.export(
+            module,
+            args=call.args,
+            kwargs=call.kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+        compiled = spaces.aoti_compile(exported)
+        spaces.aoti_apply(compiled, module)
+
+    def _maybe_apply_aoti(self, model_inputs: Dict[str, Any], cfg_scale: float) -> None:
+        if self._aoti_compiled:
+            return
+        ok, reason = self._supports_aoti()
+        if not ok:
+            print(f"AOTI: skip ({reason}).")
+            self._aoti_compiled = True
+            return
+
+        try:
+            self.load_heartmula()
+
+            prompt_tokens = model_inputs["tokens"]
+            prompt_tokens_mask = model_inputs["tokens_mask"]
+            continuous_segment = model_inputs["muq_embed"]
+            starts = model_inputs["muq_idx"]
+            prompt_pos = model_inputs["pos"]
+
+            bs_size = 2 if cfg_scale != 1.0 else 1
+            self.model.setup_caches(bs_size)
+
+            uncond_mask = None
+            if cfg_scale > 1.0 and prompt_tokens.shape[0] > 1:
+                actual_B = prompt_tokens.shape[0] // 2
+                uncond_mask = torch.cat(
+                    [
+                        torch.zeros(actual_B, dtype=torch.bool, device=prompt_tokens.device),
+                        torch.ones(actual_B, dtype=torch.bool, device=prompt_tokens.device),
+                    ]
+                )
+
+            embeds = self.model._embed_tokens(prompt_tokens, uncond_mask=uncond_mask)
+            masked_embeds = embeds * prompt_tokens_mask.unsqueeze(-1)
+            h = masked_embeds.sum(dim=2, dtype=embeds.dtype)
+            if continuous_segment is not None:
+                continuous_segment = self.model.muq_linear(continuous_segment)
+                if uncond_mask is not None:
+                    uncond_embed = self.model.unconditional_text_embedding(
+                        torch.zeros(1, device=prompt_tokens.device, dtype=torch.long)
+                    )
+                    mask_expanded = uncond_mask.view(h.shape[0], 1).expand_as(continuous_segment)
+                    continuous_segment = torch.where(mask_expanded, uncond_embed, continuous_segment)
+                batch_indices = torch.arange(h.shape[0], device=h.device)
+                h[batch_indices, starts] = continuous_segment
+
+            curr_backbone_mask = _index_causal_mask(self.model.backbone_causal_mask, prompt_pos)
+
+            backbone_params = list(inspect.signature(self.model.backbone.forward).parameters.keys())
+            backbone_name = backbone_params[1] if len(backbone_params) > 1 else "x"
+            backbone_kwargs = {
+                backbone_name: h,
+                "input_pos": prompt_pos,
+                "mask": curr_backbone_mask,
+            }
+
+            if self._aoti_dynamic_shapes and hasattr(torch, "export"):
+                seq_dim = torch.export.Dim("seq", min=4, max=self.model.backbone.max_seq_len)
+                backbone_dynamic = {
+                    backbone_name: {1: seq_dim},
+                    "input_pos": {1: seq_dim},
+                    "mask": {1: seq_dim},
+                }
+            else:
+                backbone_dynamic = None
+
+            self._maybe_fp8_quantize(self.model.backbone)
+            self._aoti_compile_module(self.model.backbone, backbone_kwargs, backbone_dynamic)
+
+            last_h = h[:, -1, :]
+            c0_sample = torch.zeros((h.shape[0], 1), device=h.device, dtype=torch.long)
+            c0_embed = self.model._embed_audio(0, c0_sample)
+            curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1)
+            curr_pos = (
+                torch.arange(0, curr_h.size(1), device=curr_h.device)
+                .unsqueeze(0)
+                .repeat(curr_h.size(0), 1)
+            )
+            curr_decoder_mask = _index_causal_mask(self.model.decoder_causal_mask, curr_pos)
+
+            decoder_params = list(inspect.signature(self.model.decoder.forward).parameters.keys())
+            decoder_name = decoder_params[1] if len(decoder_params) > 1 else "x"
+            decoder_kwargs = {
+                decoder_name: self.model.projection(curr_h),
+                "input_pos": curr_pos,
+                "mask": curr_decoder_mask,
+            }
+            decoder_dynamic = None
+
+            self._maybe_fp8_quantize(self.model.decoder)
+            self._aoti_compile_module(self.model.decoder, decoder_kwargs, decoder_dynamic)
+            self._aoti_compiled = True
+            print("AOTI: compiled backbone and decoder successfully.")
+        except Exception as exc:
+            self._aoti_compiled = True
+            print(f"AOTI: compilation failed, falling back to eager. Error: {exc}")
+
     def preprocess(self, inputs: Dict[str, Any], cfg_scale: float):
         self.load_heartmula()
         tags = inputs["tags"]
@@ -258,6 +439,7 @@ class HeartMuLaGenPipeline:
         offload_mode: str = "auto",
     ) -> Iterator[torch.Tensor]:
         model_inputs = self.preprocess(inputs, cfg_scale=cfg_scale)
+        self._maybe_apply_aoti(model_inputs, cfg_scale)
         self.load_heartcodec()
 
         prompt_tokens = model_inputs["tokens"]
@@ -352,6 +534,7 @@ class HeartMuLaGenPipeline:
         offload_mode: str = "auto",
     ):
         model_inputs = self.preprocess(inputs, cfg_scale=cfg_scale)
+        self._maybe_apply_aoti(model_inputs, cfg_scale)
         frames = self._generate_frames(model_inputs, max_audio_length_ms, temperature, topk, cfg_scale)
         wav = self._decode_frames(frames)
         torchaudio.save(save_path, wav, 48000)
@@ -369,6 +552,9 @@ class HeartMuLaGenPipeline:
         lazy_load: bool = False,
         use_deepspeed: bool = False,
         ds_inference_config: Optional[Dict[str, Any]] = None,
+        torchao_quantize: bool = False,
+        aoti_enabled: Optional[bool] = None,
+        aoti_dynamic_shapes: Optional[bool] = None,
     ):
         if os.path.isfile(vocab_path := os.path.join(pretrained_path, "tokenizer.json")):
             tokenizer = Tokenizer.from_file(vocab_path)
@@ -415,4 +601,7 @@ class HeartMuLaGenPipeline:
             lazy_load=lazy_load,
             use_deepspeed=use_deepspeed,
             ds_inference_config=ds_inference_config,
+            torchao_quantize=torchao_quantize,
+            aoti_enabled=aoti_enabled,
+            aoti_dynamic_shapes=aoti_dynamic_shapes,
         )