OpenMOSS-Team/MOSS-SoundEffect-v2.0

1import functools2import os3import time4 5# Disable Dynamo by default for Space stability; the CLI script does the same.6os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")7 8import gradio as gr9import numpy as np10import torch11 12try:13    import spaces14except ImportError:15    class _SpacesFallback:16        @staticmethod17        def GPU(*_args, **_kwargs):18            def _decorator(func):19                return func20 21            return _decorator22 23    spaces = _SpacesFallback()24 25from moss_soundeffect_v2 import MossSoundEffectPipeline26 27MODEL_PATH = "OpenMOSS-Team/MOSS-SoundEffect-v2.0"28DEFAULT_DEVICE = "cuda"29MAX_INFERENCE_SECONDS = 3030 31 32@functools.lru_cache(maxsize=1)33def load_backend(model_dir: str, device_str: str):34    device = torch.device(device_str if torch.cuda.is_available() else "cpu")35    pipe = MossSoundEffectPipeline.from_pretrained(36        model_dir,37        torch_dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,38        device=str(device),39    )40    return pipe, device41 42 43@spaces.GPU(duration=180)44def run_inference(prompt, seconds, steps, cfg_scale, sigma_shift, seed):45    if not (prompt or "").strip():46        raise ValueError("Please enter a prompt describing the audio you want to generate.")47 48    seconds = round(float(seconds), 1)49    if seconds <= 0:50        raise ValueError("Duration must be greater than 0.")51    if seconds > MAX_INFERENCE_SECONDS:52        raise ValueError(f"Duration must be no greater than {MAX_INFERENCE_SECONDS}s.")53 54    started_at = time.monotonic()55    pipe, _ = load_backend(MODEL_PATH, DEFAULT_DEVICE)56 57    audio = pipe(58        prompt=prompt,59        seconds=seconds,60        num_inference_steps=int(steps),61        cfg_scale=float(cfg_scale),62        sigma_shift=float(sigma_shift),63        seed=int(seed),64    )65 66    audio_np = audio[0].detach().float().cpu().numpy()67    if audio_np.ndim > 1 and audio_np.shape[0] == 1:68        audio_np = audio_np.squeeze(0)69    elif audio_np.ndim > 1:70        audio_np = audio_np.T71    audio_np = audio_np.astype(np.float32, copy=False)72 73    elapsed = time.monotonic() - started_at74    status = (75        f"Done | elapsed: {elapsed:.2f}s | "76        f"duration={seconds:.1f}s, steps={int(steps)}, "77        f"cfg_scale={float(cfg_scale):.2f}, sigma_shift={float(sigma_shift):.2f}, "78        f"seed={int(seed)}"79    )80    return (pipe.sample_rate, audio_np), status81 82 83with gr.Blocks(title="MOSS-SoundEffect v2.0") as demo:84    gr.Markdown(85        """86        # MOSS-SoundEffect v2.087        Text-to-audio diffusion demo.88        """89    )90 91    with gr.Row():92        with gr.Column(scale=3):93            prompt = gr.Textbox(94                label="Prompt",95                lines=8,96                value="The crisp, rhythmic click-clack of fast typing on a mechanical keyboard.",97            )98            seconds = gr.Slider(1, MAX_INFERENCE_SECONDS, step=0.1, value=10, label="Duration (seconds)")99            with gr.Accordion("Sampling Parameters", open=True):100                steps = gr.Slider(10, 150, step=1, value=50, label="num_inference_steps")101                cfg_scale = gr.Slider(1.0, 8.0, step=0.1, value=4.0, label="cfg_scale")102                sigma_shift = gr.Slider(0.0, 10.0, step=0.1, value=5.0, label="sigma_shift")103                seed = gr.Number(value=0, label="seed", precision=0)104            run_btn = gr.Button("Generate Sound Effect", variant="primary")105        with gr.Column(scale=2):106            output_audio = gr.Audio(label="Output Audio", type="numpy")107            status = gr.Textbox(label="Status", lines=4, interactive=False)108 109    run_btn.click(110        fn=run_inference,111        inputs=[prompt, seconds, steps, cfg_scale, sigma_shift, seed],112        outputs=[output_audio, status],113    )114 115demo.queue(max_size=16, default_concurrency_limit=1)116 117if __name__ == "__main__":118    demo.launch()119