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
This commit is contained in:
steve-adminandClaude Sonnet 5 committed
1 parent
f038ba609d
commit
f4e5b9e20a
2 files changed
+143
No files matched your search
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
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()
|
||||
@@ -49,6 +49,7 @@ AUK_FLASH_CKPT = AUK_DIR / "ckpts" / "AuK-Flash" / "auk_flash.safetensors"
|
||||
AUK_QWEN_PATH = AUK_DIR / "ckpts" / "Qwen2.5-Omni-3B"
|
||||
TTS_PYTHON = TTS_DIR / ".venv" / "bin" / "python3"
|
||||
TTS_GENERATE = TTS_DIR / "generate.py"
|
||||
TTS_GENERATE_SFX = TTS_DIR / "generate_sfx.py"
|
||||
|
||||
MODEL_BASE = "Base (full quality)"
|
||||
MODEL_FLASH = "Flash (fast draft, 4-step)"
|
||||
@@ -418,6 +419,30 @@ def generate_speech(reference: str, text: str, language: str):
|
||||
return str(out_path), transcript
|
||||
|
||||
|
||||
def generate_sound_effect(description: str):
|
||||
"""MOSS-SoundEffect - text-only ambience/SFX generation, no reference
|
||||
audio or character voice involved (different generation primitive from
|
||||
the AuK/MOSS-TTS voice-cloning flow above, so it's a standalone section
|
||||
rather than wired into that pipeline's history/reference state).
|
||||
No whisper verification here - there's no speech content to check."""
|
||||
if not description or not description.strip():
|
||||
raise gr.Error("Enter a description of the sound effect/ambience.")
|
||||
|
||||
_free_comfyui_memory()
|
||||
out_path = WORKDIR / f"sfx_{int(time.time())}.wav"
|
||||
env = None
|
||||
if TTS_LD_LIBRARY_PATH:
|
||||
env = os.environ.copy()
|
||||
env["LD_LIBRARY_PATH"] = TTS_LD_LIBRARY_PATH
|
||||
|
||||
_run(
|
||||
[str(TTS_PYTHON), str(TTS_GENERATE_SFX), "--text", description.strip(), "--output", str(out_path)],
|
||||
cwd=TTS_DIR,
|
||||
env=env,
|
||||
)
|
||||
return str(out_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Characters (the-service-station cast)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -781,6 +806,24 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
|
||||
save_char_btn = gr.Button("Save as new character (status: draft)")
|
||||
save_char_status = gr.Markdown()
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
gr.Markdown(
|
||||
"## 5. Sound effects (MOSS-SoundEffect)\n"
|
||||
"Text-only ambience/SFX generation - no reference audio or character "
|
||||
"voice involved, a separate model from the voice pipeline above."
|
||||
)
|
||||
sfx_description = gr.Textbox(
|
||||
label="Description",
|
||||
placeholder=(
|
||||
'e.g. "A door creaking open slowly, then closing with a soft thud." '
|
||||
'or "A busy outdoor market with distant chatter and footsteps."'
|
||||
),
|
||||
lines=2,
|
||||
)
|
||||
sfx_btn = gr.Button("Generate sound effect", variant="primary")
|
||||
sfx_audio = gr.Audio(label="Generated sound effect", type="filepath")
|
||||
|
||||
strip_btn.click(strip_vocals, inputs=[input_audio], outputs=[vocals_out, no_vocals_out])
|
||||
use_vocals_btn.click(
|
||||
use_as_reference, inputs=[vocals_out, history_state],
|
||||
@@ -819,6 +862,7 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
|
||||
inputs=[new_char_name, new_char_role, new_char_prose, gen_language, current_ref_audio, history_state],
|
||||
outputs=[save_char_status],
|
||||
)
|
||||
sfx_btn.click(generate_sound_effect, inputs=[sfx_description], outputs=[sfx_audio])
|
||||
|
||||
with gr.Tab("Characters"):
|
||||
gr.Markdown(
|
||||
|
||||
Reference in new issue
Block a user