stable_audio_tools/interface/gradio.py

import gc
import platform
import os
import subprocess as sp
import gradio as gr
import json 
import torch
import torchaudio

from aeiou.viz import audio_spectrogram_image
from einops import rearrange
from safetensors.torch import load_file
from torch.nn import functional as F
from torchaudio import transforms as T

from ..inference.generation import generate_diffusion_cond, generate_diffusion_uncond
from ..models.factory import create_model_from_config
from ..models.pretrained import get_pretrained_model
from ..models.utils import load_ckpt_state_dict
from ..inference.utils import prepare_audio
from ..training.utils import copy_state_dict
from ..data.utils import read_video, merge_video_audio


import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import warnings
warnings.filterwarnings("ignore", category=UserWarning)


device = torch.device("cpu")

os.environ['TMPDIR'] = './tmp'

current_model_name = None
current_model = None
current_sample_rate = None
current_sample_size = None



def load_model(model_name, model_config=None, model_ckpt_path=None, pretrained_name=None, pretransform_ckpt_path=None, device="cuda", model_half=False):
    global model_configurations
    
    if pretrained_name is not None:
        print(f"Loading pretrained model {pretrained_name}")
        model, model_config = get_pretrained_model(pretrained_name)
    elif model_config is not None and model_ckpt_path is not None:
        print(f"Creating model from config")
        model = create_model_from_config(model_config)
        print(f"Loading model checkpoint from {model_ckpt_path}")
        copy_state_dict(model, load_ckpt_state_dict(model_ckpt_path))
    sample_rate = model_config["sample_rate"]
    sample_size = model_config["sample_size"]
    if pretransform_ckpt_path is not None:
        print(f"Loading pretransform checkpoint from {pretransform_ckpt_path}")
        model.pretransform.load_state_dict(load_ckpt_state_dict(pretransform_ckpt_path), strict=False)
        print(f"Done loading pretransform")
    model.to(device).eval().requires_grad_(False)
    if model_half:
        model.to(torch.float16)
    print(f"Done loading model")
    return model, model_config, sample_rate, sample_size

def load_and_process_audio(audio_path, sample_rate, seconds_start, seconds_total):
    if audio_path is None:
        return torch.zeros((2, int(sample_rate * seconds_total)))
    audio_tensor, sr = torchaudio.load(audio_path)
    start_index = int(sample_rate * seconds_start)
    target_length = int(sample_rate * seconds_total)
    end_index = start_index + target_length
    audio_tensor = audio_tensor[:, start_index:end_index]
    if audio_tensor.shape[1] < target_length:
        pad_length = target_length - audio_tensor.shape[1]
        audio_tensor = F.pad(audio_tensor, (pad_length, 0))
    return audio_tensor

