-
Notifications
You must be signed in to change notification settings - Fork 2.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
audiogen_app.py #185
Open
cocktailpeanut
wants to merge
6
commits into
facebookresearch:main
Choose a base branch
from
cocktailpeanut:main
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
audiogen_app.py #185
Changes from all commits
Commits
Show all changes
6 commits
Select commit
Hold shift + click to select a range
216e90c
no xformers
cocktailpeanut 1f92d7f
audiogen
cocktailpeanut 0a75e71
remove multiband for now
cocktailpeanut 852db42
music => audio
cocktailpeanut bb8ff0d
Merge branch 'facebookresearch:main' into main
cocktailpeanut a6b772f
remove requirements-no-xformers.txt
cocktailpeanut File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,249 @@ | ||
# Copyright (c) Meta Platforms, Inc. and affiliates. | ||
# All rights reserved. | ||
|
||
# This source code is licensed under the license found in the | ||
# LICENSE file in the root directory of this source tree. | ||
|
||
# Updated to account for UI changes from https://github.com/rkfg/audiocraft/blob/long/app.py | ||
# also released under the MIT license. | ||
|
||
import argparse | ||
from concurrent.futures import ProcessPoolExecutor | ||
import os | ||
from pathlib import Path | ||
import subprocess as sp | ||
from tempfile import NamedTemporaryFile | ||
import time | ||
import typing as tp | ||
import warnings | ||
|
||
import torch | ||
import gradio as gr | ||
|
||
from audiocraft.data.audio_utils import convert_audio | ||
from audiocraft.data.audio import audio_write | ||
from audiocraft.models import AudioGen, MultiBandDiffusion | ||
|
||
|
||
MODEL = None # Last used model | ||
INTERRUPTING = False | ||
# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform | ||
_old_call = sp.call | ||
|
||
|
||
def _call_nostderr(*args, **kwargs): | ||
# Avoid ffmpeg vomiting on the logs. | ||
kwargs['stderr'] = sp.DEVNULL | ||
kwargs['stdout'] = sp.DEVNULL | ||
_old_call(*args, **kwargs) | ||
|
||
|
||
sp.call = _call_nostderr | ||
# Preallocating the pool of processes. | ||
pool = ProcessPoolExecutor(4) | ||
pool.__enter__() | ||
|
||
|
||
def interrupt(): | ||
global INTERRUPTING | ||
INTERRUPTING = True | ||
|
||
|
||
class FileCleaner: | ||
def __init__(self, file_lifetime: float = 3600): | ||
self.file_lifetime = file_lifetime | ||
self.files = [] | ||
|
||
def add(self, path: tp.Union[str, Path]): | ||
self._cleanup() | ||
self.files.append((time.time(), Path(path))) | ||
|
||
def _cleanup(self): | ||
now = time.time() | ||
for time_added, path in list(self.files): | ||
if now - time_added > self.file_lifetime: | ||
if path.exists(): | ||
path.unlink() | ||
self.files.pop(0) | ||
else: | ||
break | ||
|
||
|
||
file_cleaner = FileCleaner() | ||
|
||
|
||
def make_waveform(*args, **kwargs): | ||
# Further remove some warnings. | ||
be = time.time() | ||
with warnings.catch_warnings(): | ||
warnings.simplefilter('ignore') | ||
out = gr.make_waveform(*args, **kwargs) | ||
print("Make a video took", time.time() - be) | ||
return out | ||
|
||
|
||
def load_model(version='facebook/audiogen-medium'): | ||
global MODEL | ||
print("Loading model", version) | ||
if MODEL is None or MODEL.name != version: | ||
MODEL = AudioGen.get_pretrained(version) | ||
|
||
|
||
def load_diffusion(): | ||
global MBD | ||
print("loading MBD") | ||
MBD = MultiBandDiffusion.get_mbd_musicgen() | ||
|
||
|
||
def _do_predictions(texts, duration, progress=False, **gen_kwargs): | ||
MODEL.set_generation_params(duration=duration, **gen_kwargs) | ||
be = time.time() | ||
target_sr = 32000 | ||
target_ac = 1 | ||
|
||
outputs = MODEL.generate(texts, progress=progress) | ||
if USE_DIFFUSION: | ||
outputs_diffusion = MBD.tokens_to_wav(outputs[1]) | ||
outputs = torch.cat([outputs[0], outputs_diffusion], dim=0) | ||
outputs = outputs.detach().cpu().float() | ||
pending_videos = [] | ||
out_wavs = [] | ||
for output in outputs: | ||
with NamedTemporaryFile("wb", suffix=".wav", delete=False) as file: | ||
audio_write( | ||
file.name, output, MODEL.sample_rate, strategy="loudness", | ||
loudness_headroom_db=16, loudness_compressor=True, add_suffix=False) | ||
pending_videos.append(pool.submit(make_waveform, file.name)) | ||
out_wavs.append(file.name) | ||
file_cleaner.add(file.name) | ||
out_videos = [pending_video.result() for pending_video in pending_videos] | ||
for video in out_videos: | ||
file_cleaner.add(video) | ||
print("batch finished", len(texts), time.time() - be) | ||
print("Tempfiles currently stored: ", len(file_cleaner.files)) | ||
return out_videos, out_wavs | ||
|
||
|
||
|
||
def predict_full(model, decoder, text, duration, topk, topp, temperature, cfg_coef, progress=gr.Progress()): | ||
global INTERRUPTING | ||
global USE_DIFFUSION | ||
INTERRUPTING = False | ||
if temperature < 0: | ||
raise gr.Error("Temperature must be >= 0.") | ||
if topk < 0: | ||
raise gr.Error("Topk must be non-negative.") | ||
if topp < 0: | ||
raise gr.Error("Topp must be non-negative.") | ||
|
||
topk = int(topk) | ||
if decoder == "MultiBand_Diffusion": | ||
USE_DIFFUSION = True | ||
load_diffusion() | ||
else: | ||
USE_DIFFUSION = False | ||
load_model(model) | ||
|
||
def _progress(generated, to_generate): | ||
progress((min(generated, to_generate), to_generate)) | ||
if INTERRUPTING: | ||
raise gr.Error("Interrupted.") | ||
MODEL.set_custom_progress_callback(_progress) | ||
|
||
videos, wavs = _do_predictions( | ||
[text], duration, progress=True, | ||
top_k=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef) | ||
if USE_DIFFUSION: | ||
return videos[0], wavs[0], videos[1], wavs[1] | ||
return videos[0], wavs[0], None, None | ||
return videos[0], wavs[0] | ||
|
||
|
||
|
||
def toggle_diffusion(choice): | ||
if choice == "MultiBand_Diffusion": | ||
return [gr.update(visible=True)] * 2 | ||
else: | ||
return [gr.update(visible=False)] * 2 | ||
|
||
|
||
def ui_full(launch_kwargs): | ||
with gr.Blocks() as interface: | ||
gr.Markdown( | ||
""" | ||
# AudioGen | ||
This is your private demo for [AudioGen](https://github.com/facebookresearch/audiocraft/blob/main/docs/AUDIOGEN.md), | ||
a simple and controllable model for audio generation | ||
""" | ||
) | ||
with gr.Row(): | ||
with gr.Column(): | ||
with gr.Row(): | ||
text = gr.Text(label="Input Text", interactive=True) | ||
with gr.Row(): | ||
submit = gr.Button("Submit") | ||
# Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license. | ||
_ = gr.Button("Interrupt").click(fn=interrupt, queue=False) | ||
with gr.Row(): | ||
model = gr.Radio(["facebook/audiogen-medium"], label="Model", value="facebook/audiogen-medium", interactive=True) | ||
with gr.Row(): | ||
decoder = gr.Radio(["Default"], label="Decoder", value="Default", interactive=False) | ||
with gr.Row(): | ||
duration = gr.Slider(minimum=1, maximum=120, value=10, label="Duration", interactive=True) | ||
with gr.Row(): | ||
topk = gr.Number(label="Top-k", value=250, interactive=True) | ||
topp = gr.Number(label="Top-p", value=0, interactive=True) | ||
temperature = gr.Number(label="Temperature", value=1.0, interactive=True) | ||
cfg_coef = gr.Number(label="Classifier Free Guidance", value=3.0, interactive=True) | ||
with gr.Column(): | ||
output = gr.Video(label="Generated Audio") | ||
audio_output = gr.Audio(label="Generated Audio (wav)", type='filepath') | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Would be great to add some examples as well similarly to the MusicGen demo for users to get some ideas about possible prompts. |
||
submit.click(predict_full, inputs=[model, decoder, text, duration, topk, topp, temperature, cfg_coef], outputs=[output, audio_output]) | ||
|
||
interface.queue().launch(**launch_kwargs) | ||
|
||
|
||
|
||
if __name__ == "__main__": | ||
parser = argparse.ArgumentParser() | ||
parser.add_argument( | ||
'--listen', | ||
type=str, | ||
default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1', | ||
help='IP to listen on for connections to Gradio', | ||
) | ||
parser.add_argument( | ||
'--username', type=str, default='', help='Username for authentication' | ||
) | ||
parser.add_argument( | ||
'--password', type=str, default='', help='Password for authentication' | ||
) | ||
parser.add_argument( | ||
'--server_port', | ||
type=int, | ||
default=0, | ||
help='Port to run the server listener on', | ||
) | ||
parser.add_argument( | ||
'--inbrowser', action='store_true', help='Open in browser' | ||
) | ||
parser.add_argument( | ||
'--share', action='store_true', help='Share the gradio UI' | ||
) | ||
|
||
args = parser.parse_args() | ||
|
||
launch_kwargs = {} | ||
launch_kwargs['server_name'] = args.listen | ||
|
||
if args.username and args.password: | ||
launch_kwargs['auth'] = (args.username, args.password) | ||
if args.server_port: | ||
launch_kwargs['server_port'] = args.server_port | ||
if args.inbrowser: | ||
launch_kwargs['inbrowser'] = args.inbrowser | ||
if args.share: | ||
launch_kwargs['share'] = args.share | ||
|
||
# Show the interface | ||
ui_full(launch_kwargs) |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We don't have a MultiBandDiffusion model for AudioGen, let's remove this bit from the demo for the merge?