tts/generate_sfx.py

@@ -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()

webui/app.py

@@ -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(