| import spaces |
| import gradio as gr |
| 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=7860, share=True) |
