294 lines
12 KiB
Python
294 lines
12 KiB
Python
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("<h1 style='text-align: center; font-size: 300%;'>💀🔊 StableAudioWebUI 💀🔊</h1>")
|
|
|
|
# # 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("<p style='text-align: center;'><a href='https://github.com/Saganaki22/StableAudioWebUI'>Github Repository</a></p>")
|
|
|
|
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
|
gr.Markdown("<h1 style='text-align: center; font-size: 300%;'>💀🔊 StableAudioWebUI 💀🔊</h1>")
|
|
|
|
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("<p style='text-align: center;'><a href='https://github.com/Saganaki22/StableAudioWebUI' target='_blank'>Github Repository</a></p>")
|
|
|
|
# Launch the Gradio demo
|
|
demo.launch()
|