Fix ACE-Step ZeroGPU singing pipeline

audio_engine.py

@@ -5,17 +5,20 @@ from __future__ import annotations
 import logging
 import os
 import re
+import tempfile
 import time
 import traceback
+from typing import Any
 
-# Import spaces before torch/diffusers so ZeroGPU can install its CUDA shim.
+# IMPORTANT:
+# Import spaces before torch/diffusers so Hugging Face ZeroGPU can install
+# its CUDA shim before PyTorch is imported.
 try:
     import spaces
 except ModuleNotFoundError:
     if os.environ.get("SPACE_ID"):
-        raise  # A Space needs the ZeroGPU runtime and its spaces module.
+        raise
 
-    # Local CPU machines can still open the form and write rhymes.
     class _LocalSpaces:
         @staticmethod
         def GPU(**_kwargs):
@@ -23,101 +26,233 @@ except ModuleNotFoundError:
 
     spaces = _LocalSpaces()
 
+
 import numpy as np
+import soundfile as sf
+
 from rhyme_engine import LANGUAGES, validate_lyric_for_audio
 
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("kids_rhyme.audio")
 
+
 MODEL_ID = "ACE-Step/acestep-v15-xl-turbo-diffusers"
-VOCAL_LANGUAGE = {key: value["code"] for key, value in LANGUAGES.items()}
-FALLBACK_SAMPLE_RATE = 48000
 
-# Turn this on temporarily in Hugging Face Settings -> Variables and secrets:
-# KIDS_DEBUG=1
-# Leave it unset (or set it to 0) when the public app is working.
-DEBUG_ERRORS = os.environ.get("KIDS_DEBUG", "0") == "1"
+VOCAL_LANGUAGE = {
+    key: value["code"]
+    for key, value in LANGUAGES.items()
+}
 
-# ZeroGPU expects the model to be placed on CUDA at module import time. Its
-# CUDA shim makes this possible before the real GPU is allocated to a request.
-_pipe = None
-_load_error = None
-try:
-    import torch
-except ModuleNotFoundError:
-    torch = None
+FALLBACK_SAMPLE_RATE = 48000
 
-if torch is not None and torch.cuda.is_available():
-    try:
-        from diffusers import AceStepPipeline
+DEBUG_ERRORS = os.environ.get(
+    "KIDS_DEBUG",
+    "1",
+) == "1"
 
-        # diffusers versions have used both names; support either without
-        # changing the rest of the app.
-        try:
-            _pipe = AceStepPipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
-        except TypeError:
-            _pipe = AceStepPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
 
-        _pipe = _pipe.to("cuda")
-        logger.info("Singing model loaded: %s", MODEL_ID)
-    except Exception as exc:
-        _load_error = exc
-        logger.error("Singing model failed to load:\n%s", traceback.format_exc())
+# ---------------------------------------------------------------------
+# DO NOT create/move the pipeline to CUDA at module import time.
+#
+# On ZeroGPU there may be no CUDA device allocated while the module is
+# importing. CUDA becomes available inside the @spaces.GPU function.
+# ---------------------------------------------------------------------
+
+_pipe = None
 
 
 def _debug_detail(exc: BaseException) -> str:
-    """Return a short browser-safe debug suffix only while KIDS_DEBUG=1."""
     if not DEBUG_ERRORS:
         return ""
+
     message = str(exc).replace("\n", " ").strip()
-    return f" [debug: {type(exc).__name__}: {message[:500]}]"
 
+    return (
+        f" [debug: {type(exc).__name__}: "
+        f"{message[:1000]}]"
+    )
 
