Files
audio-demo/tts/generate_sfx.py
steve-adminandClaude Sonnet 5 f4e5b9e20a Add MOSS-SoundEffect generation (CLI + webui section)
New tts/generate_sfx.py mirrors generate.py (MOSS-TTS) but for the
MOSS-SoundEffect checkpoint: text-only ambience/SFX generation, no
reference audio or voice cloning. Needed PYTORCH_ALLOC_CONF=
expandable_segments:True baked in - this checkpoints weights alone use
~22.3GB of this boxs ~23.56GB usable VRAM, leaving generate() itself no
room without it (confirmed via OOM otherwise).

webui/app.py gets a new standalone "5. Sound effects" section in the
Pipeline Studio tab (separate from the AuK/MOSS-TTS voice flow - this is
a different generation primitive with no reference/history state).
Verified end-to-end: three real test clips (door creak, market ambience,
car pass) generated cleanly, approved by Steve as "definitely passable,"
then wired into the GUI and confirmed working through app.py itself.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01SN15Vxn4svKuhYxFsVtMJf
2026-10-09 14:57:39 +00:00

100 lines
3.7 KiB
Python

"""
MOSS-SoundEffect CLI: generate an ambient/sound-effect clip from a text
description. Sibling of generate.py (MOSS-TTS voice cloning) - same
backend setup and model-loading pattern, different checkpoint and a
text-only build_user_message call (no reference audio, no cloning).
"""
import argparse
import importlib.util
import os
# This checkpoint's weights alone use ~22.3GB of this box's ~23.56GB usable
# VRAM, leaving generate()'s own small allocations (a vocab-sized mask
# tensor, etc.) no room without this - confirmed via OOM otherwise. Must be
# set before torch's CUDA allocator initializes, so it's set here rather
# than relying on the caller's shell environment.
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import soundfile as sf
import torch
from transformers import AutoModel, AutoProcessor
PRETRAINED = "OpenMOSS-Team/MOSS-SoundEffect"
def resolve_attn_implementation(device: str, dtype: torch.dtype) -> str:
if (
device == "cuda"
and importlib.util.find_spec("flash_attn") is not None
and dtype in {torch.float16, torch.bfloat16}
):
major, _ = torch.cuda.get_device_capability()
if major >= 8:
return "flash_attention_2"
return "sdpa" if device == "cuda" else "eager"
def main() -> None:
parser = argparse.ArgumentParser(description="MOSS-SoundEffect text-to-sound-effect")
parser.add_argument("--text", required=True, help="Description of the sound effect/ambience")
parser.add_argument("--output", required=True, help="Output wav path")
parser.add_argument(
"--seed", type=int, default=None,
help="Random seed for reproducibility. Unset = nondeterministic.",
)
parser.add_argument("--max_new_tokens", type=int, default=4096)
args = parser.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
if args.seed is not None:
torch.manual_seed(args.seed)
if device == "cuda":
torch.cuda.manual_seed_all(args.seed)
torch.backends.cuda.enable_cudnn_sdp(False)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)
attn_implementation = resolve_attn_implementation(device, dtype)
print(f"[INFO] attn_implementation={attn_implementation}")
processor = AutoProcessor.from_pretrained(PRETRAINED, trust_remote_code=True)
processor.audio_tokenizer = processor.audio_tokenizer.to(device)
model = AutoModel.from_pretrained(
PRETRAINED,
trust_remote_code=True,
attn_implementation=attn_implementation,
dtype=dtype,
).to(device)
model.eval()
conversation = [processor.build_user_message(ambient_sound=args.text)]
with torch.no_grad():
batch = processor([conversation], mode="generation")
# This checkpoint's generate() doesn't accept the TTS model's sampling
# kwargs (do_sample/audio_temperature/etc raise TypeError) - matches
# the official README example, which only passes these three.
outputs = model.generate(
input_ids=batch["input_ids"].to(device),
attention_mask=batch["attention_mask"].to(device),
max_new_tokens=args.max_new_tokens,
)
message = processor.decode(outputs)[0]
audio = message.audio_codes_list[0]
data = audio.detach().cpu().to(torch.float32).numpy()
if data.ndim == 2:
data = data.T
sf.write(args.output, data, processor.model_config.sampling_rate)
print(f"[INFO] Saved {args.output}")
if __name__ == "__main__":
main()