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