import torch import torchaudio from einops import rearrange import gradio as gr import spaces import os import uuid # Importing the model-related functions from stable_audio_tools import get_pretrained_model from stable_audio_tools.inference.generation import generate_diffusion_cond PAGE_SIZE = 10 FILE_DIR_PATH = "/data" theme = gr.themes.Base( font=[gr.themes.GoogleFont('Libre Franklin'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif'], ) # Load the model outside of the GPU-decorated function def load_model(): model, model_config = get_pretrained_model("stabilityai/stable-audio-open-1.0") print("Loading model...Done") return model, model_config # Function to set up, generate, and process the audio @spaces.GPU(duration=120) # Allocate GPU only when this function is called def generate_audio(prompt, sampler_type_dropdown, seconds_total=30, steps=100, cfg_scale=7,sigma_min_slider=0.3,sigma_max_slider=500, progress=gr.Progress(track_tqdm=True)): print(f"Prompt received: {prompt}") print(f"Settings: Duration={seconds_total}s, Steps={steps}, CFG Scale={cfg_scale}") device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # Fetch the Hugging Face token from the environment variable hf_token = os.getenv('HF_TOKEN') print(f"Hugging Face token: {hf_token}") # Use pre-loaded model and configuration model, model_config = load_model() sample_rate = model_config["sample_rate"] sample_size = model_config["sample_size"] print(f"Sample rate: {sample_rate}, Sample size: {sample_size}") model = model.to(device) print("Model moved to device.") # Set up text and timing conditioning conditioning = [{ "prompt": prompt, "seconds_start": 0, "seconds_total": seconds_total }] print(f"Conditioning: {conditioning}") # Generate stereo audio print("Generating audio...") output = generate_diffusion_cond( model, steps=steps, cfg_scale=cfg_scale, conditioning=conditioning, sample_size=sample_size, sigma_min=sigma_min_slider, sigma_max=sigma_max_slider, sampler_type=sampler_type_dropdown,#"dpmpp-3m-sde", device=device ) print("Audio generated.") # Rearrange audio batch to a single sequence output = rearrange(output, "b d n -> d (b n)") print("Audio rearranged.") # Peak normalize, clip, convert to int16 output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu() max_length = sample_rate * seconds_total if output.shape[1] > max_length: output = output[:, :max_length] print(f"Audio trimmed to {seconds_total} seconds.") # Generate a unique filename for the output random_uuid = uuid.uuid4().hex unique_filename = f"/data/output_{random_uuid}.wav" unique_textfile = f"/data/output_{random_uuid}.txt" print(f"Saving audio to file: {unique_filename}") # Save to file torchaudio.save(unique_filename, output, sample_rate) print(f"Audio saved: {unique_filename}") with open(unique_textfile, "w") as file: file.write(prompt) # Return the path to the generated audio file return unique_filename def list_all_outputs(generation_history): directory_path = FILE_DIR_PATH files_in_directory = os.listdir(directory_path) wav_files = [os.path.join(directory_path, file) for file in files_in_directory if file.endswith('.wav')] wav_files.sort(key=lambda x: os.path.getmtime(os.path.join(directory_path, x)), reverse=True) history_list = generation_history.split(',') if generation_history else [] updated_files = [file for file in wav_files if file not in history_list] updated_history = updated_files + history_list return ','.join(updated_history), gr.update(visible=True) def increase_list_size(list_size): return list_size+PAGE_SIZE css = ''' #live_gen:before { content: ''; animation: svelte-z7cif2-pulseStart 1s cubic-bezier(.4,0,.6,1), svelte-z7cif2-pulse 2s cubic-bezier(.4,0,.6,1) 1s infinite; border: 2px solid var(--color-accent); background: transparent; z-index: var(--layer-1); pointer-events: none; position: absolute; height: 100%; width: 100%; border-radius: 7px; } #live_gen_items{ max-height: 570px; overflow-y: scroll; } ''' examples = [ [ "A serene soundscape of a quiet beach at sunset.", # Text prompt "dpmpp-2m-sde", # Sampler type 45, # Duration in Seconds 100, # Number of Diffusion Steps 10, # CFG Scale 0.5, # Sigma min 800 # Sigma max ], [ "clapping crowd", # Text prompt "dpmpp-3m-sde", # Sampler type 30, # Duration in Seconds 100, # Number of Diffusion Steps 7, # CFG Scale 0.5, # Sigma min 500 # Sigma max ], [ "A forest ambiance with birds chirping and wind rustling through the leaves.", # Text prompt "k-dpm-fast", # Sampler type 60, # Duration in Seconds 140, # Number of Diffusion Steps 7.5, # CFG Scale 0.3, # Sigma min 700 # Sigma max ], [ "A gentle rainfall with distant thunder.", # Text prompt "dpmpp-3m-sde", # Sampler type 35, # Duration in Seconds 110, # Number of Diffusion Steps 8, # CFG Scale 0.1, # Sigma min 500 # Sigma max ], [ "A jazz cafe environment with soft music and ambient chatter.", # Text prompt "k-lms", # Sampler type 25, # Duration in Seconds 90, # Number of Diffusion Steps 6, # CFG Scale 0.4, # Sigma min 650 # Sigma max ], ["Rock beat played in a treated studio, session drumming on an acoustic kit.", "dpmpp-2m-sde", # Sampler type 30, # Duration in Seconds 100, # Number of Diffusion Steps 7, # CFG Scale 0.3, # Sigma min 500 # Sigma max ] ] with gr.Blocks(theme=theme, css=css) as demo: gr.Markdown("# Stable Audio Multiplayer Live") gr.Markdown("Generate audio with text, share and learn from others how to best prompt this new model") generation_history = gr.Textbox(visible=False) list_size = gr.Number(value=PAGE_SIZE, visible=False) with gr.Row(): with gr.Column(): prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here") btn_run = gr.Button("Generate") with gr.Accordion("Parameters", open=True): with gr.Row(): duration = gr.Slider(0, 47, value=20, step=1, label="Duration in Seconds") with gr.Accordion("Advanced parameters", open=False): steps = gr.Slider(10, 150, value=80, step=10, label="Number of Diffusion Steps") sampler_type = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde") with gr.Row(): cfg_scale = gr.Slider(1, 15, value=7, step=0.1, label="CFG Scale") sigma_min = gr.Slider(0.0, 5.0, step=0.01, value=0.3, label="Sigma min") sigma_max = gr.Slider(0.0, 1000.0, step=0.1, value=500, label="Sigma max") with gr.Column() as output_list: output = gr.Audio(type="filepath", label="Generated Audio") with gr.Column(elem_id="live_gen") as community_list: gr.Markdown("# Community generations") with gr.Column(elem_id="live_gen_items"): @gr.render(inputs=[generation_history, list_size]) def show_output_list(generation_history, list_size): history_list = generation_history.split(',') if generation_history else [] history_list_latest = history_list[:list_size] for generation in history_list_latest: generation_prompt_file = generation.replace('.wav', '.txt') with open(generation_prompt_file, 'r') as file: generation_prompt = file.read() with gr.Group(): gr.Markdown(value=f"### {generation_prompt}") gr.Audio(value=generation) load_more = gr.Button("Load more") load_more.click(fn=increase_list_size, inputs=list_size, outputs=list_size) gr.Examples( fn=generate_audio, examples=examples, inputs=[prompt, sampler_type, duration, steps, cfg_scale, sigma_min, sigma_max], outputs=output, cache_examples="lazy" ) gr.on( triggers=[btn_run.click, prompt.submit], fn=generate_audio, inputs=[prompt, sampler_type, duration, steps, cfg_scale, sigma_min, sigma_max], outputs=output ) demo.load(fn=list_all_outputs, inputs=generation_history, outputs=[generation_history, community_list], every=2) model, model_config = load_model() demo.launch()