Update app.py

app.py

@@ -1,351 +1,28 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under thmage license found in the
-# LICENSE file in the root directory of this source tree.
-
-import argparse
-from concurrent.futures import ProcessPoolExecutor
-import logging
-import os
-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 audiocraft.data.audio import audio_write
+import torchaudio
 from audiocraft.models import MAGNeT
+from audiocraft. data. audio import audio_write
 
-
-MODEL = None  # Last used model
-SPACE_ID = os.environ.get('SPACE_ID', '')
-MAX_BATCH_SIZE = 12
-N_REPEATS = 2
-INTERRUPTING = False
-MBD = None
-# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform
-_old_call = sp.call
-
-PROD_STRIDE_1 = "prod-stride1 (new!)"
-
-
-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/magnet-small-10secs'):
-    global MODEL
-    print("Loading model", version)
-    if MODEL is None or MODEL.name != version:
-        MODEL = None  # in case loading would crash
-        MODEL = MAGNeT.get_pretrained(version)
+model = MAGNeT.get_pretrained('facebook/magnet-small-10secs')
 
 
-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()
+def infer(description):
+    descriptions = ['disco beat', 'energetic EDM']
 
-    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(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()):
-    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)
+    wav = model.generate(descriptions) 
     
-    videos, wavs = _do_predictions(
-        [text] * N_REPEATS, 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)
-
-    outputs_ = [videos[0]] + [wav for wav in wavs]
-    return tuple(outputs_)
-
-def ui_full(launch_kwargs):
-    with gr.Blocks() as interface:
-        gr.Markdown(
-            """
-            # MAGNeT
-            This is your private demo for [MAGNeT](https://github.com/facebookresearch/audiocraft),
-            A fast text-to-music model, consists of a single, non-autoregressive transformer.
-            presented at: ["Masked Audio Generation using a Single Non-Autoregressive Transformer"] (https://huggingface.co/papers/2401.04577)
-            """
-        )
-        with gr.Row():
-            with gr.Column():
-                with gr.Row():
-                    text = gr.Text(label="Input Text", value="80s electronic track with melodic synthesizers, catchy beat and groovy bass", 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/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/magnet-small-10secs', interactive=True)
-                    model_path = gr.Text(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 - variation 1")
-                audio_outputs = [gr.Audio(label=f"Generated Audio - variation {i+1}", type='filepath') for i in range(N_REPEATS)]
-        submit.click(fn=predict_full, 
-                        inputs=[model, model_path, text, 
-                                    temperature, topp,
-                                    max_cfg_coef, min_cfg_coef,
-                                    decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4,
-                                    span_score],
-                                    outputs=[output] + [o for o in audio_outputs])
-        gr.Examples(
-            fn=predict_full,
-            examples=[
-                [
-                    "80s electronic track with melodic synthesizers, catchy beat and groovy bass",
-                    'facebook/magnet-small-10secs',
-                    20, 3.0, 0.9, 10.0,
-                ],
-                [
-                    "80s electronic track with melodic synthesizers, catchy beat and groovy bass. 170 bpm",
-                    'facebook/magnet-small-10secs',
-                    20, 3.0, 0.9, 10.0,
-                ],
-                [
-                    "Earthy tones, environmentally conscious, ukulele-infused, harmonic, breezy, easygoing, organic instrumentation, gentle grooves",
-                    'facebook/magnet-medium-10secs',
-                    20, 3.0, 0.9, 10.0,
-                ],
-                [   "Funky groove with electric piano playing blue chords rhythmically",
-                    'facebook/magnet-medium-10secs',
-                    20, 3.0, 0.9, 10.0,
-                ],
-                [
-                    "Rock with saturated guitars, a heavy bass line and crazy drum break and fills.",
-                    'facebook/magnet-small-30secs',
-                    60, 3.0, 0.9, 10.0,
-                ],
-                [   "A grand orchestral arrangement with thunderous percussion, epic brass fanfares, and soaring strings, creating a cinematic atmosphere fit for a heroic battle",
-                    'facebook/magnet-medium-30secs',
-                    60, 3.0, 0.9, 10.0,
-                ],
-                [   "Seagulls squawking as ocean waves crash while wind blows heavily into a microphone.",
-                    'facebook/audio-magnet-small', 
-                    20, 3.5, 0.8, 20.0,
-                ],
-                [   "A toilet flushing as music is playing and a man is singing in the distance.",
-                    'facebook/audio-magnet-medium', 
-                    20, 3.5, 0.8, 20.0,
-                ],
-            ],
-
-            inputs=[text, model, decoding_steps1, temperature, topp, max_cfg_coef],
-            outputs=[output]
-        )
-
-        gr.Markdown(
-            """
-            ### More details
-            
-            #### Music Generation
-            "magnet" models will generate a short music extract based on the textual description you provided.
-            These models can generate either 10 seconds or 30 seconds of music.
-            These models were trained with descriptions from a stock music catalog. Descriptions that will work best
-            should include some level of details on the instruments present, along with some intended use case
-            (e.g. adding "perfect for a commercial" can somehow help).
-
-            We present 4 model variants:
-            1. facebook/magnet-small-10secs - a 300M non-autoregressive transformer capable of generating 10-second music conditioned
-                on text.
-            2. facebook/magnet-medium-10secs - 1.5B parameters, 10 seconds audio.
-            3. facebook/magnet-small-30secs - 300M parameters, 30 seconds audio.
-            4. facebook/magnet-medium-30secs - 1.5B parameters, 30 seconds audio.
+    for idx, one_wav in enumerate(wav):
+        print(idx)
+        audio_write(f'{idx}', 
+                    one_wav.cpu(), 
+                    model.sample_rate, 
+                    strategy="loudness", 
+                    loudness_compressor=True)
         
-            #### Sound-Effect Generation
-            "audio-magnet" models will generate a 10-second sound effect based on the description you provide. 
-
-            These models were trained on the following data sources: a subset of AudioSet (Gemmeke et al., 2017), 
-            [BBC sound effects](https://sound-effects.bbcrewind.co.uk/), AudioCaps (Kim et al., 2019), 
-            Clotho v2 (Drossos et al., 2020), VGG-Sound (Chen et al., 2020), FSD50K (Fonseca et al., 2021), 
-            [Free To Use Sounds](https://www.freetousesounds.com/all-in-one-bundle/), [Sonniss Game Effects](https://sonniss.com/gameaudiogdc), 
-            [WeSoundEffects](https://wesoundeffects.com/we-sound-effects-bundle-2020/), 
-            [Paramount Motion - Odeon Cinematic Sound Effects](https://www.paramountmotion.com/odeon-sound-effects).
-
-            We present 2 model variants:
-            1. facebook/audio-magnet-small - 10 second sound effect generation, 300M parameters.
-            2. facebook/audio-magnet-medium - 10 second sound effect generation, 1.5B parameters.
-
-            See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft/blob/main/docs/MAGNET.md)
-            for more details.
-            """
-        )
-
-        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
-
-    logging.basicConfig(level=logging.INFO, stream=sys.stderr)
+    return "done"
 
-    # Show the interface
-    ui_full(launch_kwargs)
+gr.Interface(
+    fn = infer,
+    inputs = gr.Textbox(value="gogo"),
+    outputs = gr.Textbox()
+).launch()