boatmanandsphinxkkkbc
a0918ce583
Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com> Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>
331 lines
11 KiB
Python
331 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Example script for text-to-audio generation using Stable Audio Open.
|
|
|
|
This script demonstrates how to generate audio from text prompts using
|
|
the Stable Audio Open model with vLLM-Omni.
|
|
|
|
Usage:
|
|
python text_to_audio.py --prompt "The sound of a dog barking"
|
|
python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
|
|
python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
|
|
python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache
|
|
"""
|
|
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Generate audio with Stable Audio Open.")
|
|
parser.add_argument(
|
|
"--model",
|
|
default="stabilityai/stable-audio-open-1.0",
|
|
help="Stable Audio model name or local path.",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt",
|
|
default="The sound of a hammer hitting a wooden surface.",
|
|
help="Text prompt for audio generation.",
|
|
)
|
|
parser.add_argument(
|
|
"--negative-prompt",
|
|
default="Low quality.",
|
|
help="Negative prompt for classifier-free guidance.",
|
|
)
|
|
parser.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
default=42,
|
|
help="Random seed for deterministic results.",
|
|
)
|
|
parser.add_argument(
|
|
"--guidance-scale",
|
|
type=float,
|
|
default=7.0,
|
|
help="Classifier-free guidance scale.",
|
|
)
|
|
parser.add_argument(
|
|
"--audio-start",
|
|
type=float,
|
|
default=0.0,
|
|
help="Audio start time in seconds.",
|
|
)
|
|
parser.add_argument(
|
|
"--audio-length",
|
|
type=float,
|
|
default=10.0,
|
|
help="Audio length in seconds (max ~47s for stable-audio-open-1.0).",
|
|
)
|
|
parser.add_argument(
|
|
"--num-inference-steps",
|
|
type=int,
|
|
default=100,
|
|
help="Number of denoising steps for the diffusion sampler.",
|
|
)
|
|
parser.add_argument(
|
|
"--num-waveforms",
|
|
type=int,
|
|
default=1,
|
|
help="Number of audio waveforms to generate for the given prompt.",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=str,
|
|
default="stable_audio_output.wav",
|
|
help="Path to save the generated audio (WAV format).",
|
|
)
|
|
parser.add_argument(
|
|
"--sample-rate",
|
|
type=int,
|
|
default=44100,
|
|
help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
|
|
)
|
|
parser.add_argument(
|
|
"--cache-backend",
|
|
type=str,
|
|
default=None,
|
|
choices=["tea_cache"],
|
|
help=(
|
|
"Cache backend to use for acceleration. "
|
|
"Stable Audio currently supports 'tea_cache'. "
|
|
"Default: None (no cache acceleration)."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--tea-cache-rel-l1-thresh",
|
|
type=float,
|
|
default=0.2,
|
|
help="[tea_cache] Threshold for accumulated relative L1 distance.",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-diffusion-pipeline-profiler",
|
|
action="store_true",
|
|
help="Enable diffusion pipeline profiler to display stage durations.",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-cpu-offload", action="store_true", help="Enable ModelWise cpu offloading to save gpu memory"
|
|
)
|
|
parser.add_argument(
|
|
"--enable-layerwise-offload", action="store_true", help="Enable Layerwise cpu offloading to save gpu memory"
|
|
)
|
|
parser.add_argument(
|
|
"--use-hsdp",
|
|
action="store_true",
|
|
help="Enable HSDP for Stable Audio DiT weight sharding.",
|
|
)
|
|
parser.add_argument(
|
|
"--hsdp-shard-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs to shard Stable Audio DiT weights across when HSDP is enabled.",
|
|
)
|
|
parser.add_argument(
|
|
"--hsdp-replicate-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of HSDP replica groups. Default 1 means pure sharding.",
|
|
)
|
|
parser.add_argument(
|
|
"--tensor-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
|
)
|
|
parser.add_argument(
|
|
"--ulysses-degree",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for ulysses sequence parallelism.",
|
|
)
|
|
parser.add_argument(
|
|
"--ulysses-mode",
|
|
type=str,
|
|
default="strict",
|
|
choices=["strict", "advanced_uaa"],
|
|
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
|
|
)
|
|
parser.add_argument(
|
|
"--ring-degree",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for ring sequence parallelism.",
|
|
)
|
|
parser.add_argument(
|
|
"--cfg-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
choices=[1, 2],
|
|
help="Number of GPUs used for classifier free guidance parallel size.",
|
|
)
|
|
parser.add_argument(
|
|
"--vae-patch-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 44100):
|
|
"""Save audio data to a WAV file."""
|
|
try:
|
|
import soundfile as sf
|
|
|
|
sf.write(output_path, audio_data, sample_rate)
|
|
except ImportError:
|
|
try:
|
|
import scipy.io.wavfile as wav
|
|
|
|
# Ensure audio is in the correct format for scipy
|
|
if audio_data.dtype == np.float32 or audio_data.dtype == np.float64:
|
|
# Normalize to int16 range
|
|
audio_data = np.clip(audio_data, -1.0, 1.0)
|
|
audio_data = (audio_data * 32767).astype(np.int16)
|
|
wav.write(output_path, sample_rate, audio_data)
|
|
except ImportError:
|
|
raise ImportError(
|
|
"Either 'soundfile' or 'scipy' is required to save audio files. "
|
|
"Install with: pip install soundfile or pip install scipy"
|
|
)
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
|
cache_config = None
|
|
if args.cache_backend == "tea_cache":
|
|
cache_config = {
|
|
"rel_l1_thresh": args.tea_cache_rel_l1_thresh,
|
|
}
|
|
|
|
print(f"\n{'=' * 60}")
|
|
print("Stable Audio Open - Text-to-Audio Generation")
|
|
print(f"{'=' * 60}")
|
|
print(f" Model: {args.model}")
|
|
print(f" Prompt: {args.prompt}")
|
|
print(f" Negative prompt: {args.negative_prompt}")
|
|
print(f" Audio length: {args.audio_length}s")
|
|
print(f" Inference steps: {args.num_inference_steps}")
|
|
print(f" Guidance scale: {args.guidance_scale}")
|
|
print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
|
|
print(f" ModelWise Offload: {'Enabled' if args.enable_cpu_offload else 'None'}")
|
|
print(f" LayerWise Offload: {'Enabled' if args.enable_layerwise_offload else 'None'}")
|
|
if args.use_hsdp:
|
|
print(f" HSDP: enabled (shard_size={args.hsdp_shard_size}, replicate_size={args.hsdp_replicate_size})")
|
|
else:
|
|
print(" HSDP: disabled")
|
|
print(f" Seed: {args.seed}")
|
|
print(f"{'=' * 60}\n")
|
|
|
|
parallel_config = DiffusionParallelConfig(
|
|
use_hsdp=args.use_hsdp,
|
|
hsdp_shard_size=args.hsdp_shard_size,
|
|
hsdp_replicate_size=args.hsdp_replicate_size,
|
|
)
|
|
|
|
# Initialize Omni with Stable Audio model
|
|
omni = Omni(
|
|
model=args.model,
|
|
parallel_config=parallel_config,
|
|
cache_backend=args.cache_backend,
|
|
cache_config=cache_config,
|
|
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
|
enable_cpu_offload=args.enable_cpu_offload,
|
|
enable_layerwise_offload=args.enable_layerwise_offload,
|
|
)
|
|
|
|
# Calculate audio end time
|
|
audio_end_in_s = args.audio_start + args.audio_length
|
|
|
|
# Time profiling for generation
|
|
generation_start = time.perf_counter()
|
|
|
|
# Generate audio
|
|
outputs = omni.generate(
|
|
{
|
|
"prompt": args.prompt,
|
|
"negative_prompt": args.negative_prompt,
|
|
},
|
|
OmniDiffusionSamplingParams(
|
|
generator=generator,
|
|
guidance_scale=args.guidance_scale,
|
|
num_inference_steps=args.num_inference_steps,
|
|
num_outputs_per_prompt=args.num_waveforms,
|
|
extra_args={
|
|
"audio_start_in_s": args.audio_start,
|
|
"audio_end_in_s": audio_end_in_s,
|
|
},
|
|
),
|
|
)
|
|
|
|
generation_end = time.perf_counter()
|
|
generation_time = generation_end - generation_start
|
|
|
|
print(f"Total generation time: {generation_time:.2f} seconds")
|
|
|
|
# Process and save audio
|
|
output_path = Path(args.output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
suffix = output_path.suffix or ".wav"
|
|
stem = output_path.stem or "stable_audio_output"
|
|
|
|
# Extract audio from omni.generate() outputs
|
|
if not outputs:
|
|
raise ValueError("No output generated from omni.generate()")
|
|
|
|
output = outputs[0]
|
|
if not hasattr(output, "request_output") or not output.request_output:
|
|
raise ValueError("No request_output found in OmniRequestOutput")
|
|
request_output = output.request_output
|
|
if not hasattr(request_output, "multimodal_output"):
|
|
raise ValueError("No multimodal_output found in request_output")
|
|
|
|
audio = request_output.multimodal_output.get("audio")
|
|
if audio is None:
|
|
raise ValueError("No audio output found in request_output")
|
|
|
|
# Handle different output formats
|
|
if isinstance(audio, torch.Tensor):
|
|
audio = audio.cpu().float().numpy()
|
|
|
|
# Audio shape is typically [batch, channels, samples] or [channels, samples]
|
|
if audio.ndim == 3:
|
|
# [batch, channels, samples]
|
|
if args.num_waveforms <= 1:
|
|
audio_data = audio[0].T # [samples, channels]
|
|
save_audio(audio_data, str(output_path), args.sample_rate)
|
|
print(f"Saved generated audio to {output_path}")
|
|
else:
|
|
for idx in range(audio.shape[0]):
|
|
audio_data = audio[idx].T # [samples, channels]
|
|
save_path = output_path.parent / f"{stem}_{idx}{suffix}"
|
|
save_audio(audio_data, str(save_path), args.sample_rate)
|
|
print(f"Saved generated audio to {save_path}")
|
|
elif audio.ndim == 2:
|
|
# [channels, samples]
|
|
audio_data = audio.T # [samples, channels]
|
|
save_audio(audio_data, str(output_path), args.sample_rate)
|
|
print(f"Saved generated audio to {output_path}")
|
|
else:
|
|
# [samples] - mono audio
|
|
save_audio(audio, str(output_path), args.sample_rate)
|
|
print(f"Saved generated audio to {output_path}")
|
|
|
|
print(f"\nGenerated {args.audio_length}s of audio at {args.sample_rate} Hz")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|