Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import json | |
| import numpy as np | |
| import torch | |
| import soundfile as sf | |
| import gradio as gr | |
| from diffusers import DDPMScheduler | |
| from pico_model import PicoDiffusion | |
| from audioldm.variational_autoencoder.autoencoder import AutoencoderKL | |
| from llm_preprocess import get_event, preprocess_gemini, preprocess_gpt | |
| import spaces | |
| class dotdict(dict): | |
| """dot.notation access to dictionary attributes""" | |
| __getattr__ = dict.get | |
| __setattr__ = dict.__setitem__ | |
| __delattr__ = dict.__delitem__ | |
| class InferRunner: | |
| def __init__(self, device): | |
| vae_config = json.load(open("ckpts/ldm/vae_config.json")) | |
| self.vae = AutoencoderKL(**vae_config) | |
| vae_weights = torch.load("ckpts/ldm/pytorch_model_vae.bin", map_location="cpu") | |
| self.vae.load_state_dict(vae_weights) | |
| self.vae = self.vae.to(device) | |
| train_args = dotdict(json.loads(open("ckpts/pico_model/summary.jsonl").readlines()[0])) | |
| self.pico_model = PicoDiffusion( | |
| scheduler_name=train_args.scheduler_name, | |
| unet_model_config_path=train_args.unet_model_config, | |
| snr_gamma=train_args.snr_gamma, | |
| freeze_text_encoder_ckpt="ckpts/laion_clap/630k-audioset-best.pt", | |
| diffusion_pt="ckpts/pico_model/diffusion.pt", | |
| ).eval().to(device) | |
| self.scheduler = DDPMScheduler.from_pretrained(train_args.scheduler_name, subfolder="scheduler") | |
| #device = "cuda" if torch.cuda.is_available() else "cpu" | |
| device = "cuda" | |
| runner = InferRunner(device) | |
| event_list = get_event() | |
| def infer(caption, num_steps=200, guidance_scale=3.0, audio_len=16000*10): | |
| with torch.no_grad(): | |
| latents = runner.pico_model.demo_inference(caption, runner.scheduler, num_steps=num_steps, guidance_scale=guidance_scale, num_samples_per_prompt=1, disable_progress=True) | |
| mel = runner.vae.decode_first_stage(latents) | |
| wave = runner.vae.decode_to_waveform(mel)[0][:audio_len] | |
| outpath = f"output.wav" | |
| sf.write(outpath, wave, samplerate=16000, subtype='PCM_16') | |
| return outpath | |
| def preprocess(caption): | |
| output = preprocess_gemini(caption) | |
| return output, output | |
| def update_textbox(event_name, current_text): | |
| event = event_name + ' two times.' | |
| if current_text: | |
| return current_text.strip('.') + ' then ' + event | |
| else: | |
| return event | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| gr.Markdown("## PicoAudio") | |
| with gr.Row(): | |
| gr.Markdown(""" | |
| [](https://arxiv.org/abs/2407.02869v2) | |
| [](https://zeyuxie29.github.io/PicoAudio.github.io/) | |
| [](https://github.com/zeyuxie29/PicoAudio) | |
| [](https://huggingface.co/datasets/ZeyuXie/PicoAudio/tree/main) | |
| """) | |
| with gr.Row(): | |
| description_text = f"18 events supported:" | |
| gr.Markdown(description_text) | |
| btn_event = [] | |
| with gr.Row(): | |
| for i in range(6): | |
| event_name = f"{event_list[i]}" | |
| btn_event.append(gr.Button(event_name)) | |
| with gr.Row(): | |
| for i in range(6, 12): | |
| event_name = f"{event_list[i]}" | |
| btn_event.append(gr.Button(event_name)) | |
| with gr.Row(): | |
| for i in range(12, 18): | |
| event_name = f"{event_list[i]}" | |
| btn_event.append(gr.Button(event_name)) | |
| with gr.Row(): | |
| gr.Markdown("## Step1-Preprocess") | |
| with gr.Row(): | |
| preprocess_description_text = f"Transfer free-text into timestamp caption via LLM. "+\ | |
| "This demo uses Gemini as the preprocessor. If any errors occur, please try a few more times. "+\ | |
| "We also provide the GPT version consistent with the paper in the file 'Files/llm_reprocessing.py'. You can use your own api_key to modify and run 'Files/inference.py' for local inference." | |
| gr.Markdown(preprocess_description_text) | |
| with gr.Row(): | |
| with gr.Column(): | |
| freetext_prompt = gr.Textbox(label="Free-text Prompt: Input your free-text caption here. (e.g. a dog barks three times.)", | |
| value="a dog barks three times.",) | |
| with gr.Row(): | |
| preprocess_run_button = gr.Button() | |
| preprocess_run_clear = gr.ClearButton([freetext_prompt]) | |
| prompt = None | |
| with gr.Column(): | |
| freetext_prompt_out = gr.Textbox(label="Timestamp Caption: Preprocess output") | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Examples( | |
| examples = [["spraying two times then gunshot three times."], | |
| ["a dog barks three times."], | |
| ["cow mooing two times."],], | |
| inputs = [freetext_prompt], | |
| outputs = [prompt] | |
| ) | |
| with gr.Column(): | |
| pass | |
| with gr.Row(): | |
| gr.Markdown("## Step2-Generate") | |
| with gr.Row(): | |
| generate_description_text = f"Generate audio based on timestamp caption." | |
| gr.Markdown(generate_description_text) | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="Timestamp Caption: Specify your timestamp caption formatted as 'event1 at onset1-offset1_onset2-offset2 and event2 at onset1-offset1'.", | |
| value="spraying at 0.38-1.176_3.06-3.856 and gunshot at 1.729-3.729_4.367-6.367_7.031-9.031.",) | |
| with gr.Row(): | |
| generate_run_button = gr.Button() | |
| generate_run_clear = gr.ClearButton([prompt]) | |
| with gr.Accordion("Advanced options", open=False): | |
| num_steps = gr.Slider(label="num_steps", minimum=1, maximum=300, value=200, step=1) | |
| guidance_scale = gr.Slider(label="guidance_scale", minimum=0.1, maximum=8.0, value=3.0, step=0.1) | |
| with gr.Column(): | |
| outaudio = gr.Audio() | |
| for i in range(18): | |
| event_name = f"{event_list[i]}" | |
| btn_event[i].click(fn=update_textbox, inputs=[gr.State(event_name), freetext_prompt], outputs=freetext_prompt) | |
| preprocess_run_button.click(fn=preprocess, inputs=[freetext_prompt], outputs=[prompt, freetext_prompt_out]) | |
| generate_run_button.click(fn=infer, inputs=[prompt, num_steps, guidance_scale], outputs=[outaudio]) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Examples( | |
| examples = [["spraying at 0.38-1.176_3.06-3.856 and gunshot at 1.729-3.729_4.367-6.367_7.031-9.031."], | |
| ["dog_barking at 0.562-2.562_4.25-6.25."], | |
| ["cow_mooing at 0.958-3.582_5.272-7.896."],], | |
| inputs = [prompt, num_steps, guidance_scale], | |
| outputs = [outaudio] | |
| ) | |
| with gr.Column(): | |
| pass | |
| demo.launch() | |
| # description_text = f"18 events: {', '.join(event_list)}" | |
| # prompt = gr.Textbox(label="Prompt: Input your caption formatted as 'event1 at onset1-offset1_onset2-offset2 and event2 at onset1-offset1'.", | |
| # value="spraying at 0.38-1.176_3.06-3.856 and gunshot at 1.729-3.729_4.367-6.367_7.031-9.031.",) | |
| # outaudio = gr.Audio() | |
| # num_steps = gr.Slider(label="num_steps", minimum=1, maximum=300, value=200, step=1) | |
| # guidance_scale = gr.Slider(label="guidance_scale", minimum=0.1, maximum=8.0, value=3.0, step=0.1) | |
| # gr_interface = gr.Interface( | |
| # fn=infer, | |
| # inputs=[prompt, num_steps, guidance_scale], | |
| # outputs=[outaudio], | |
| # title="PicoAudio", | |
| # description=description_text, | |
| # allow_flagging=False, | |
| # examples=[ | |
| # ["spraying at 0.38-1.176_3.06-3.856 and gunshot at 1.729-3.729_4.367-6.367_7.031-9.031."], | |
| # ["dog_barking at 0.562-2.562_4.25-6.25."], | |
| # ["cow_mooing at 0.958-3.582_5.272-7.896."], | |
| # ], | |
| # cache_examples="lazy", # Turn on to cache. | |
| # ) | |
| # gr_interface.queue(10).launch() |