import torch
import torchaudio
from einops import rearrange
from stable_audio_tools import get_pretrained_model
from omegaconf import OmegaConf
from stable_audio_tools.models.factory import create_model_from_config
from stable_audio_tools.inference.generation import generate_diffusion_cond
from safetensors.torch import load_file as load_safetensors
from pydub import AudioSegment
import re
import os
from datetime import datetime
import gradio as gr
# Define a function to toggle the visibility of the seed slider
def toggle_seed_slider(x):
return gr.Slider(interactive=not x)
# Define a function to set up the model and device
def setup_model(model_path, model_half):
"""
Sets up the model and device.
Args:
model_path (str): Path to a local model .ckpt or .safetensors file. If empty, downloads the default model.
model_half (bool): Whether to use float16 half-precision.
"""
device = "cuda" if torch.cuda.is_available() else "cpu"
# If no path is provided, or path doesn't exist, download the default model
if not model_path or not os.path.exists(model_path):
if model_path:
print(f"Warning: Model path '{model_path}' not found. Falling back to default model.")
model_id = "audo/stable-audio-open-1.0"
print(f"Loading default model from Hugging Face: {model_id}")
model, model_config = get_pretrained_model(model_id)
# Otherwise, load the model from the local filesystem
else:
print(f"Loading local model from: {model_path}")
# Find the model_config.json file in the same directory as the model
model_dir = os.path.dirname(model_path)
config_path = os.path.join(model_dir, "model_config.json")
if not os.path.exists(config_path):
raise FileNotFoundError(f"Error: Could not find 'model_config.json' in the same directory as the model: {model_dir}")
print(f"Loading model config from: {config_path}")
model_config = OmegaConf.load(config_path)
# Create the model structure from the config
model = create_model_from_config(model_config)
# Load the weights from the checkpoint
if model_path.endswith(".safetensors"):
print("Loading weights from .safetensors file.")
state_dict = load_safetensors(model_path)
elif model_path.endswith(".ckpt"):
print("Loading weights from .ckpt file.")
state_dict = torch.load(model_path, map_location="cpu")["state_dict"]
else:
raise ValueError("Unsupported model file type. Please use .safetensors or .ckpt")
model.load_state_dict(state_dict)
model = model.to(device)
# Convert model to float16 if model_half is True
if model_half:
model = model.to(torch.float16)
print("Model data type:", next(model.parameters()).dtype)
return model, model_config, device
# Define the function to generate audio based on a prompt
def generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, sampler_type, model_half, model, model_config, device):
# Set up text and timing conditioning
conditioning = [{
"prompt": prompt,
"seconds_start": 0,
"seconds_total": generation_time
}]
# Generate stereo audio
output = generate_diffusion_cond(
model,
steps=steps,
cfg_scale=cfg_scale,
conditioning=conditioning,
sample_size=model_config["sample_size"],
sigma_min=sigma_min,
sigma_max=sigma_max,
sampler_type=sampler_type,
device=device,
seed=seed
)
# Rearrange audio batch to a single sequence
output = rearrange(output, "b d n -> d (b n)")
# Peak normalize, clip, and convert to int16 directly if model_half is used
output = output.div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767)
if model_half:
output = output.to(torch.int16).cpu()
else:
output = output.to(torch.float32).to(torch.int16).cpu()
torchaudio.save("temp_output.wav", output, model_config["sample_rate"])
# Convert to MP3 format using pydub
audio = AudioSegment.from_wav("temp_output.wav")
# Create Output folder and dated subfolder if they do not exist
output_folder = "Output"
date_folder = datetime.now().strftime("%Y-%m-%d")
save_path = os.path.join(output_folder, date_folder)
os.makedirs(save_path, exist_ok=True)
# Set a maximum filename length (e.g., 50 characters)
max_length = 50
if len(prompt) > max_length:
prompt = prompt[:max_length] + "_truncated"
# Sanitize the prompt to create a safe filename
filename = re.sub(r'\W+', '_', prompt) + ".mp3"
full_path = os.path.join(save_path, filename)
# Ensure the filename is unique by appending a number if the file already exists
base_filename = filename
counter = 1
while os.path.exists(full_path):
filename = f"{base_filename[:-4]}_{counter}.mp3"
full_path = os.path.join(save_path, filename)
counter += 1
# Export the audio to MP3 format
audio.export(full_path, format="mp3")
return full_path
def audio_generator(prompt, model_path, sampler_type, steps, cfg_scale, sigma_min, sigma_max, generation_time, random_seed, seed, model_half):
"""
Main function called by the Gradio UI to orchestrate audio generation.
"""
try:
print("Generating audio with parameters:")
print("Prompt:", prompt)
print("Sampler Type:", sampler_type)
print("Steps:", steps)
print("CFG Scale:", cfg_scale)
print("Sigma Min:", sigma_min)
print("Sigma Max:", sigma_max)
print("Generation Time:", generation_time)
print("Random Seed:", "Random" if random_seed else "Fixed")
print("Seed:", seed)
print("Model Half Precision:", model_half)
# Set up the model and device
model, model_config, device = setup_model(model_path, model_half)
if random_seed:
seed = torch.randint(0, 1000000, (1,)).item()
filename = generate_audio(prompt, steps, cfg_scale, sigma_min, sigma_max, generation_time, seed, sampler_type, model_half, model, model_config, device)
return gr.Audio(filename), f"Generated: {filename}"
except Exception as e:
return str(e)
# Create Gradio interface
# with gr.Blocks() as demo:
# gr.Markdown("💀🔊 StableAudioWebUI 💀🔊
")
# # Main input components
# prompt_textbox = gr.Textbox(lines=5, label="Prompt")
# sampler_dropdown = gr.Dropdown(
# label="Sampler Type",
# choices=[
# "dpmpp-3m-sde",
# "dpmpp-2m-sde",
# "k-heun",
# "k-lms",
# "k-dpmpp-2s-ancestral",
# "k-dpm-2",
# "k-dpm-fast"
# ],
# value="dpmpp-3m-sde"
# )
# steps_slider = gr.Slider(minimum=0, maximum=200, label="Steps", step=1, value=100)
# generation_time_slider = gr.Slider(minimum=0, maximum=47, label="Generation Time (seconds)", step=1, value=47)
# random_seed_checkbox = gr.Checkbox(label="Random Seed")
# seed_slider = gr.Slider(minimum=-1, maximum=999999, label="Seed", step=1, value=123456)
# # Advanced parameters accordion
# with gr.Accordion("Advanced Parameters", open=False):
# cfg_scale_slider = gr.Slider(minimum=0, maximum=15, label="CFG Scale", step=0.1, value=7)
# sigma_min_slider = gr.Slider(minimum=0, maximum=50, label="Sigma Min", step=0.1, value=0.3)
# sigma_max_slider = gr.Slider(minimum=0, maximum=1000, label="Sigma Max", step=0.1, value=500)
# # Low VRAM checkbox and submit button
# model_half_checkbox = gr.Checkbox(label="Low VRAM (float16)", value=False)
# submit_button = gr.Button("Generate")
# # Define the output components
# audio_output = gr.Audio()
# output_textbox = gr.Textbox(label="Output")
# # Link the button and the function
# random_seed_checkbox.change(fn=toggle_seed_slider, inputs=[random_seed_checkbox], outputs=[seed_slider])
# submit_button.click(audio_generator,
# inputs=[prompt_textbox, sampler_dropdown, steps_slider, cfg_scale_slider,sigma_min_slider, sigma_max_slider, generation_time_slider, random_seed_checkbox, seed_slider, model_half_checkbox],
# outputs=[audio_output, output_textbox])
# # GitHub link at the bottom
# gr.Markdown("Github Repository
")
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("💀🔊 StableAudioWebUI 💀🔊
")
with gr.Row():
with gr.Column(scale=2):
# Main input components
prompt_textbox = gr.Textbox(lines=5, label="Prompt", placeholder="A beautiful orchestral piece with violins, piano, and a choir...")
# NEW: Textbox for local model path
model_path_textbox = gr.Textbox(
label="Local Model Path (Optional)",
placeholder="e.g., /home/user/models/stable-audio-open-1.0.ckpt. Leave blank for default."
)
sampler_dropdown = gr.Dropdown(
label="Sampler Type",
choices=[
"dpmpp-3m-sde",
"dpmpp-2m-sde",
"k-heun",
"k-lms",
"k-dpmpp-2s-ancestral",
"k-dpm-2",
"k-dpm-fast"
],
value="dpmpp-3m-sde"
)
with gr.Row():
steps_slider = gr.Slider(minimum=10, maximum=200, label="Steps", step=1, value=100)
generation_time_slider = gr.Slider(minimum=1, maximum=47, label="Generation Time (seconds)", step=1, value=47)
with gr.Row():
random_seed_checkbox = gr.Checkbox(label="Random Seed", value=True)
seed_slider = gr.Slider(minimum=-1, maximum=999999, label="Seed", step=1, value=12345, interactive=False)
# Advanced parameters accordion
with gr.Accordion("Advanced Parameters", open=False):
cfg_scale_slider = gr.Slider(minimum=0, maximum=25, label="CFG Scale", step=0.1, value=7)
sigma_min_slider = gr.Slider(minimum=0.01, maximum=50, label="Sigma Min", step=0.01, value=0.3)
sigma_max_slider = gr.Slider(minimum=1, maximum=1000, label="Sigma Max", step=1, value=500)
# Low VRAM checkbox and submit button
model_half_checkbox = gr.Checkbox(label="Low VRAM (float16)", value=False)
submit_button = gr.Button("Generate", variant="primary")
with gr.Column(scale=1):
# Define the output components
audio_output = gr.Audio(label="Generated Audio")
output_textbox = gr.Textbox(label="Status", interactive=False)
# Link the button and the function
random_seed_checkbox.change(fn=toggle_seed_slider, inputs=[random_seed_checkbox], outputs=[seed_slider])
# MODIFIED: Added model_path_textbox to the list of inputs
submit_button.click(
fn=audio_generator,
inputs=[
prompt_textbox,
model_path_textbox,
sampler_dropdown,
steps_slider,
cfg_scale_slider,
sigma_min_slider,
sigma_max_slider,
generation_time_slider,
random_seed_checkbox,
seed_slider,
model_half_checkbox
],
outputs=[audio_output, output_textbox]
)
# GitHub link at the bottom
gr.Markdown("Github Repository
")
# Launch the Gradio demo
demo.launch()