ZeroTwo3/videoshop-backend
2
1import argparse2from concurrent.futures import ProcessPoolExecutor3import os4from pathlib import Path5import subprocess as sp6from tempfile import NamedTemporaryFile7import time8import typing as tp9import warnings10 11import torch12import gradio as gr13 14from audiocraft.data.audio_utils import convert_audio15from audiocraft.data.audio import audio_write16from audiocraft.models import MusicGen, MultiBandDiffusion17 18 19MODEL = None # Last used model20IS_BATCHED = "facebook/MusicGen" in os.environ.get('SPACE_ID', '')21print(IS_BATCHED)22MAX_BATCH_SIZE = 1223BATCHED_DURATION = 1524INTERRUPTING = False25MBD = None26# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform27_old_call = sp.call28 29 30def _call_nostderr(*args, **kwargs):31 # Avoid ffmpeg vomiting on the logs.32 kwargs['stderr'] = sp.DEVNULL33 kwargs['stdout'] = sp.DEVNULL34 _old_call(*args, **kwargs)35 36 37sp.call = _call_nostderr38# Preallocating the pool of processes.39pool = ProcessPoolExecutor(4)40pool.__enter__()41 42 43def interrupt():44 global INTERRUPTING45 INTERRUPTING = True46 47 48class FileCleaner:49 def __init__(self, file_lifetime: float = 3600):50 self.file_lifetime = file_lifetime51 self.files = []52 53 def add(self, path: tp.Union[str, Path]):54 self._cleanup()55 self.files.append((time.time(), Path(path)))56 57 def _cleanup(self):58 now = time.time()59 for time_added, path in list(self.files):60 if now - time_added > self.file_lifetime:61 if path.exists():62 path.unlink()63 self.files.pop(0)64 else:65 break66 67 68file_cleaner = FileCleaner()69 70 71def make_waveform(*args, **kwargs):72 # Further remove some warnings.73 be = time.time()74 with warnings.catch_warnings():75 warnings.simplefilter('ignore')76 out = gr.make_waveform(*args, **kwargs)77 print("Make a video took", time.time() - be)78 return out79 80 81def load_model(version='facebook/musicgen-melody'):82 global MODEL83 print("Loading model", version)84 if MODEL is None or MODEL.name != version:85 MODEL = MusicGen.get_pretrained(version)86 87 88def load_diffusion():89 global MBD90 if MBD is None:91 print("loading MBD")92 MBD = MultiBandDiffusion.get_mbd_musicgen()93 94 95def _do_predictions(texts, melodies, duration, progress=False, **gen_kwargs):96 MODEL.set_generation_params(duration=duration, **gen_kwargs)97 print("new batch", len(texts), texts, [None if m is None else (m[0], m[1].shape) for m in melodies])98 be = time.time()99 processed_melodies = []100 target_sr = 32000101 target_ac = 1102 for melody in melodies:103 if melody is None:104 processed_melodies.append(None)105 else:106 sr, melody = melody[0], torch.from_numpy(melody[1]).to(MODEL.device).float().t()107 if melody.dim() == 1:108 melody = melody[None]109 melody = melody[..., :int(sr * duration)]110 melody = convert_audio(melody, sr, target_sr, target_ac)111 processed_melodies.append(melody)112 113 if any(m is not None for m in processed_melodies):114 outputs = MODEL.generate_with_chroma(115 descriptions=texts,116 melody_wavs=processed_melodies,117 melody_sample_rate=target_sr,118 progress=progress,119 return_tokens=USE_DIFFUSION120 )121 else:122 outputs = MODEL.generate(texts, progress=progress, return_tokens=USE_DIFFUSION)123 if USE_DIFFUSION:124 outputs_diffusion = MBD.tokens_to_wav(outputs[1])125 outputs = torch.cat([outputs[0], outputs_diffusion], dim=0)126 outputs = outputs.detach().cpu().float()127 pending_videos = []128 out_wavs = []129 for output in outputs:130 with NamedTemporaryFile("wb", suffix=".wav", delete=False) as file:131 audio_write(132 file.name, output, MODEL.sample_rate, strategy="loudness",133 loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)134 pending_videos.append(pool.submit(make_waveform, file.name))135 out_wavs.append(file.name)136 file_cleaner.add(file.name)137 out_videos = [pending_video.result() for pending_video in pending_videos]138 for video in out_videos:139 file_cleaner.add(video)140 print("batch finished", len(texts), time.time() - be)141 print("Tempfiles currently stored: ", len(file_cleaner.files))142 return out_videos, out_wavs143 144 145def predict_batched(texts, melodies):146 max_text_length = 512147 texts = [text[:max_text_length] for text in texts]148 load_model('facebook/musicgen-melody')149 res = _do_predictions(texts, melodies, BATCHED_DURATION)150 return res151 152 153def predict_full(model, decoder, text, melody, duration, topk, topp, temperature, cfg_coef, progress=gr.Progress()):154 global INTERRUPTING155 global USE_DIFFUSION156 INTERRUPTING = False157 if temperature < 0:158 raise gr.Error("Temperature must be >= 0.")159 if topk < 0:160 raise gr.Error("Topk must be non-negative.")161 if topp < 0:162 raise gr.Error("Topp must be non-negative.")163 164 topk = int(topk)165 if decoder == "MultiBand_Diffusion":166 USE_DIFFUSION = True167 load_diffusion()168 else:169 USE_DIFFUSION = False170 load_model(model)171 172 def _progress(generated, to_generate):173 progress((min(generated, to_generate), to_generate))174 if INTERRUPTING:175 raise gr.Error("Interrupted.")176 MODEL.set_custom_progress_callback(_progress)177 178 videos, wavs = _do_predictions(179 [text], [melody], duration, progress=True,180 top_k=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef)181 if USE_DIFFUSION:182 return videos[0], wavs[0], videos[1], wavs[1]183 return videos[0], wavs[0], None, None184 185 186def toggle_audio_src(choice):187 if choice == "mic":188 return gr.update(source="microphone", value=None, label="Microphone")189 else:190 return gr.update(source="upload", value=None, label="File")191 192 193def toggle_diffusion(choice):194 if choice == "MultiBand_Diffusion":195 return [gr.update(visible=True)] * 2196 else:197 return [gr.update(visible=False)] * 2198 199 200def ui_full(launch_kwargs):201 with gr.Blocks() as interface:202 203 with gr.Row():204 with gr.Column():205 with gr.Row():206 text = gr.Text(label="Input Text", interactive=True)207 # with gr.Column():208 # radio = gr.Radio(["file", "mic"], value="file",209 # label="Condition on a melody (optional) File or Mic")210 # melody = gr.Audio(source="upload", type="numpy", label="File",211 # interactive=True, elem_id="melody-input")212 with gr.Row():213 submit = gr.Button("Submit")214 # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.215 _ = gr.Button("Interrupt").click(fn=interrupt, queue=False)216 with gr.Row():217 model = gr.Radio(["facebook/musicgen-melody", "facebook/musicgen-medium", "facebook/musicgen-small",218 "facebook/musicgen-large"],219 label="Model", value="facebook/musicgen-melody", interactive=True)220 # with gr.Row():221 # decoder = gr.Radio(["Default", "MultiBand_Diffusion"],222 # label="Decoder", value="Default", interactive=True)223 # decoder = "Default"224 with gr.Row():225 duration = gr.Slider(minimum=1, maximum=120, value=10, label="Duration", interactive=True)226 # with gr.Row():227 # topk = gr.Number(label="Top-k", value=250, interactive=True)228 # topp = gr.Number(label="Top-p", value=0, interactive=True)229 # temperature = gr.Number(label="Temperature", value=1.0, interactive=True)230 # cfg_coef = gr.Number(label="Classifier Free Guidance", value=3.0, interactive=True)231 232 with gr.Column():233 output = gr.Video(label="Generated Music")234 audio_output = gr.Audio(label="Generated Music (wav)", type='filepath')235 # diffusion_output = gr.Video(label="MultiBand Diffusion Decoder")236 # audio_diffusion = gr.Audio(label="MultiBand Diffusion Decoder (wav)", type='filepath')237 238 melody = gr.Audio(source= None, type="numpy", label="File",239 interactive=False, visible= False, elem_id="melody-input")240 decoder = gr.Radio(["Default", "MultiBand_Diffusion"],241 label="Decoder", value="Default", interactive=True, visible= False)242 # duration = gr.Slider(minimum=1, maximum=120, value=10, label="Duration", interactive=True, visible= False)243 topk = gr.Number(label="Top-k", value=250, interactive=True, visible= False)244 topp = gr.Number(label="Top-p", value=0, interactive=True, visible= False)245 temperature = gr.Number(label="Temperature", value=1.0, interactive=True, visible= False)246 cfg_coef = gr.Number(label="Classifier Free Guidance", value=3.0, interactive=True, visible= False)247 diffusion_output = gr.Video(label="MultiBand Diffusion Decoder" , visible=False)248 audio_diffusion = gr.Audio(label="MultiBand Diffusion Decoder (wav)", type='filepath', visible= False)249 250 print("melody", melody)251 print("decoder", decoder)252 print("topk", topk)253 print("topp", topp)254 print("cfg_coef", cfg_coef)255 print("diffusion_output" , diffusion_output)256 print("audio_diffusion" , audio_diffusion)257 258 submit.click(toggle_diffusion, decoder, [diffusion_output, audio_diffusion], queue=False,259 show_progress=False).then(predict_full, inputs=[model, decoder, text, melody, duration, topk, topp,260 temperature, cfg_coef],261 outputs=[output, audio_output, diffusion_output, audio_diffusion])262 # radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)263 264 # gr.Examples(265 # fn=predict_full,266 # examples=[267 # [268 # "An 80s driving pop song with heavy drums and synth pads in the background",269 # "./assets/bach.mp3",270 # "facebook/musicgen-melody",271 # "Default"272 # ],273 # [274 # "A cheerful country song with acoustic guitars",275 # "./assets/bolero_ravel.mp3",276 # "facebook/musicgen-melody",277 # "Default"278 # ],279 # [280 # "90s rock song with electric guitar and heavy drums",281 # None,282 # "facebook/musicgen-medium",283 # "Default"284 # ],285 # [286 # "a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions",287 # "./assets/bach.mp3",288 # "facebook/musicgen-melody",289 # "Default"290 # ],291 # [292 # "lofi slow bpm electro chill with organic samples",293 # None,294 # "facebook/musicgen-medium",295 # "Default"296 # ],297 # [298 # "Punk rock with loud drum and power guitar",299 # None,300 # "facebook/musicgen-medium",301 # "MultiBand_Diffusion"302 # ],303 # ],304 # inputs=[text, melody, model, decoder],305 # outputs=[output]306 # )307 gr.Markdown(308 """309 """310 )311 312 interface.queue().launch(**launch_kwargs)313 314def ui_batched(launch_kwargs):315 with gr.Blocks() as demo:316 gr.Markdown(317 """318 This project generate Music from prompt.319 """320 )321 with gr.Row():322 with gr.Column():323 with gr.Row():324 text = gr.Text(label="Describe your music", lines=2, interactive=True)325 with gr.Column():326 radio = gr.Radio(["file", "mic"], value="file",327 label="Condition on a melody (optional) File or Mic")328 melody = gr.Audio(source="upload", type="numpy", label="File",329 interactive=True, elem_id="melody-input")330 with gr.Row():331 submit = gr.Button("Generate")332 with gr.Column():333 output = gr.Video(label="Generated Music")334 audio_output = gr.Audio(label="Generated Music (wav)", type='filepath')335 submit.click(predict_batched, inputs=[text, melody],336 outputs=[output, audio_output], batch=True, max_batch_size=MAX_BATCH_SIZE)337 radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)338 gr.Examples(339 fn=predict_batched,340 # examples=[341 # [342 # "An 80s driving pop song with heavy drums and synth pads in the background",343 # "./assets/bach.mp3",344 # ],345 # [346 # "A cheerful country song with acoustic guitars",347 # "./assets/bolero_ravel.mp3",348 # ],349 # [350 # "90s rock song with electric guitar and heavy drums",351 # None,352 # ],353 # [354 # "a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions bpm: 130",355 # "./assets/bach.mp3",356 # ],357 # [358 # "lofi slow bpm electro chill with organic samples",359 # None,360 # ],361 # ],362 examples=[363 364 ],365 inputs=[text, melody],366 outputs=[output]367 )368 gr.Markdown("""369 370 """)371 372 demo.queue(max_size=8 * 4).launch(**launch_kwargs)373 374 375if __name__ == "__main__":376 parser = argparse.ArgumentParser()377 parser.add_argument(378 '--listen',379 type=str,380 default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1',381 help='IP to listen on for connections to Gradio',382 )383 parser.add_argument(384 '--username', type=str, default='', help='Username for authentication'385 )386 parser.add_argument(387 '--password', type=str, default='', help='Password for authentication'388 )389 parser.add_argument(390 '--server_port',391 type=int,392 default=0,393 help='Port to run the server listener on',394 )395 parser.add_argument(396 '--inbrowser', action='store_true', help='Open in browser'397 )398 parser.add_argument(399 '--share', action='store_true', help='Share the gradio UI'400 )401 402 args = parser.parse_args()403 404 launch_kwargs = {}405 launch_kwargs['server_name'] = args.listen406 407 if args.username and args.password:408 launch_kwargs['auth'] = (args.username, args.password)409 if args.server_port:410 launch_kwargs['server_port'] = args.server_port411 if args.inbrowser:412 launch_kwargs['inbrowser'] = args.inbrowser413 if args.share:414 launch_kwargs['share'] = args.share415 416 # Show the interface417 if IS_BATCHED:418 global USE_DIFFUSION419 USE_DIFFUSION = False420 ui_batched(launch_kwargs)421 else:422 ui_full(launch_kwargs)423