Spaces:
Running
Running
| """ | |
| Woosh-DFlow on HARP: four-step text-to-audio sound effect generation. | |
| The model runs in this Space. It used to run in a second Space that this one | |
| called over the Gradio client, because Woosh needs Gradio 6 and PyHARP pinned | |
| Gradio 5.28 at the time. PyHARP v0.4.1 requires Gradio 6 itself, so the two | |
| now share one process and a request no longer depends on another Space being | |
| awake. | |
| Weights are the official v1.0.0 release, CC-BY-NC 4.0, baked into the image by | |
| the Dockerfile from the teamup-tech/Woosh-weights mirror. | |
| """ | |
| import os | |
| import tempfile | |
| import gradio as gr | |
| import numpy as np | |
| import soundfile as sf | |
| import torch | |
| from pyharp import (Channels, ModelCard, SampleRate, Subcategory, | |
| build_endpoint) | |
| from woosh.components.base import LoadConfig | |
| from woosh.inference.flowmap_sampler import sample_euler | |
| from woosh.model.flowmap_from_pretrained import FlowMapFromPretrained | |
| # Matches the upstream demo: 128-channel latents, ~5s at 48 kHz after the AE. | |
| SAMPLE_RATE = 48000 | |
| CHANNELS = 1 | |
| LATENT_CHANNELS = 128 | |
| LATENT_FRAMES = 501 | |
| # Woosh-DFlow is the distilled model, which is what makes four steps enough. | |
| # The renoise schedule is the one the authors distilled against, so it has a | |
| # step for each and is not a knob to expose. | |
| NUM_STEPS = 4 | |
| RENOISE = [0, 0.5, 0.5, 0.3] | |
| CHECKPOINT = os.environ.get("WOOSH_CHECKPOINT", "checkpoints/Woosh-DFlow") | |
| # ZeroGPU is offered to Gradio SDK Spaces only, so there is no @spaces.GPU to | |
| # apply here. Swapping the torch wheel index in the Dockerfile is what moves | |
| # this Space to a dedicated GPU. | |
| if torch.cuda.is_available(): | |
| device = "cuda" | |
| elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available(): | |
| device = "mps" | |
| else: | |
| device = "cpu" | |
| model_card = ModelCard( | |
| name="Woosh", | |
| description=( | |
| "Generates a five-second, 48 kHz sound effect from a text prompt. " | |
| "Woosh-DFlow is the distilled member of Sony AI's Woosh family: it " | |
| "reaches a result in four sampling steps where the undistilled models " | |
| "take tens, which is what makes it quick enough to audition ideas " | |
| "against. It writes a single effect rather than a scene, so a prompt " | |
| "naming one source and one action, such as \"sportscar engine revving " | |
| "and driving away\", lands closer than a description of a whole " | |
| "environment. The weights are released under CC-BY-NC 4.0, which " | |
| "permits non-commercial use only. Repository: " | |
| "https://github.com/SonyResearch/Woosh Paper: " | |
| "https://arxiv.org/abs/2604.01929" | |
| ), | |
| author=("Gaëtan Hadjeres, Marc Ferras, Khaled Koutini, Benno Weck, " | |
| "Alexandre Bittar, Thomas Hummel, Zineb Lahrichi, Hakim Missoum, " | |
| "Joan Serrà, Yuki Mitsufuji"), | |
| tags=[Subcategory.HOLISTIC, SampleRate(SAMPLE_RATE), Channels(CHANNELS), | |
| "sound-effects"], | |
| ) | |
| # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the | |
| # server, so the model is loaded there alone | |
| if __name__ != "__main__": | |
| model = FlowMapFromPretrained(LoadConfig(path=CHECKPOINT)).eval().to(device) | |
| def process_fn(prompt, cfg_scale, seed): | |
| """ | |
| Generate one sound effect from a text prompt. | |
| Args: | |
| prompt (str): What the effect should be. | |
| cfg_scale (float): How strongly the prompt steers sampling. | |
| seed (float): Seed for the noise, or negative to draw a fresh one. | |
| Returns: | |
| output_path (str): The generated audio. | |
| """ | |
| if not (prompt or "").strip(): | |
| raise gr.Error("Describe the sound effect to generate.") | |
| if seed is None or int(seed) < 0: | |
| seed = int.from_bytes(os.urandom(4), "big") % (2 ** 31) | |
| torch.manual_seed(int(seed)) | |
| noise = torch.randn(1, LATENT_CHANNELS, LATENT_FRAMES).to(device) | |
| cond = model.get_cond({"audio": None, "description": [prompt.strip()]}, | |
| no_dropout=True, device=device) | |
| sampled = sample_euler(model=model, noise=noise, cond=cond, | |
| num_steps=NUM_STEPS, renoise=RENOISE, | |
| cfg=float(cfg_scale)) | |
| audio = model.autoencoder.inverse(sampled).cpu().float() | |
| # Peak-normalize only where the result would clip, so a quiet effect keeps | |
| # its level rather than being pushed up to full scale | |
| peak = audio.abs().amax(dim=-1, keepdim=True).clamp(min=1.0) | |
| audio = (audio / peak).clamp(-1.0, 1.0) | |
| output_path = os.path.join(tempfile.mkdtemp(), "woosh_dflow.wav") | |
| sf.write(output_path, np.ascontiguousarray(audio[0].squeeze().numpy()), | |
| SAMPLE_RATE) | |
| return output_path | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| input_components = [ | |
| gr.Textbox(label="Prompt", | |
| value="sportscar engine revving and driving away quickly", | |
| info=( | |
| "The effect to generate. Name the source and what " | |
| "it does. The model writes one effect at a time, so " | |
| "a single action works better than a scene with " | |
| "several things happening at once." | |
| )), | |
| gr.Slider(minimum=0.0, maximum=15.0, step=0.1, value=4.5, | |
| label="Prompt strength", | |
| info=( | |
| "How closely sampling follows the prompt. Low values " | |
| "wander off it, and high values follow the wording so " | |
| "literally that the result can turn harsh." | |
| )), | |
| gr.Number(value=-1, precision=0, label="Seed", | |
| info=( | |
| "Seeds the noise the four steps start from, so the " | |
| "same seed and prompt reproduce a result. Leave it " | |
| "negative to draw a new one each time." | |
| )), | |
| ] | |
| output_components = [ | |
| gr.Audio(type="filepath", label="Generated audio").set_info(( | |
| "The generated effect, five seconds at 48 kHz." | |
| )), | |
| ] | |
| build_endpoint( | |
| model_card=model_card, | |
| input_components=input_components, | |
| output_components=output_components, | |
| process_fn=process_fn, | |
| ) | |
| demo.queue().launch(server_name="0.0.0.0", | |
| server_port=int(os.environ.get("PORT", "7860")), | |
| show_error=True, pwa=True) | |