def generate_cond(
        prompt,
        negative_prompt=None,
        video_file=None,
        video_path=None,
        audio_prompt_file=None,
        audio_prompt_path=None,
        seconds_start=0,
        seconds_total=10,
        cfg_scale=6.0,
        steps=250,
        preview_every=None,
        seed=-1,
        sampler_type="dpmpp-3m-sde",
        sigma_min=0.03,
        sigma_max=1000,
        cfg_rescale=0.0,
        use_init=False,
        init_audio=None,
        init_noise_level=1.0,
        mask_cropfrom=None,
        mask_pastefrom=None,
        mask_pasteto=None,
        mask_maskstart=None,
        mask_maskend=None,
        mask_softnessL=None,
        mask_softnessR=None,
        mask_marination=None,
        batch_size=1
    ):
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()
    print(f"Prompt: {prompt}")
    preview_images = []
    if preview_every == 0:
        preview_every = None

    try:
        has_mps = platform.system() == "Darwin" and torch.backends.mps.is_available()
    except Exception:
        has_mps = False
    if has_mps:
        device = torch.device("mps")
    elif torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    model_name = 'default'
    cfg = model_configurations[model_name]
    model_config_path = cfg.get("model_config")
    ckpt_path = cfg.get("ckpt_path")
    pretrained_name = cfg.get("pretrained_name")
    pretransform_ckpt_path = cfg.get("pretransform_ckpt_path")
    model_type = cfg.get("model_type", "diffusion_cond")
    if model_config_path:
        with open(model_config_path) as f:
            model_config = json.load(f)
    else:
        model_config = None
    target_fps = model_config.get("video_fps", 5)
    global current_model_name, current_model, current_sample_rate, current_sample_size
    if current_model is None or model_name != current_model_name:
        current_model, model_config, sample_rate, sample_size = load_model(
            model_name=model_name,
            model_config=model_config,
            model_ckpt_path=ckpt_path,
            pretrained_name=pretrained_name,
            pretransform_ckpt_path=pretransform_ckpt_path,
            device=device,
            model_half=False
        )
        current_model_name = model_name
        model = current_model
        current_sample_rate = sample_rate
        current_sample_size = sample_size
    else:
        model = current_model
        sample_rate = current_sample_rate
        sample_size = current_sample_size
    if video_file is not None:
        video_path = video_file.name
    elif video_path:
        video_path = video_path.strip()
    else:
        video_path = None
        
    if audio_prompt_file is not None:
        print(f'audio_prompt_file: {audio_prompt_file}')
        audio_path = audio_prompt_file.name
    elif audio_prompt_path:
        audio_path = audio_prompt_path.strip()
    else:
        audio_path = None
    
    Video_tensors = read_video(video_path, seek_time=seconds_start, duration=seconds_total, target_fps=target_fps)        
    audio_tensor = load_and_process_audio(audio_path, sample_rate, seconds_start, seconds_total)
    
    audio_tensor = audio_tensor.to(device)
    seconds_input = sample_size / sample_rate
    print(f'video_path: {video_path}')
    
    if not prompt:
        prompt = ""
        
    conditioning = [{
        "video_prompt": [Video_tensors.unsqueeze(0)],        
        "text_prompt": prompt,
        "audio_prompt": audio_tensor.unsqueeze(0),
        "seconds_start": seconds_start,
        "seconds_total": seconds_input
    }] * batch_size
    if negative_prompt:
        negative_conditioning = [{
            "video_prompt": [Video_tensors.unsqueeze(0)],        
            "text_prompt": negative_prompt,
            "audio_prompt": audio_tensor.unsqueeze(0),
            "seconds_start": seconds_start,
            "seconds_total": seconds_total
        }] * batch_size
    else:
        negative_conditioning = None
    try:
        device = next(model.parameters()).device 
    except Exception as e:
        device = next(current_model.parameters()).device
    seed = int(seed)
    if not use_init:
        init_audio = None
    input_sample_size = sample_size
    if init_audio is not None:
        in_sr, init_audio = init_audio
        init_audio = torch.from_numpy(init_audio).float().div(32767)
        if init_audio.dim() == 1:
            init_audio = init_audio.unsqueeze(0)
        elif init_audio.dim() == 2:
            init_audio = init_audio.transpose(0, 1)
        if in_sr != sample_rate:
            resample_tf = T.Resample(in_sr, sample_rate).to(init_audio.device)
            init_audio = resample_tf(init_audio)
        audio_length = init_audio.shape[-1]
        if audio_length > sample_size:
            input_sample_size = audio_length + (model.min_input_length - (audio_length % model.min_input_length)) % model.min_input_length
        init_audio = (sample_rate, init_audio)
    def progress_callback(callback_info):
        nonlocal preview_images
        denoised = callback_info["denoised"]
        current_step = callback_info["i"]
        sigma = callback_info["sigma"]
        if (current_step - 1) % preview_every == 0:
            if model.pretransform is not None:
                denoised = model.pretransform.decode(denoised)
            denoised = rearrange(denoised, "b d n -> d (b n)")
            denoised = denoised.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
            audio_spectrogram = audio_spectrogram_image(denoised, sample_rate=sample_rate)
            preview_images.append((audio_spectrogram, f"Step {current_step} sigma={sigma:.3f})"))
    if mask_cropfrom is not None: 
        mask_args = {
            "cropfrom": mask_cropfrom,
            "pastefrom": mask_pastefrom,
            "pasteto": mask_pasteto,
            "maskstart": mask_maskstart,
            "maskend": mask_maskend,
            "softnessL": mask_softnessL,
            "softnessR": mask_softnessR,
            "marination": mask_marination,
        }
    else:
        mask_args = None 
    if model_type == "diffusion_cond":
        audio = generate_diffusion_cond(
            model, 
            conditioning=conditioning,
            negative_conditioning=negative_conditioning,
            steps=steps,
            cfg_scale=cfg_scale,
            batch_size=batch_size,
            sample_size=input_sample_size,
            sample_rate=sample_rate,
            seed=seed,
            device=device,
            sampler_type=sampler_type,
            sigma_min=sigma_min,
            sigma_max=sigma_max,
            init_audio=init_audio,
            init_noise_level=init_noise_level,
            mask_args=mask_args,
            callback=progress_callback if preview_every is not None else None,
            scale_phi=cfg_rescale
        )
    elif model_type == "diffusion_uncond":
        audio = generate_diffusion_uncond(
            model, 
            steps=steps,
            batch_size=batch_size,
            sample_size=input_sample_size,
            seed=seed,
            device=device,
            sampler_type=sampler_type,
            sigma_min=sigma_min,
            sigma_max=sigma_max,
            init_audio=init_audio,
            init_noise_level=init_noise_level,
            callback=progress_callback if preview_every is not None else None
        )
    else:
        raise ValueError(f"Unsupported model type: {model_type}")
    audio = rearrange(audio, "b d n -> d (b n)")
    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
    file_name = os.path.basename(video_path) if video_path else "output"
    output_dir = f"demo_result"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_video_path = f"{output_dir}/{file_name}"
    torchaudio.save(f"{output_dir}/output.wav", audio, sample_rate)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    if video_path:
        merge_video_audio(video_path, f"{output_dir}/output.wav", output_video_path, seconds_start, seconds_total)
    audio_spectrogram = audio_spectrogram_image(audio, sample_rate=sample_rate)
    del video_path
    torch.cuda.empty_cache()
    gc.collect()
    return (output_video_path, f"{output_dir}/output.wav")

