peitho_model.py

"""Boot-safe Peitho model loader for ZeroGPU.

ZeroGPU has no GPU outside an `@spaces.GPU` call, and it cannot defer a
*direct* CUDA allocation (e.g. `safetensors.load_file(device="cuda")` or
`torch.empty(device="cuda")`) — those crash at import with "No CUDA GPUs are
available". So we load the processor, base model, and Mimi codec entirely on
CPU here. The actual move to GPU happens later, inside the `@spaces.GPU`
handler in `chat.py`, via the patchable `.to("cuda")` path.
"""
from __future__ import annotations

import os

# moshi/Mimi compiles its decoder with torch.compile + CUDA graphs, which clash
# with ZeroGPU's torch tensor patching (Dynamo can't treat the patched
# torch.device as a constant). Disable both BEFORE importing liquid_audio.
os.environ.setdefault("NO_TORCH_COMPILE", "1")
os.environ.setdefault("NO_CUDA_GRAPH", "1")

import spaces  # noqa: E402,F401  # imported before torch CUDA usage for ZeroGPU
import torch  # noqa: E402,F401
from accelerate import load_checkpoint_in_model  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from liquid_audio import LFM2AudioModel, LFM2AudioProcessor  # noqa: E402

BASE_REPO: str = "LiquidAI/LFM2.5-Audio-1.5B"
PEITHO_REPO: str = "jempf/peitho-1.5b-v6"

print("Loading processor + base LFM2.5-Audio-1.5B on CPU (ZeroGPU-safe)...")
proc = LFM2AudioProcessor.from_pretrained(BASE_REPO, device="cpu").eval()
lfm2_audio = LFM2AudioModel.from_pretrained(BASE_REPO).to("cpu").eval()
mimi = proc.mimi.eval()

print(f"Overlaying German v6 weights from {PEITHO_REPO}...")
_weights_dir = snapshot_download(
    repo_id=PEITHO_REPO,
    allow_patterns=["model.safetensors", "config.json"],
)
load_checkpoint_in_model(lfm2_audio, _weights_dir)
lfm2_audio.eval()
print("Peitho v6 ready on CPU; GPU placement happens inside @spaces.GPU.")