|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
from concurrent.futures import ThreadPoolExecutor |
|
import logging |
|
import os |
|
import base64 |
|
from pathlib import Path |
|
import subprocess as sp |
|
import sys |
|
from tempfile import NamedTemporaryFile |
|
import time |
|
import typing as tp |
|
import warnings |
|
|
|
import gradio as gr |
|
from pydub import AudioSegment |
|
|
|
from audiocraft.data.audio import audio_write |
|
from audiocraft.models import MAGNeT |
|
|
|
|
|
SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret') |
|
|
|
MODEL = None |
|
SPACE_ID = os.environ.get('SPACE_ID', '') |
|
MAX_BATCH_SIZE = 12 |
|
N_REPEATS = 1 |
|
INTERRUPTING = False |
|
MBD = None |
|
|
|
_old_call = sp.call |
|
|
|
PROD_STRIDE_1 = "prod-stride1 (new!)" |
|
|
|
|
|
def _call_nostderr(*args, **kwargs): |
|
|
|
kwargs['stderr'] = sp.DEVNULL |
|
kwargs['stdout'] = sp.DEVNULL |
|
_old_call(*args, **kwargs) |
|
|
|
|
|
sp.call = _call_nostderr |
|
|
|
pool = ThreadPoolExecutor(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): |
|
|
|
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/magnet-small-10secs'): |
|
global MODEL |
|
print("Loading model", version) |
|
if MODEL is None or MODEL.name != version: |
|
MODEL = None |
|
MODEL = MAGNeT.get_pretrained(version) |
|
|
|
|
|
def _do_predictions(texts, progress=False, gradio_progress=None, **gen_kwargs): |
|
MODEL.set_generation_params(**gen_kwargs) |
|
print("new batch", len(texts), texts) |
|
be = time.time() |
|
|
|
try: |
|
outputs = MODEL.generate(texts, progress=progress, return_tokens=False) |
|
except RuntimeError as e: |
|
raise gr.Error("Error while generating " + e.args[0]) |
|
outputs = outputs.detach().cpu().float() |
|
pending_videos = [] |
|
out_wavs = [] |
|
for i, output in enumerate(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) |
|
if i == 0: |
|
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_batched(texts, melodies): |
|
max_text_length = 512 |
|
texts = [text[:max_text_length] for text in texts] |
|
load_model('facebook/magnet-small-10secs') |
|
res = _do_predictions(texts, melodies) |
|
return res |
|
|
|
def predict_full(secre_token, model, model_path, text, temperature, topp, |
|
max_cfg_coef, min_cfg_coef, |
|
decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4, |
|
span_score, |
|
progress=gr.Progress()): |
|
if secret_token != SECRET_TOKEN: |
|
raise gr.Error( |
|
f'Invalid secret token. Please fork the original space if you want to use it for yourself.') |
|
|
|
global INTERRUPTING |
|
INTERRUPTING = False |
|
progress(0, desc="Loading model...") |
|
model_path = model_path.strip() |
|
if model_path: |
|
if not Path(model_path).exists(): |
|
raise gr.Error(f"Model path {model_path} doesn't exist.") |
|
if not Path(model_path).is_dir(): |
|
raise gr.Error(f"Model path {model_path} must be a folder containing " |
|
"state_dict.bin and compression_state_dict_.bin.") |
|
model = model_path |
|
if temperature < 0: |
|
raise gr.Error("Temperature must be >= 0.") |
|
|
|
load_model(model) |
|
|
|
max_generated = 0 |
|
|
|
def _progress(generated, to_generate): |
|
nonlocal max_generated |
|
max_generated = max(generated, max_generated) |
|
progress((min(max_generated, to_generate), to_generate)) |
|
if INTERRUPTING: |
|
raise gr.Error("Interrupted.") |
|
MODEL.set_custom_progress_callback(_progress) |
|
|
|
videos, wavs = _do_predictions( |
|
[text], progress=True, |
|
temperature=temperature, top_p=topp, |
|
max_cfg_coef=max_cfg_coef, min_cfg_coef=min_cfg_coef, |
|
decoding_steps=[decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4], |
|
span_arrangement='stride1' if (span_score == PROD_STRIDE_1) else 'nonoverlap', |
|
gradio_progress=progress) |
|
|
|
wav_path = wavs[0] |
|
wav_base64 = "" |
|
|
|
|
|
|
|
mp3_path = wav_path.replace(".wav", ".mp3") |
|
sound = AudioSegment.from_wav(wav_path) |
|
sound.export(mp3_path, format="mp3") |
|
|
|
|
|
mp3_base64 = "" |
|
with open(mp3_path, "rb") as mp3_file: |
|
mp3_base64 = base64.b64encode(mp3_file.read()).decode('utf-8') |
|
|
|
|
|
mp3_base64_data_uri = 'data:audio/mp3;base64,' + mp3_base64 |
|
|
|
return mp3_base64_data_uri |
|
|
|
def ui_full(launch_kwargs): |
|
with gr.Blocks() as interface: |
|
gr.HTML(""" |
|
<div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;"> |
|
<div style="text-align: center; color: black;"> |
|
<p style="color: black;">This space is a headless component of the cloud rendering engine used by AiTube.</p> |
|
<p style="color: black;">It is not available for public use, but you can use the <a href="https://huggingface.co/spaces/doevent/AnimateLCM-SVD" target="_blank">original space</a>.</p> |
|
</div> |
|
</div>""") |
|
with gr.Row(): |
|
with gr.Column(): |
|
secret_token = gr.Textbox(label="Secret Token") |
|
with gr.Row(): |
|
text = gr.Textbox(label="Input Text", value="Downtown New York, busy street, pedestrian, taxis", interactive=True) |
|
with gr.Row(): |
|
submit = gr.Button("Submit") |
|
|
|
_ = gr.Button("Interrupt").click(fn=interrupt, queue=False) |
|
with gr.Row(): |
|
model = gr.Radio(['facebook/magnet-small-10secs', 'facebook/magnet-medium-10secs', |
|
'facebook/magnet-small-30secs', 'facebook/magnet-medium-30secs', |
|
'facebook/audio-magnet-small', 'facebook/audio-magnet-medium'], |
|
label="Model", value='facebook/audio-magnet-medium', interactive=True) |
|
model_path = gr.Textbox(label="Model Path (custom models)") |
|
with gr.Row(): |
|
span_score = gr.Radio(["max-nonoverlap", PROD_STRIDE_1], |
|
label="Span Scoring", value=PROD_STRIDE_1, interactive=True) |
|
with gr.Row(): |
|
decoding_steps1 = gr.Number(label="Decoding Steps (stage 1)", value=20, interactive=True) |
|
decoding_steps2 = gr.Number(label="Decoding Steps (stage 2)", value=10, interactive=True) |
|
decoding_steps3 = gr.Number(label="Decoding Steps (stage 3)", value=10, interactive=True) |
|
decoding_steps4 = gr.Number(label="Decoding Steps (stage 4)", value=10, interactive=True) |
|
with gr.Row(): |
|
temperature = gr.Number(label="Temperature", value=3.0, step=0.25, minimum=0, interactive=True) |
|
topp = gr.Number(label="Top-p", value=0.9, step=0.1, minimum=0, maximum=1, interactive=True) |
|
max_cfg_coef = gr.Number(label="Max CFG coefficient", value=10.0, minimum=0, interactive=True) |
|
min_cfg_coef = gr.Number(label="Min CFG coefficient", value=1.0, minimum=0, interactive=True) |
|
with gr.Column(): |
|
output = gr.Video(label="Generated Audio") |
|
base64_audio_output = gr.Textbox() |
|
submit.click(fn=predict_full, |
|
inputs=[secret_token, model, model_path, text, |
|
temperature, topp, |
|
max_cfg_coef, min_cfg_coef, |
|
decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4, |
|
span_score], |
|
outputs=base64_audio_output) |
|
|
|
interface.queue(max_size=10).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 |
|
|
|
logging.basicConfig(level=logging.INFO, stream=sys.stderr) |
|
|
|
|
|
ui_full(launch_kwargs) |