def toggle_custom_model(selected_model):
    return gr.Row.update(visible=(selected_model == "Custom Model"))

def create_sampling_ui(model_config_map, inpainting=False):
    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # 🎧AudioX: Diffusion Transformer for Anything-to-Audio Generation  
            **[Project Page](https://zeyuet.github.io/AudioX/) · [Huggingface](https://huggingface.co/Zeyue7/AudioX) · [GitHub](https://github.com/ZeyueT/AudioX)**
            """
        )

        with gr.Tab("Generation"):

            with gr.Row():
                with gr.Column():
                    prompt = gr.Textbox(show_label=False, placeholder="Enter your prompt")
                    negative_prompt = gr.Textbox(show_label=False, placeholder="Negative prompt", visible=False)
                    video_path = gr.Textbox(label="Video Path", placeholder="Enter video file path")
                    video_file = gr.File(label="Upload Video File")
                    audio_prompt_file = gr.File(label="Upload Audio Prompt File", visible=False)
                    audio_prompt_path = gr.Textbox(label="Audio Prompt Path", placeholder="Enter audio file path", visible=False)
            with gr.Row():
                with gr.Column(scale=6):
                    with gr.Accordion("Video Params", open=False):                
                        seconds_start_slider = gr.Slider(minimum=0, maximum=512, step=1, value=0, label="Video Seconds Start")
                        seconds_total_slider = gr.Slider(minimum=0, maximum=10, step=1, value=10, label="Seconds Total", interactive=False)
            with gr.Row():
                with gr.Column(scale=4):
                    with gr.Accordion("Sampler Params", open=False):
                        steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=100, label="Steps")
                        preview_every_slider = gr.Slider(minimum=0, maximum=100, step=1, value=0, label="Preview Every")
                        cfg_scale_slider = gr.Slider(minimum=0.0, maximum=25.0, step=0.1, value=7.0, label="CFG Scale")
                        seed_textbox = gr.Textbox(label="Seed (set to -1 for random seed)", value="-1")
                        sampler_type_dropdown = 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"
                        )
                        sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma Min")
                        sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=500, label="Sigma Max")
                        cfg_rescale_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=0.0, label="CFG Rescale Amount")
            with gr.Row():
                with gr.Column(scale=4):
                    with gr.Accordion("Init Audio", open=False, visible=False):
                        init_audio_checkbox = gr.Checkbox(label="Use Init Audio")
                        init_audio_input = gr.Audio(label="Init Audio")
                        init_noise_level_slider = gr.Slider(minimum=0.1, maximum=100.0, step=0.01, value=0.1, label="Init Noise Level")
            gr.Markdown("## Examples")
            with gr.Accordion("Click to show examples", open=False):
                with gr.Row():
                    gr.Markdown("**📝 Task: Text-to-Audio**")                
                    with gr.Column(scale=1.2):
                        gr.Markdown("Prompt: *Typing on a keyboard*")
                        ex1 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Markdown("Prompt: *Ocean waves crashing*")
                        ex2 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Markdown("Prompt: *Footsteps in snow*")
                        ex3 = gr.Button("Load Example")
                with gr.Row():
                    gr.Markdown("**🎶 Task: Text-to-Music**")                
                    with gr.Column(scale=1.2):
                        gr.Markdown("Prompt: *An orchestral music piece for a fantasy world.*")
                        ex4 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Markdown("Prompt: *Produce upbeat electronic music for a dance party*")
                        ex5 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Markdown("Prompt: *A dreamy lo-fi beat with vinyl crackle*")
                        ex6 = gr.Button("Load Example")
                with gr.Row():
                    gr.Markdown("**🎬 Task: Video-to-Audio**\nPrompt: *Generate general audio for the video*")
                    with gr.Column(scale=1.2):
                        gr.Video("example/V2A_sample-1.mp4")
                        ex7 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Video("example/V2A_sample-2.mp4")
                        ex8 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Video("example/V2A_sample-3.mp4")
                        ex9 = gr.Button("Load Example")
                with gr.Row():
                    gr.Markdown("**🎵 Task: Video-to-Music**\nPrompt: *Generate music for the video*")                
                    with gr.Column(scale=1.2):
                        gr.Video("example/V2M_sample-1.mp4")
                        ex10 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Video("example/V2M_sample-2.mp4")
                        ex11 = gr.Button("Load Example")
                    with gr.Column(scale=1.2):
                        gr.Video("example/V2M_sample-3.mp4")
                        ex12 = gr.Button("Load Example")
            with gr.Row():
                generate_button = gr.Button("Generate", variant='primary', scale=1)
            with gr.Row():
                with gr.Column(scale=6):
                    video_output = gr.Video(label="Output Video", interactive=False)
                    audio_output = gr.Audio(label="Output Audio", interactive=False)
                    send_to_init_button = gr.Button("Send to Init Audio", scale=1, visible=False)
            send_to_init_button.click(
                fn=lambda audio: audio,
                inputs=[audio_output],
                outputs=[init_audio_input]
            )
            inputs = [
                prompt, 
                negative_prompt,
                video_file,
                video_path,
                audio_prompt_file,
                audio_prompt_path,
                seconds_start_slider, 
                seconds_total_slider, 
                cfg_scale_slider, 
                steps_slider, 
                preview_every_slider, 
                seed_textbox, 
                sampler_type_dropdown, 
                sigma_min_slider, 
                sigma_max_slider,
                cfg_rescale_slider,
                init_audio_checkbox,
                init_audio_input,
                init_noise_level_slider
            ]
            generate_button.click(
                fn=generate_cond, 
                inputs=inputs,
                outputs=[
                    video_output,
                    audio_output
                ], 
                api_name="generate"
            ) 
            ex1.click(lambda: ["Typing on a keyboard", None, None, None, None, None, 0, 10, 7.0, 100, 0, "1225575558", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex2.click(lambda: ["Ocean waves crashing", None, None, None, None, None, 0, 10, 7.0, 100, 0, "3615819170", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex3.click(lambda: ["Footsteps in snow", None, None, None, None, None, 0, 10, 7.0, 100, 0, "1703896811", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex4.click(lambda: ["An orchestral music piece for a fantasy world.", None, None, None, None, None, 0, 10, 7.0, 100, 0, "1561898939", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex5.click(lambda: ["Produce upbeat electronic music for a dance party", None, None, None, None, None, 0, 10, 7.0, 100, 0, "406022999", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex6.click(lambda: ["A dreamy lo-fi beat with vinyl crackle", None, None, None, None, None, 0, 10, 7.0, 100, 0, "807934770", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex7.click(lambda: ["Generate general audio for the video", None, None, "example/V2A_sample-1.mp4", None, None, 0, 10, 7.0, 100, 0, "3737819478", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex8.click(lambda: ["Generate general audio for the video", None, None, "example/V2A_sample-2.mp4", None, None, 0, 10, 7.0, 100, 0, "1900718499", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex9.click(lambda: ["Generate general audio for the video", None, None, "example/V2A_sample-3.mp4", None, None, 0, 10, 7.0, 100, 0, "2289822202", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex10.click(lambda: ["Generate music for the video", None, None, "example/V2M_sample-1.mp4", None, None, 0, 10, 7.0, 100, 0, "3498087420", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex11.click(lambda: ["Generate music for the video", None, None, "example/V2M_sample-2.mp4", None, None, 0, 10, 7.0, 100, 0, "3753837734", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
            ex12.click(lambda: ["Generate music for the video", None, None, "example/V2M_sample-3.mp4", None, None, 0, 10, 7.0, 100, 0, "3510832996", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs)
        return demo

def create_txt2audio_ui(model_config_map):
    with gr.Blocks(css=".gradio-container { max-width: 1120px; margin: auto; }") as ui:
        with gr.Tab("Generation"):
            create_sampling_ui(model_config_map)
    return ui

def toggle_custom_model(selected_model):
    return gr.Row.update(visible=(selected_model == "Custom Model"))

def create_ui(model_config_path=None, ckpt_path=None, pretrained_name=None, pretransform_ckpt_path=None, model_half=False):
    global model_configurations
    global device

    try:
        has_mps = platform.system() == "Darwin" and torch.backends.mps.is_available()
    except Exception:
        has_mps = False

    if has_mps:
        device = torch.device("mps")
    elif torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    print("Using device:", device)

    model_configurations = {
        "default": {
            "model_config": "./model/config.json",
            "ckpt_path": "./model/model.ckpt"
        }
    }
    ui = create_txt2audio_ui(model_configurations)
    return ui

if __name__ == "__main__":
    ui = create_ui(
        model_config_path='./model/config.json',
        share=True
    )
    ui.launch()