-def _require_pipeline():
-    if _pipe is None:
-        if _load_error is not None:
-            raise RuntimeError(
-                "The singing model could not load. Check the Space logs."
-                + _debug_detail(_load_error)
-            ) from _load_error
+
+def _load_pipeline():
+    """
+    Load ACE-Step while inside the ZeroGPU allocation.
+
+    The pipeline is cached between calls when possible, but CUDA placement
+    is never attempted during module import.
+    """
+    global _pipe
+
+    import torch
+    from diffusers import AceStepPipeline
+
+    if not torch.cuda.is_available():
         raise RuntimeError(
-            "Singing needs GPU hardware. Set the Space hardware to ZeroGPU."
+            "CUDA is unavailable inside the ZeroGPU function. "
+            "Check that the Space is using ZeroGPU hardware and that "
+            "this function is running through spaces.GPU."
+        )
+
+    if _pipe is None:
+        logger.info(
+            "Loading ACE-Step pipeline: %s",
+            MODEL_ID,
+        )
+
+        try:
+            _pipe = AceStepPipeline.from_pretrained(
+                MODEL_ID,
+                torch_dtype=torch.bfloat16,
+            )
+        except TypeError:
+            # Newer Diffusers versions prefer dtype.
+            _pipe = AceStepPipeline.from_pretrained(
+                MODEL_ID,
+                dtype=torch.bfloat16,
+            )
+
+        logger.info(
+            "ACE-Step pipeline loaded on CPU."
         )
+
+        try:
+            _pipe.vae.enable_tiling()
+            logger.info("ACE-Step VAE tiling enabled.")
+        except Exception:
+            logger.info(
+                "VAE tiling unavailable; continuing."
+            )
+
+    logger.info(
+        "Moving ACE-Step pipeline to ZeroGPU CUDA device."
+    )
+
+    _pipe.to("cuda")
+
     return _pipe
 
 
