Spaces:
Paused
Paused
File size: 4,893 Bytes
c5c143f 23df7cf c5c143f 23df7cf c5c143f 93c7079 c5c143f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | from __future__ import annotations
import subprocess
import sys
import tempfile
import threading
from pathlib import Path
MODEL_REPO = "stabilityai/stable-audio-3-small-music"
MODEL_REVISION = "0fef1392cd842149a2b6d445e181c97608faac06"
STABLE_AUDIO_TOOLS_REVISION = "3241adba4fc2a85cf5b29d9eb68d42f40a28e820"
OUTPUT_ROOT = Path(tempfile.gettempdir()) / "stable_audio_3_outputs"
_MODEL = None
_MODEL_CONFIG = None
_MODEL_LOCK = threading.Lock()
def _ensure_stable_audio_tools() -> None:
try:
import stable_audio_tools # noqa: F401
return
except ImportError:
pass
subprocess.check_call(
[
sys.executable,
"-m",
"pip",
"install",
"--quiet",
"--no-deps",
(
"git+https://github.com/Stability-AI/stable-audio-tools.git@"
f"{STABLE_AUDIO_TOOLS_REVISION}"
),
]
)
_ensure_stable_audio_tools()
def _load_model():
global _MODEL, _MODEL_CONFIG
if _MODEL is not None:
return _MODEL, _MODEL_CONFIG
with _MODEL_LOCK:
if _MODEL is not None:
return _MODEL, _MODEL_CONFIG
import torch
from stable_audio_tools.models import pretrained
original_download = pretrained.hf_hub_download
def pinned_download(repo_id, *args, **kwargs):
if repo_id == MODEL_REPO:
kwargs.setdefault("revision", MODEL_REVISION)
return original_download(repo_id, *args, **kwargs)
pretrained.hf_hub_download = pinned_download
try:
model, config = pretrained.get_pretrained_model(MODEL_REPO)
finally:
pretrained.hf_hub_download = original_download
model = model.to("cuda").to(torch.float16)
model.eval().requires_grad_(False)
_MODEL = model
_MODEL_CONFIG = config
return _MODEL, _MODEL_CONFIG
def _load_audio(path: str):
import torch
import torchaudio
audio, sample_rate = torchaudio.load(path)
if audio.shape[0] > 2:
audio = audio[:2]
return int(sample_rate), audio.to(torch.float32)
def _save_audio(output, sample_rate: int) -> str:
import soundfile as sf
import torch
output = output.permute(1, 0, 2).reshape(output.shape[1], -1)
output = output.to(torch.float32)
peak = output.abs().max().clamp(min=1e-9)
output = output.div(peak).clamp(-1, 1).cpu().numpy().T
OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
output_dir = Path(tempfile.mkdtemp(dir=OUTPUT_ROOT))
output_path = output_dir / "stable_audio_3_edit.wav"
sf.write(output_path, output, sample_rate, subtype="PCM_16")
return str(output_path)
def edit_audio(
audio_path: str,
prompt: str,
mode: str,
edit_start: float,
edit_end: float,
continuation_length: float,
strength: float,
seed: int,
) -> str:
import torch
import torchaudio
model, config = _load_model()
from stable_audio_tools.inference.generation import (
generate_diffusion_cond_inpaint,
)
sample_rate = int(config["sample_rate"])
sample_size = int(config["sample_size"])
source_rate, source = _load_audio(audio_path)
source_duration = source.shape[-1] / source_rate
if source_rate != sample_rate:
source = torchaudio.functional.resample(source, source_rate, sample_rate)
model_dtype = next(model.parameters()).dtype
source_tuple = (sample_rate, source.to(model_dtype))
if mode == "Continue":
output_duration = source_duration + continuation_length
else:
output_duration = source_duration
conditioning = [{"prompt": prompt, "seconds_total": output_duration}]
kwargs = {
"steps": 8,
"cfg_scale": 1.0,
"conditioning": conditioning,
"sample_size": sample_size,
"sampler_type": "pingpong",
"seed": int(seed) if seed > 0 else -1,
"device": "cuda",
"sigma_max": 1.0,
"apg_scale": 1.0,
"duration_padding_sec": 6.0,
}
if mode == "Restyle":
kwargs["init_audio"] = source_tuple
kwargs["init_noise_level"] = float(strength)
elif mode == "Inpaint":
kwargs["inpaint_audio"] = source_tuple
kwargs["inpaint_mask_start_seconds"] = float(edit_start)
kwargs["inpaint_mask_end_seconds"] = float(edit_end)
elif mode == "Continue":
kwargs["inpaint_audio"] = source_tuple
kwargs["inpaint_mask_start_seconds"] = float(source_duration)
kwargs["inpaint_mask_end_seconds"] = float(output_duration)
else:
raise ValueError(f"Unknown edit mode: {mode}")
with torch.inference_mode():
output = generate_diffusion_cond_inpaint(model, **kwargs)
output = output[..., : int(output_duration * sample_rate)]
return _save_audio(output, sample_rate)
|