| | import torch |
| | import torchaudio |
| | import os |
| | from einops import rearrange |
| | import gc |
| | import spaces |
| | import gradio as gr |
| | import torch |
| | import torchaudio |
| | import os |
| | from einops import rearrange |
| | from stable_audio_tools import get_pretrained_model |
| | from stable_audio_tools.inference.generation import generate_diffusion_cond |
| | from stable_audio_tools.data.utils import read_video, merge_video_audio, load_and_process_audio |
| | import stat |
| | import platform |
| | import logging |
| | from transformers import logging as transformers_logging |
| |
|
| | transformers_logging.set_verbosity_error() |
| | logging.getLogger("transformers").setLevel(logging.ERROR) |
| |
|
| | model, model_config = get_pretrained_model('HKUSTAudio/AudioX') |
| | sample_rate = model_config["sample_rate"] |
| | sample_size = model_config["sample_size"] |
| |
|
| | TEMP_DIR = "tmp/gradio" |
| | os.makedirs(TEMP_DIR, exist_ok=True) |
| | os.chmod(TEMP_DIR, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) |
| |
|
| | VIDEO_TEMP_DIR = os.path.join(TEMP_DIR, "videos") |
| | os.makedirs(VIDEO_TEMP_DIR, exist_ok=True) |
| | os.chmod(VIDEO_TEMP_DIR, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) |
| |
|
| |
|
| |
|
| | @spaces.GPU(duration=10) |
| | def generate_cond( |
| | prompt, |
| | negative_prompt=None, |
| | video_file=None, |
| | audio_prompt_file=None, |
| | audio_prompt_path=None, |
| | seconds_start=0, |
| | seconds_total=10, |
| | cfg_scale=7.0, |
| | steps=100, |
| | preview_every=0, |
| | seed=-1, |
| | sampler_type="dpmpp-3m-sde", |
| | sigma_min=0.03, |
| | sigma_max=500, |
| | cfg_rescale=0.0, |
| | use_init=False, |
| | init_audio=None, |
| | init_noise_level=0.1, |
| | 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") |
| | |
| | global model |
| | model = model.to(device) |
| |
|
| | target_fps = model_config.get("video_fps", 5) |
| | model_type = model_config.get("model_type", "diffusion_cond") |
| |
|
| | if video_file is not None: |
| | actual_video_path = video_file['name'] if isinstance(video_file, dict) else video_file.name |
| | else: |
| | actual_video_path = None |
| |
|
| | if audio_prompt_file is not None: |
| | audio_path = audio_prompt_file.name |
| | elif audio_prompt_path: |
| | audio_path = audio_prompt_path.strip() |
| | else: |
| | audio_path = None |
| |
|
| | Video_tensors = read_video(actual_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 |
| | |
| | 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 |
| | }] |
| | 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 |
| | }] * 1 |
| | else: |
| | negative_conditioning = None |
| |
|
| | seed = int(seed) |
| | if not use_init: |
| | init_audio = None |
| | input_sample_size = sample_size |
| |
|
| | 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 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=None, |
| | callback=progress_callback if preview_every is not None else None, |
| | scale_phi=cfg_rescale |
| | ) |
| |
|
| | audio = rearrange(audio, "b d n -> d (b n)") |
| |
|
| | samples_10s = 10 * sample_rate |
| | audio = audio[:, :samples_10s] |
| | audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu() |
| |
|
| | output_dir = "demo_result" |
| | os.makedirs(output_dir, exist_ok=True) |
| | output_audio_path = f"{output_dir}/output.wav" |
| | torchaudio.save(output_audio_path, audio, sample_rate) |
| |
|
| | if actual_video_path: |
| | output_video_path = f"{output_dir}/{os.path.basename(actual_video_path)}" |
| | target_width = 1280 |
| | target_height = 720 |
| | merge_video_audio( |
| | actual_video_path, |
| | output_audio_path, |
| | output_video_path, |
| | seconds_start, |
| | seconds_total |
| | ) |
| | else: |
| | output_video_path = None |
| |
|
| | del actual_video_path |
| | torch.cuda.empty_cache() |
| | gc.collect() |
| |
|
| | return output_video_path, output_audio_path |
| |
|
| |
|
| | with gr.Blocks() as interface: |
| | gr.Markdown( |
| | """ |
| | # 🎧AudioX: Diffusion Transformer for Anything-to-Audio Generation |
| | **[Paper](https://arxiv.org/abs/2503.10522) · [Project Page](https://zeyuet.github.io/AudioX/) · [Huggingface](https://huggingface.co/HKUSTAudio/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_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 = gr.Slider( |
| | minimum=0, |
| | maximum=512, |
| | step=1, |
| | value=0, |
| | label="Video Seconds Start" |
| | ) |
| | seconds_total = 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 = gr.Slider( |
| | minimum=1, |
| | maximum=500, |
| | step=1, |
| | value=100, |
| | label="Steps" |
| | ) |
| | preview_every = gr.Slider( |
| | minimum=0, |
| | maximum=100, |
| | step=1, |
| | value=0, |
| | label="Preview Every" |
| | ) |
| | cfg_scale = gr.Slider( |
| | minimum=0.0, |
| | maximum=25.0, |
| | step=0.1, |
| | value=7.0, |
| | label="CFG Scale" |
| | ) |
| | seed = gr.Textbox( |
| | label="Seed (set to -1 for random seed)", |
| | value="-1" |
| | ) |
| | sampler_type = gr.Dropdown( |
| | choices=[ |
| | "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 = gr.Slider( |
| | minimum=0.0, |
| | maximum=2.0, |
| | step=0.01, |
| | value=0.03, |
| | label="Sigma Min" |
| | ) |
| | sigma_max = gr.Slider( |
| | minimum=0.0, |
| | maximum=1000.0, |
| | step=0.1, |
| | value=500, |
| | label="Sigma Max" |
| | ) |
| | cfg_rescale = 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 = gr.Slider( |
| | minimum=0.1, |
| | maximum=100.0, |
| | step=0.01, |
| | value=0.1, |
| | label="Init Noise Level" |
| | ) |
| |
|
| | with gr.Row(): |
| | generate_button = gr.Button("Generate", variant="primary") |
| |
|
| | 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) |
| |
|
| | inputs = [ |
| | prompt, |
| | negative_prompt, |
| | video_file, |
| | audio_prompt_file, |
| | audio_prompt_path, |
| | seconds_start, |
| | seconds_total, |
| | cfg_scale, |
| | steps, |
| | preview_every, |
| | seed, |
| | sampler_type, |
| | sigma_min, |
| | sigma_max, |
| | cfg_rescale, |
| | init_audio_checkbox, |
| | init_audio_input, |
| | init_noise_level |
| | ] |
| |
|
| | generate_button.click( |
| | fn=generate_cond, |
| | inputs=inputs, |
| | outputs=[video_output, audio_output] |
| | ) |
| |
|
| | 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") |
| |
|
| | ex1.click(lambda: ["Typing on a keyboard", 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, 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, 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, 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, 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, 0, 10, 7.0, 100, 0, "807934770", "dpmpp-3m-sde", 0.03, 500, 0.0, False, None, 0.1], inputs=[], outputs=inputs) |
| |
|
| | interface.queue(5).launch(server_name="0.0.0.0", server_port=2159, share=True) |
