Replace Kids Rhyme Studio with clean updated build

audio_engine.py

@@ -0,0 +1,276 @@
+"""Generate a short song with sung vocals from the edited rhyme."""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+import time
+import traceback
+
+# Import spaces before torch/diffusers so ZeroGPU can install its CUDA shim.
+try:
+    import spaces
+except ModuleNotFoundError:
+    if os.environ.get("SPACE_ID"):
+        raise  # A Space needs the ZeroGPU runtime and its spaces module.
+
+    # Local CPU machines can still open the form and write rhymes.
+    class _LocalSpaces:
+        @staticmethod
+        def GPU(**_kwargs):
+            return lambda fn: fn
+
+    spaces = _LocalSpaces()
+
+import numpy as np
+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"
+
+# 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
+
+if torch is not None and torch.cuda.is_available():
+    try:
+        from diffusers import AceStepPipeline
+
+        # 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())
+
+
+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]}]"
+
+
+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
+        raise RuntimeError(
+            "Singing needs GPU hardware. Set the Space hardware to ZeroGPU."
+        )
+    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,
+        )
+        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)
+
+
+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()]
+    if len(stanzas) >= 2:
+        first = stanzas[0]
+        second = "\n".join(stanzas[1:])
+    else:
+        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.")
+        midpoint = (len(lines) + 1) // 2
+        first, second = "\n".join(lines[:midpoint]), "\n".join(lines[midpoint:])
+    return f"[verse]\n{first}\n[chorus]\n{second}"
+
+
+def _song_settings(
+    text: str,
+    language: str,
+    mood: str,
+    theme: str = "",
+    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 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."
+        )
+        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."
+        )
+        bpm = 112
+
+    language_name = LANGUAGES[language]["name"]
+    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())
+    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)),
+        "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]."""
+    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"):
+            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__}.")
+
+    if hasattr(value, "detach"):
+        value = value.detach().float().cpu().numpy()
+
+    arr = np.asarray(value, dtype=np.float32)
+
+    # Typical result: [batch, channels, samples]. Take first batch.
+    while arr.ndim > 2:
+        arr = arr[0]
+    if arr.ndim == 1:
+        arr = arr[None, :]
+    if arr.ndim != 2:
+        raise ValueError(f"Unexpected audio rank/shape {arr.shape}.")
+
+    # If it came back [samples, channels], transpose it.
+    if arr.shape[0] > 2 and arr.shape[1] in (1, 2):
+        arr = arr.T
+
+    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)
+    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)
+    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()
+
+    result = pipe(**settings)
+    waveform = _extract_audio(result)
+    sr = _sample_rate()
+
+    logger.info("ACE-Step returned audio shape %s at %s Hz", waveform.shape, sr)
+
+    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")
+
+    waveform = _polish(waveform, sr)
+    pcm = (np.clip(waveform.T, -1.0, 1.0) * 32767).astype(np.int16)
+
+    logger.info(
+        "Generated %.1fs of audio at %s Hz in %.1fs",
+        len(pcm) / sr,
+        sr,
+        time.monotonic() - started,
+    )
+    return sr, pcm
+
+
+@spaces.GPU(duration=120)
+def _generate_on_gpu(settings: dict) -> tuple[int, np.ndarray]:
+    return generate_song(settings)
+
+
+def make_sung_song(
+    text: str,
+    language: str,
+    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()
+
+    try:
+        audio = _generate_on_gpu(settings)
+    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
+
+    return audio, "Your sung song is ready to listen to and download."