-def _sample_rate() -> int:
-    pipe = _require_pipeline()
-    rate = getattr(pipe, "sample_rate", None)
-    if rate is None:
-        vae_cfg = getattr(getattr(pipe, "vae", None), "config", None)
-        rate = getattr(vae_cfg, "sampling_rate", None)
-    if rate is None:
-        logger.warning(
-            "No sample rate found on pipeline; using %s Hz for %s",
-            FALLBACK_SAMPLE_RATE,
-            MODEL_ID,
+def _find_sample_rate(
+    pipe: Any,
+    result: Any = None,
+) -> int:
+    """
+    Try known locations for the model/output sample rate instead of blindly
+    assuming 48 kHz.
+    """
+
+    candidates = []
+
+    if result is not None:
+        for attr in (
+            "sample_rate",
+            "sampling_rate",
+            "audio_sample_rate",
+        ):
+            candidates.append(
+                getattr(result, attr, None)
+            )
+
+        if isinstance(result, dict):
+            for key in (
+                "sample_rate",
+                "sampling_rate",
+                "audio_sample_rate",
+            ):
+                candidates.append(result.get(key))
+
+    for attr in (
+        "sample_rate",
+        "sampling_rate",
+        "audio_sample_rate",
+    ):
+        candidates.append(
+            getattr(pipe, attr, None)
         )
-        rate = FALLBACK_SAMPLE_RATE
-    if isinstance(rate, (bool, np.bool_)) or not isinstance(rate, (int, np.integer)) or rate <= 0:
-        raise ValueError(f"The pipeline returned an invalid sample rate: {rate!r}")
-    return int(rate)
+
+    vae = getattr(pipe, "vae", None)
+    vae_config = getattr(vae, "config", None)
+
+    if vae_config is not None:
+        for attr in (
+            "sample_rate",
+            "sampling_rate",
+            "audio_sample_rate",
+        ):
+            candidates.append(
+                getattr(vae_config, attr, None)
+            )
+
+    config = getattr(pipe, "config", None)
+
+    if config is not None:
+        for attr in (
+            "sample_rate",
+            "sampling_rate",
+            "audio_sample_rate",
+        ):
+            candidates.append(
+                getattr(config, attr, None)
+            )
+
+    for rate in candidates:
+        if (
+            not isinstance(rate, (bool, np.bool_))
+            and isinstance(rate, (int, np.integer))
+            and int(rate) > 0
+        ):
+            logger.info(
+                "Detected ACE-Step sample rate: %s Hz",
+                rate,
+            )
+            return int(rate)
+
+    logger.warning(
+        "ACE-Step did not expose a sample rate; "
+        "falling back to %s Hz.",
+        FALLBACK_SAMPLE_RATE,
+    )
+
+    return FALLBACK_SAMPLE_RATE
 
 
 def _song_lyrics(text: str) -> str:
-    """Preserve every edited line and mark its two parts for the song model."""
-    text = text.replace("\r\n", "\n").replace("\r", "\n")
-    stanzas = [part.strip() for part in re.split(r"\n\s*\n", text) if part.strip()]
+    text = (
+        text.replace("\r\n", "\n")
+        .replace("\r", "\n")
+    )
+
+    stanzas = [
+        part.strip()
+        for part in re.split(r"\n\s*\n", text)
+        if part.strip()
+    ]
+
     if len(stanzas) >= 2:
         first = stanzas[0]
         second = "\n".join(stanzas[1:])
+
     else:
-        lines = [line.strip() for line in text.splitlines() if line.strip()]
+        lines = [
+            line.strip()
+            for line in text.splitlines()
+            if line.strip()
+        ]
+
         if len(lines) < 2:
-            raise ValueError("Add at least two short lines to sing.")
+            raise ValueError(
+                "Add at least two short lines to sing."
+            )
+
         midpoint = (len(lines) + 1) // 2
-        first, second = "\n".join(lines[:midpoint]), "\n".join(lines[midpoint:])
-    return f"[verse]\n{first}\n[chorus]\n{second}"
+
+        first = "\n".join(lines[:midpoint])
+        second = "\n".join(lines[midpoint:])
+
+    return (
+        f"[verse]\n{first}\n"
+        f"[chorus]\n{second}"
+    )
 
 
 def _song_settings(
@@ -128,130 +263,398 @@ def _song_settings(
     theme_prompt: str = "",
 ) -> dict:
     text = validate_lyric_for_audio(text)
-    if language not in VOCAL_LANGUAGE or mood not in ("Bouncy", "Calm"):
-        raise ValueError("Write a rhyme first to select its language and music mood.")
+
+    if (
+        language not in VOCAL_LANGUAGE
+        or mood not in ("Bouncy", "Calm")
+    ):
+        raise ValueError(
+            "Write a rhyme first to select its "
+            "language and music mood."
+        )
 
     if mood == "Calm":
         prompt = (
-            "Gentle original children's lullaby, a clear warm voice SINGING a simple "
-            "memorable melody in the language of the lyrics. Soft piano, glockenspiel, "
-            "light acoustic guitar, slow swaying rhythm. Vocal-forward mix. "
-            "Sing the supplied lyrics; no spoken words or narration."
+            "Gentle original children's lullaby, "
+            "a clear warm voice SINGING a simple "
+            "memorable melody in the language of the lyrics. "
+            "Soft piano, glockenspiel, light acoustic guitar, "
+            "slow swaying rhythm. Vocal-forward mix. "
+            "Sing the supplied lyrics; "
+            "no spoken words or narration."
         )
         bpm = 82
+
     else:
         prompt = (
-            "Playful original children's sing-along, a clear cheerful voice SINGING "
-            "a simple catchy melody in the language of the lyrics. Ukulele, handclaps, "
-            "toy piano, bright steady beat. Vocal-forward mix. "
-            "Sing the supplied lyrics; no spoken words or narration."
+            "Playful original children's sing-along, "
+            "a clear cheerful voice SINGING "
+            "a simple catchy melody in the language of the lyrics. "
+            "Ukulele, handclaps, toy piano, bright steady beat. "
+            "Vocal-forward mix. "
+            "Sing the supplied lyrics; "
+            "no spoken words or narration."
         )
         bpm = 112
 
     language_name = LANGUAGES[language]["name"]
-    prompt = f"Sing all vocals in {language_name}. " + prompt
+
+    prompt = (
+        f"Sing all vocals in {language_name}. "
+        + prompt
+    )
+
     if theme:
         prompt += f" Song theme: {theme}."
+
     if theme_prompt:
         prompt += f" Topic: {theme_prompt}."
 
-    line_count = sum(bool(line.strip()) for line in text.splitlines())
+    line_count = sum(
+        bool(line.strip())
+        for line in text.splitlines()
+    )
+
     return {
         "prompt": prompt,
         "lyrics": _song_lyrics(text),
         "vocal_language": VOCAL_LANGUAGE[language],
-        # 40 seconds for a normal rhyme, with a little extra room for longer text.
-        "audio_duration": min(56.0, 40.0 + 4.0 * max(0, line_count - 8)),
+        "audio_duration": min(
+            56.0,
+            40.0 + 4.0 * max(0, line_count - 8),
+        ),
         "num_inference_steps": 8,
         "bpm": bpm,
         "task_type": "text2music",
-        "output_type": "np",
     }
 
 
-def _extract_audio(result) -> np.ndarray:
-    """Extract ACE-Step audio and normalize it to [channels, samples]."""
+def _extract_audio(result: Any) -> np.ndarray:
+    """
+    Normalize ACE-Step output into float32 [channels, samples].
+
+    Handles tensors, numpy arrays, lists/batches and common Diffusers
+    pipeline output containers.
+    """
+
     value = None
 
-    # Current diffusers ACE-Step output uses `audios`.
     if hasattr(result, "audios"):
         value = result.audios
+
     elif hasattr(result, "audio"):
         value = result.audio
+
     elif isinstance(result, dict):
-        for key in ("audios", "audio", "waveform", "sample"):
+        for key in (
+            "audios",
+            "audio",
+            "waveform",
+            "sample",
+        ):
             if key in result:
                 value = result[key]
                 break
+
     elif isinstance(result, (tuple, list)) and result:
         value = result[0]
 
     if value is None:
-        raise ValueError(f"No audio was returned by {type(result).__name__}.")
+        raise RuntimeError(
+            "ACE-Step returned no audio. "
+            f"Result type: {type(result).__name__}"
+        )
+
+    # result.audios may itself be a batch/list.
+    if isinstance(value, (list, tuple)):
+        if not value:
+            raise RuntimeError(
+                "ACE-Step returned an empty audio list."
+            )
+        value = value[0]
 
     if hasattr(value, "detach"):
-        value = value.detach().float().cpu().numpy()
+        value = (
+            value.detach()
+            .float()
+            .cpu()
+            .numpy()
+        )
+
+    arr = np.asarray(value)
+
+    logger.info(
+        "Raw ACE-Step audio: type=%s shape=%s dtype=%s",
+        type(value).__name__,
+        getattr(arr, "shape", None),
+        getattr(arr, "dtype", None),
+    )
 
-    arr = np.asarray(value, dtype=np.float32)
+    arr = arr.astype(
+        np.float32,
+        copy=False,
+    )
 
-    # Typical result: [batch, channels, samples]. Take first batch.
+    # Typical batched forms:
+    # [batch, channels, samples]
+    # [batch, samples]
     while arr.ndim > 2:
         arr = arr[0]
+
     if arr.ndim == 1:
-        arr = arr[None, :]
+        arr = arr[np.newaxis, :]
+
     if arr.ndim != 2:
-        raise ValueError(f"Unexpected audio rank/shape {arr.shape}.")
+        raise RuntimeError(
+            "Unexpected ACE-Step audio shape: "
+            f"{arr.shape}"
+        )
 
-    # If it came back [samples, channels], transpose it.
-    if arr.shape[0] > 2 and arr.shape[1] in (1, 2):
+    # Normalize to [channels, samples].
+    #
+    # If first dimension clearly looks like samples and the second
+    # dimension is mono/stereo, transpose it.
+    if (
+        arr.shape[0] > 2
+        and arr.shape[1] in (1, 2)
+    ):
         arr = arr.T
 
+    if arr.shape[0] not in (1, 2):
+        raise RuntimeError(
+            "Unexpected ACE-Step channel layout: "
+            f"{arr.shape}"
+        )
+
+    if not np.isfinite(arr).all():
+        raise RuntimeError(
+            "ACE-Step produced NaN or infinite audio values."
+        )
+
+    peak = float(np.max(np.abs(arr)))
+
+    if peak < 0.001:
+        raise RuntimeError(
+            "ACE-Step returned silent audio."
+        )
+
     return np.ascontiguousarray(arr)
 
 
-def _polish(wave: np.ndarray, sr: int) -> np.ndarray:
-    """Peak-normalize and fade the tail so playback ends cleanly."""
-    wave = np.array(wave, dtype=np.float32, copy=True)
+def _polish(
+    wave: np.ndarray,
+    sr: int,
+) -> np.ndarray:
+    wave = np.array(
+        wave,
+        dtype=np.float32,
+        copy=True,
+    )
+
     peak = float(np.max(np.abs(wave)))
+
     if peak > 0:
         wave *= 0.89 / peak
-    n = min(int(0.6 * sr), wave.shape[1] // 4)
-    if n > 0:
-        wave[:, -n:] *= np.linspace(1.0, 0.0, n, dtype=np.float32)
+
+    fade_samples = min(
+        int(0.6 * sr),
+        wave.shape[1] // 4,
+    )
+
+    if fade_samples > 0:
+        wave[:, -fade_samples:] *= np.linspace(
+            1.0,
+            0.0,
+            fade_samples,
+            dtype=np.float32,
+        )
+
     return wave
 
 
-def generate_song(settings: dict) -> tuple[int, np.ndarray]:
-    """Inference core. The caller must provide GPU access."""
-    pipe = _require_pipeline()
-    started = time.monotonic()
+def _save_wav(
+    waveform: np.ndarray,
+    sample_rate: int,
+) -> str:
+    """
+    Save an actual WAV file and return its path.
 
-    result = pipe(**settings)
-    waveform = _extract_audio(result)
-    sr = _sample_rate()
+    Returning a filepath is reliable for Gradio Audio outputs and gives
+    the user a downloadable WAV.
+    """
 
-    logger.info("ACE-Step returned audio shape %s at %s Hz", waveform.shape, sr)
+    if waveform.ndim != 2:
+        raise RuntimeError(
+            f"Invalid waveform shape before WAV save: "
+            f"{waveform.shape}"
+        )
+
+    # soundfile expects:
+    # mono   -> [samples]
+    # stereo -> [samples, channels]
+    if waveform.shape[0] == 1:
+        output = waveform[0]
+    else:
+        output = waveform.T
+
+    output = np.ascontiguousarray(
+        np.clip(
+            output,
+            -1.0,
+            1.0,
+        ),
+        dtype=np.float32,
+    )
 
-    if waveform.shape[0] not in (1, 2) or waveform.shape[1] < sr:
-        raise ValueError(f"Unexpected audio shape {waveform.shape} at {sr} Hz")
-    if not np.isfinite(waveform).all() or float(np.max(np.abs(waveform))) < 0.005:
-        raise ValueError("Model returned empty or invalid audio")
+    temp = tempfile.NamedTemporaryFile(
+        suffix=".wav",
+        delete=False,
+    )
+
+    path = temp.name
+    temp.close()
 
-    waveform = _polish(waveform, sr)
-    pcm = (np.clip(waveform.T, -1.0, 1.0) * 32767).astype(np.int16)
+    sf.write(
+        path,
+        output,
+        samplerate=sample_rate,
+        subtype="PCM_16",
+        format="WAV",
+    )
+
+    if (
+        not os.path.isfile(path)
+        or os.path.getsize(path) <= 44
+    ):
+        raise RuntimeError(
+            "WAV file creation failed."
+        )
 
     logger.info(
-        "Generated %.1fs of audio at %s Hz in %.1fs",
-        len(pcm) / sr,
-        sr,
-        time.monotonic() - started,
+        "Song WAV saved: %s (%d bytes)",
+        path,
+        os.path.getsize(path),
     )
-    return sr, pcm
+
+    return path
 
 
 @spaces.GPU(duration=120)
-def _generate_on_gpu(settings: dict) -> tuple[int, np.ndarray]:
-    return generate_song(settings)
+def _generate_on_gpu(
+    settings: dict,
+) -> str:
+    """
+    EVERYTHING requiring CUDA happens after ZeroGPU allocation.
+    """
+
+    import torch
+
+    started = time.monotonic()
+
+    logger.info(
+        "ZeroGPU allocation entered. "
+        "cuda_available=%s torch=%s",
+        torch.cuda.is_available(),
+        torch.__version__,
+    )
+
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "ZeroGPU allocation did not expose CUDA."
+        )
+
+    try:
+        logger.info(
+            "CUDA device: %s",
+            torch.cuda.get_device_name(0),
+        )
+    except Exception:
+        logger.info(
+            "CUDA device name unavailable."
+        )
+
+    pipe = _load_pipeline()
+
+    logger.info(
+        "Starting ACE-Step inference with settings: %r",
+        settings,
+    )
+
+    try:
+        with torch.inference_mode():
+            result = pipe(**settings)
+
+        logger.info(
+            "ACE-Step result type: %s",
+            type(result).__name__,
+        )
+
+        waveform = _extract_audio(result)
+        sample_rate = _find_sample_rate(
+            pipe,
+            result,
+        )
+
+        logger.info(
+            "ACE-Step normalized audio shape=%s "
+            "sample_rate=%s",
+            waveform.shape,
+            sample_rate,
+        )
+
+        if waveform.shape[1] < sample_rate:
+            raise RuntimeError(
+                "ACE-Step returned less than one second "
+                f"of audio: shape={waveform.shape}, "
+                f"sample_rate={sample_rate}"
+            )
+
+        waveform = _polish(
+            waveform,
+            sample_rate,
+        )
+
+        wav_path = _save_wav(
+            waveform,
+            sample_rate,
+        )
+
+        logger.info(
+            "Generated %.2f seconds of audio in %.2f seconds.",
+            waveform.shape[1] / sample_rate,
+            time.monotonic() - started,
+        )
+
+        return wav_path
+
+    except Exception:
+        # This is deliberately logger.exception rather than a generic
+        # "Singing failed" message. Hugging Face runtime logs will contain
+        # the complete traceback and original exception.
+        logger.exception(
+            "ACE-Step inference failed."
+        )
+        raise
+
+    finally:
+        # Do not delete _pipe here. Keeping the CPU-side object cached can
+        # avoid re-downloading/reconstructing it. Move it back off the
+        # leased ZeroGPU CUDA device before leaving the GPU scope.
+        if _pipe is not None:
+            try:
+                _pipe.to("cpu")
+                logger.info(
+                    "ACE-Step pipeline moved back to CPU."
+                )
+            except Exception:
+                logger.exception(
+                    "Could not move ACE-Step pipeline back to CPU."
+                )
+
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
 
 
 def make_sung_song(
@@ -260,17 +663,41 @@ def make_sung_song(
     mood: str,
     theme: str = "",
     theme_prompt: str = "",
-) -> tuple[tuple[int, np.ndarray], str]:
-    """Validate first, then generate. Keep the app's five-argument API."""
-    settings = _song_settings(text, language, mood, theme, theme_prompt)
-    _require_pipeline()
+):
+    """
+    Called by app.py.
+
+    Keep this exact five-argument signature.
+    """
 
     try:
-        audio = _generate_on_gpu(settings)
+        settings = _song_settings(
+            text=text,
+            language=language,
+            mood=mood,
+            theme=theme,
+            theme_prompt=theme_prompt,
+        )
+
+        wav_path = _generate_on_gpu(
+            settings
+        )
+
+        return (
+            wav_path,
+            "Your sung song is ready to listen to and download.",
+        )
+
     except Exception as exc:
-        logger.error("Song generation failed:\n%s", traceback.format_exc())
-        raise RuntimeError(
-            "Singing failed this time. Please try again." + _debug_detail(exc)
-        ) from exc
+        # Full original traceback in Hugging Face logs.
+        logger.error(
+            "Singing generation failed with full traceback:\n%s",
+            traceback.format_exc(),
+        )
 
-    return audio, "Your sung song is ready to listen to and download."
+        # DEBUG_ERRORS=1 also exposes the underlying exception in Gradio,
+        # which is useful while fixing the Space.
+        raise RuntimeError(
+            "Singing failed."
+            + _debug_detail(exc)
+        ) from exc