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()