Initial Kids Rhyme Studio setup

audio_engine.py

@@ -0,0 +1,183 @@
+"""Generate a short song with sung vocals from the edited rhyme."""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+import time
+
+# 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()}
+# Released ACE-Step checkpoints use 48 kHz; prefer the pipeline's own value.
+# https://huggingface.co/docs/diffusers/api/pipelines/ace_step
+FALLBACK_SAMPLE_RATE = 48000
+
+# ZeroGPU expects the model to be placed on CUDA at module import time. Its
+# CUDA shim makes this possible even before a real GPU has been allocated.
+# On a local machine without CUDA, lyric writing still works and the song
+# button explains what is needed.
+_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
+
+        _pipe = AceStepPipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda")
+        logger.info("Singing model loaded: %s", MODEL_ID)
+    except Exception as exc:
+        _load_error = exc
+        logger.exception("Singing model failed to load")
+
+
+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.") from _load_error
+        raise RuntimeError("Singing needs GPU hardware. Set the Space hardware to ZeroGPU.")
+    return _pipe
+
+
+def _sample_rate() -> int:
+    rate = getattr(_require_pipeline(), "sample_rate", None)
+    if rate is None:
+        logger.warning("Pipeline has no sample_rate; 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("The pipeline returned an invalid sample rate.")
+    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}."
+    return dict(
+        prompt=prompt,
+        lyrics=_song_lyrics(text),
+        vocal_language=VOCAL_LANGUAGE[language],
+        # Give longer rhymes time for every line: 40 seconds for 8 lines,
+        # up to 56 seconds for 12 lines when all ten words are supplied.
+        audio_duration=min(56.0, 40.0 + 4.0 * max(0, sum(bool(line.strip()) for line in text.splitlines()) - 8)),
+        num_inference_steps=8,
+        bpm=bpm,
+        task_type="text2music",
+        output_type="np",
+    )
+
+
+def _polish(wave: np.ndarray, sr: int) -> np.ndarray:
+    """Peak-normalize and fade the tail; this cannot recover omitted lyrics."""
+    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 without a UI decorator; caller must provide GPU access.
+
+    A future API server needs CUDA hardware or the ZeroGPU wrapper below.
+    This function is not a CPU inference path.
+    """
+    pipe = _require_pipeline()
+    sr = _sample_rate()
+    started = time.monotonic()
+    result = pipe(**settings)
+    # ACE-Step returns [batch, channels, samples]; Gradio takes [samples, channels].
+    waveform = np.asarray(result.audios[0], dtype=np.float32)
+    if waveform.ndim != 2 or waveform.shape[0] not in (1, 2) or waveform.shape[1] < sr:
+        raise ValueError(f"Unexpected audio shape {waveform.shape}")
+    if not np.isfinite(waveform).all() or 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)  # Keep this budget until real deployed runs have been timed.
+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 before requesting a GPU; retain 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.exception("Song generation failed")
+        raise RuntimeError("Singing failed this time. Please try again.") from exc
+    return audio, "Your sung song is ready to listen to and download."