app.py

import torch
import random
import numpy as np
import gradio as gr
import librosa
# import spaces
from accelerate import Accelerator
from transformers import T5Tokenizer, T5EncoderModel
from diffusers import DDIMScheduler
from src.models.conditioners import MaskDiT
from src.models.controlnet import DiTControlNet
from src.models.conditions import Conditioner
from src.modules.autoencoder_wrapper import Autoencoder
from src.inference_controlnet import inference
from src.utils import load_yaml_with_includes


# Load model and configs
def load_models(config_name, ckpt_path, controlnet_path, vae_path, device):
    params = load_yaml_with_includes(config_name)

    # Load codec model
    autoencoder = Autoencoder(ckpt_path=vae_path,
                              model_type=params['autoencoder']['name'],
                              quantization_first=params['autoencoder']['q_first']).to(device)
    autoencoder.eval()

    # Load text encoder
    tokenizer = T5Tokenizer.from_pretrained(params['text_encoder']['model'])
    text_encoder = T5EncoderModel.from_pretrained(params['text_encoder']['model']).to(device)
    text_encoder.eval()

    # Load main U-Net model
    unet = MaskDiT(**params['model']).to(device)
    unet.load_state_dict(torch.load(ckpt_path, map_location='cpu')['model'])
    unet.eval()

    controlnet_config = params['model'].copy()
    controlnet_config.update(params['controlnet'])
    controlnet = DiTControlNet(**controlnet_config).to(device)
    controlnet.eval()
    controlnet.load_state_dict(torch.load(controlnet_path, map_location='cpu')['model'])
    conditioner = Conditioner(**params['conditioner']).to(device)

    accelerator = Accelerator(mixed_precision="fp16")
    unet, controlnet = accelerator.prepare(unet, controlnet)

    # Load noise scheduler
    noise_scheduler = DDIMScheduler(**params['diff'])

    latents = torch.randn((1, 128, 128), device=device)
    noise = torch.randn_like(latents)
    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (1,), device=device)
    _ = noise_scheduler.add_noise(latents, noise, timesteps)

    return autoencoder, unet, controlnet, conditioner, tokenizer, text_encoder, noise_scheduler, params


MAX_SEED = np.iinfo(np.int32).max

# Model and config paths
config_name = 'ckpts/controlnet/energy_l.yml'
ckpt_path = 'ckpts/s3/ezaudio_s3_l.pt'
controlnet_path = 'ckpts/controlnet/s3_l_energy.pt'
vae_path = 'ckpts/vae/1m.pt'
# save_path = 'output/'
# os.makedirs(save_path, exist_ok=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

(autoencoder, unet, controlnet, conditioner, 
 tokenizer, text_encoder, noise_scheduler, params) = load_models(config_name, ckpt_path, controlnet_path, vae_path, device)


# @spaces.GPU
def generate_audio(text,
                   audio_path, surpass_noise,
                   guidance_scale, guidance_rescale, 
                   ddim_steps, eta,
                   conditioning_scale,
                   random_seed, randomize_seed):
    sr = params['autoencoder']['sr']

    gt, _ = librosa.load(audio_path, sr=sr)
    gt = gt / (np.max(np.abs(gt)) + 1e-9)  # Normalize audio

    if surpass_noise > 0:
        mask = np.abs(gt) <= surpass_noise
        gt[mask] = 0

    original_length = len(gt)
    # Ensure the audio is of the correct length by padding or trimming
    duration_seconds = min(len(gt) / sr, 10)
    quantized_duration = np.ceil(duration_seconds * 2) / 2  # This rounds to the nearest 0.5 seconds
    num_samples = int(quantized_duration * sr)
    audio_frames = round(num_samples / sr * params['autoencoder']['latent_sr'])

    if len(gt) < num_samples:
        padding = num_samples - len(gt)
        gt = np.pad(gt, (0, padding), 'constant')
    else:
        gt = gt[:num_samples]

    gt_audio = torch.tensor(gt).unsqueeze(0).unsqueeze(1).to(device)
    gt = autoencoder(audio=gt_audio)
    condition = conditioner(gt_audio.squeeze(1), gt.shape)

    # Handle random seed
    if randomize_seed:
        random_seed = random.randint(0, MAX_SEED)

    # Perform inference
    pred = inference(autoencoder, unet, controlnet,
                     None, None, condition,
                     tokenizer, text_encoder, 
                     params, noise_scheduler,
                     text, neg_text=None,
                     audio_frames=audio_frames, 
                     guidance_scale=guidance_scale, guidance_rescale=guidance_rescale, 
                     ddim_steps=ddim_steps, eta=eta, random_seed=random_seed, 
                     conditioning_scale=conditioning_scale, device=device)

    pred = pred.cpu().numpy().squeeze(0).squeeze(0)[:original_length]

    return sr, pred

# CSS styling (optional)
css = """
#col-container {
    margin: 0 auto;
    max-width: 1280px;
}
"""

examples_energy = [
    ["Dog barking in the background", "reference.mp3"],
    ["Duck quacking", "reference2.mp3"],
    ["Truck honking on the street", "reference3.mp3"]
]


# Gradio Blocks layout
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
        # EzAudio-ControlNet: Interactive and Creative Control for Text-to-Audio Generation
        EzAudio-ControlNet enables control over the timing of sound effects within audio generation.
        
        Learn more about 🟣**EzAudio** on the [EzAudio Homepage](https://haidog-yaqub.github.io/EzAudio-Page/).
        
        Explore **Vanilla Text-to-Audio**, **Editing**, and **Inpainting** features on the [🤗EzAudio Space](https://huggingface.co/spaces/OpenSound/EzAudio).
    """)
    with gr.Row():
        # Input for the text prompt (used for generating new audio)
        text_input = gr.Textbox(
            label="Text Prompt",
            show_label=True,
            max_lines=2,
            placeholder="Describe the sound you want to generate",
            value="Truck honking on the street",
            scale=4
        )
        # Button to generate the audio
        generate_button = gr.Button("Generate")
        
    # Audio input to use as base
    audio_file_input = gr.Audio(label="Upload Reference Audio (less than 10s)", value='reference3.mp3', type="filepath")
    
    # Output Component for the generated audio
    generated_audio_output = gr.Audio(label="Generated Audio", type="numpy")
    
    with gr.Accordion("Advanced Settings", open=False):
        # Length of the generated audio
        surpass_noise = gr.Slider(minimum=0, maximum=0.1, step=0.01, value=0.0, label="Noise Threshold (Amplitude)")
        guidance_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.5, value=5.0, label="Guidance Scale")
        guidance_rescale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Guidance Rescale")
        ddim_steps = gr.Slider(minimum=25, maximum=200, step=5, value=50, label="DDIM Steps")
        eta = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="Eta")
        conditioning_scale = gr.Slider(minimum=0.0, maximum=2.0, step=0.25, value=1.0, label="Conditioning Scale")
        random_seed = gr.Slider(minimum=0, maximum=10000, step=1, value=0, label="Random Seed")
        randomize_seed = gr.Checkbox(label="Randomize Seed (Disable Seed)", value=True)

    gr.Examples(
    examples=examples_energy,
    inputs=[text_input, audio_file_input]
    )
    
    # Link the inputs to the function
    generate_button.click(
        fn=generate_audio,
        inputs=[
            text_input, audio_file_input, surpass_noise, guidance_scale, guidance_rescale,
            ddim_steps, eta, conditioning_scale, random_seed, randomize_seed
        ],
        outputs=[generated_audio_output]
    )

    text_input.submit(
        fn=generate_audio,
        inputs=[
            text_input, audio_file_input, surpass_noise, guidance_scale, guidance_rescale,
            ddim_steps, eta, conditioning_scale, random_seed, randomize_seed
        ],
        outputs=[generated_audio_output]
    )

# Launch the Gradio demo
demo.launch()