Upload 6 files

tools/gradio/app.py

@@ -0,0 +1,246 @@
+import sys
+import gradio as gr
+import json
+from datetime import datetime
+import yaml
+import time
+import re
+import os.path as op
+import torch
+from levo_inference_lowmem import LeVoInference
+
+EXAMPLE_LYRICS = """
+[intro-short]
+
+[verse]
+夜晚的街灯闪烁
+我漫步在熟悉的角落
+回忆像潮水般涌来
+你的笑容如此清晰
+在心头无法抹去
+那些曾经的甜蜜
+如今只剩我独自回忆
+
+[verse]
+手机屏幕亮起
+是你发来的消息
+简单的几个字
+却让我泪流满面
+曾经的拥抱温暖
+如今却变得遥远
+我多想回到从前
+重新拥有你的陪伴
+
+[chorus]
+回忆的温度还在
+你却已不在
+我的心被爱填满
+却又被思念刺痛
+音乐的节奏奏响
+我的心却在流浪
+没有你的日子
+我该如何继续向前
+
+[outro-short]
+""".strip()
+
+APP_DIR = op.dirname(op.dirname(op.dirname(op.abspath(__file__))))
+MODEL = LeVoInference(sys.argv[1])
+with open(op.join(APP_DIR, 'conf/vocab.yaml'), 'r', encoding='utf-8') as file:
+    STRUCTS = yaml.safe_load(file)
+
+
+def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_coef=None, temperature=None, top_k=None, gen_type="mixed", progress=gr.Progress(track_tqdm=True)):
+    global MODEL
+    global STRUCTS
+    params = {'cfg_coef':cfg_coef, 'temperature':temperature, 'top_k':top_k}
+    params = {k:v for k,v in params.items() if v is not None}
+    vocal_structs = ['[verse]', '[chorus]', '[bridge]']
+    sample_rate = MODEL.cfg.sample_rate
+    
+    # format lyric
+    lyric = lyric.replace("[intro]", "[intro-short]").replace("[inst]", "[inst-short]").replace("[outro]", "[outro-short]")
+    paragraphs = [p.strip() for p in lyric.strip().split('\n\n') if p.strip()]
+    if len(paragraphs) < 1:
+        return None, json.dumps("Lyrics can not be left blank")
+    paragraphs_norm = []
+    vocal_flag = False
+    for para in paragraphs:
+        lines = para.splitlines()
+        struct_tag = lines[0].strip().lower()
+        if struct_tag not in STRUCTS:
+            return None, json.dumps(f"Segments should start with a structure tag in {STRUCTS}")
+        if struct_tag in vocal_structs:
+            vocal_flag = True
+            if len(lines) < 2 or not [line.strip() for line in lines[1:] if line.strip()]:
+                return None, json.dumps("The following segments require lyrics: [verse], [chorus], [bridge]")
+            else:
+                new_para_list = []
+                for line in lines[1:]:
+                    new_para_list.append(re.sub(r"[^\w\s\[\]\-\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af\u00c0-\u017f]", "", line))
+                new_para_str = f"{struct_tag} {'.'.join(new_para_list)}"
+        else:
+            if len(lines) > 1:
+                return None, json.dumps("The following segments should not contain lyrics: [intro], [intro-short], [intro-medium], [inst], [inst-short], [inst-medium], [outro], [outro-short], [outro-medium]")
+            else:
+                new_para_str = struct_tag
+        paragraphs_norm.append(new_para_str)
+    if not vocal_flag:
+        return None, json.dumps(f"The lyric must contain at least one of the following structures: {vocal_structs}")
+    lyric_norm = " ; ".join(paragraphs_norm)
+
+    # format prompt 
+    if prompt_audio is not None:
+        genre = None
+        description = None
+    elif description is not None and description != "":
+        genre = None
+        if description[-1] != ".":
+            description = description + "."
+
+    progress(0.0, "Start Generation")
+    start = time.time()
+    
+    audio_data = MODEL(lyric_norm, description, prompt_audio, genre, op.join(APP_DIR, "tools/new_auto_prompt.pt"), gen_type, params).cpu().permute(1, 0).float().numpy()
+
+    end = time.time()
+    
+    # 创建输入配置的JSON
+    input_config = {
+        "lyric": lyric_norm,
+        "genre": genre,
+        "prompt_audio": prompt_audio,
+        "description": description,
+        "params": params,
+        "inference_duration": end - start,
+        "timestamp": datetime.now().isoformat(),
+    }
+    
+    return (sample_rate, audio_data), json.dumps(input_config, indent=2)
+
+
+# 创建Gradio界面
+with gr.Blocks(title="SongGeneration Demo Space") as demo:
+    gr.Markdown("# 🎵 SongGeneration Demo Space")
+    gr.Markdown("Demo interface for the song generation model. Provide a lyrics, and optionally an audio or text prompt, to generate a custom song. The code is in [GIT](https://github.com/tencent-ailab/SongGeneration)")
+    
+    with gr.Row():
+        with gr.Column():
+            lyric = gr.Textbox(
+                label="Lyrics",
+                lines=5,
+                max_lines=15,
+                value=EXAMPLE_LYRICS,
+                info="Each paragraph represents a segment starting with a structure tag and ending with a blank line, each line is a sentence without punctuation, segments [intro], [inst], [outro] should not contain lyrics, while [verse], [chorus], and [bridge] require lyrics.",
+                placeholder="""Lyric Format
+'''
+[structure tag]
+lyrics
+
+[structure tag]
+lyrics
+'''
+1. One paragraph represents one segments, starting with a structure tag and ending with a blank line
+2. One line represents one sentence, punctuation is not recommended inside the sentence
+3. The following segments should not contain lyrics: [intro-short], [intro-medium], [inst-short], [inst-medium], [outro-short], [outro-medium]
+4. The following segments require lyrics: [verse], [chorus], [bridge]
+"""
+            )
+
+            with gr.Tabs(elem_id="extra-tabs"):
+                with gr.Tab("Genre Select"):
+                    genre = gr.Radio(
+                        choices=["Auto", "Pop", "Latin", "Rock", "Electronic", "Metal", "Country", "R&B/Soul", "Ballad", "Jazz", "World", "Hip-Hop", "Funk", "Soundtrack"],
+                        label="Genre Select(Optional)",
+                        value="Pop",
+                        interactive=True,
+                        elem_id="single-select-radio"
+                    )
+                with gr.Tab("Audio Prompt"):
+                    prompt_audio = gr.Audio(
+                        label="Prompt Audio (Optional)",
+                        type="filepath",
+                        elem_id="audio-prompt"
+                    )
+                with gr.Tab("Text Prompt"):
+                    gr.Markdown("For detailed usage, please refer to [here](https://github.com/tencent-ailab/SongGeneration?tab=readme-ov-file#-description-input-format)")
+                    description = gr.Textbox(
+                        label="Song Description (Optional)",
+                        info="Describe the gender, timbre, genre, emotion, instrument and bpm of the song. Only English is supported currently.​",
+                        placeholder="female, dark, pop, sad, piano and drums, the bpm is 125.",
+                        lines=1,
+                        max_lines=2
+                    )
+
+            with gr.Accordion("Advanced Config", open=False):
+                cfg_coef = gr.Slider(
+                    label="CFG Coefficient",
+                    minimum=0.1,
+                    maximum=3.0,
+                    step=0.1,
+                    value=1.5,
+                    interactive=True,
+                    elem_id="cfg-coef",
+                )
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=0.1,
+                    maximum=2.0,
+                    step=0.1,
+                    value=0.9,
+                    interactive=True,
+                    elem_id="temperature",
+                )
+                top_k = gr.Slider(
+                    label="Top-K",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=50,
+                    interactive=True,
+                    elem_id="top_k",
+                )
+            with gr.Row():
+                generate_btn = gr.Button("Generate Song", variant="primary")
+                generate_bgm_btn = gr.Button("Generate Pure Music", variant="primary")
+        
+        with gr.Column():
+            output_audio = gr.Audio(label="Generated Song", type="numpy")
+            output_json = gr.JSON(label="Generated Info")
+    
+        # # 示例按钮
+        # examples = gr.Examples(
+        #     examples=[
+        #         ["male, bright, rock, happy, electric guitar and drums, the bpm is 150."],
+        #         ["female, warm, jazz, romantic, synthesizer and piano, the bpm is 100."]
+        #     ],
+        #     inputs=[description],
+        #     label="Text Prompt examples"
+        # )
+
+        # examples = gr.Examples(
+        #     examples=[
+        #     "[intro-medium]\n\n[verse]\n在这个疯狂的世界里\n谁不渴望一点改变\n在爱情面前\n我们都显得那么不安全\n你紧紧抱着我\n告诉我再靠近一点\n别让这璀璨的夜晚白白浪费\n我那迷茫的眼睛\n看不见未来的路\n在情感消散之前\n我们对爱的渴望永不熄灭\n你给我留下一句誓言\n想知道我们的爱是否能持续到永远\n[chorus]\n\n约定在那最后的夜晚\n不管命运如何摆布\n我们的心是否依然如初\n我会穿上红衬衫\n带着摇滚的激情\n回到我们初遇的地方\n约定在那最后的夜晚\n就算全世界都变了样\n我依然坚守诺言\n铭记这一天\n你永远是我心中的爱恋\n\n[outro-medium]\n",
+        #     "[intro-short]\n\n[verse]\nThrough emerald canyons where fireflies dwell\nCerulean berries kiss morning's first swell\nCrystalline dew crowns each Vitamin Dawn's confection dissolves slowly on me\nAmbrosia breezes through honeycomb vines\nNature's own candy in Fibonacci lines\n[chorus] Blueberry fruit so sweet\n takes you higher\n can't be beat\n In your lungs\n it starts to swell\n You're under its spell\n [verse] Resin of sunlight in candied retreat\nMarmalade moonbeams melt under bare feet\nNectar spirals bloom chloroplast champagne\nPhotosynthesis sings through my veins\nChlorophyll rhythms pulse warm in my blood\nThe forest's green pharmacy floods every bud[chorus] Blueberry fruit so sweet\n takes you higher\n can't be beat\n In your lungs\n it starts to swell\n You're under its spell\n feel the buzz\n ride the wave\n Limey me\n blueberry\n your mind's enslaved\n In the haze\n lose all time\n floating free\n feeling fine\n Blueberry\n fruit so sweet\n takes you higher\n can't be beat\n In your lungs\n it starts to swell\n cry\n You're under its spell\n\n[outro-short]\n",
+        #     ],
+        #     inputs=[lyric],
+        #     label="Lyrics examples",
+        # )
+    
+    # 生成按钮点击事件
+    generate_btn.click(
+        fn=generate_song,
+        inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k],
+        outputs=[output_audio, output_json]
+    )
+    generate_bgm_btn.click(
+        fn=generate_song,
+        inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k, gr.State("bgm")],
+        outputs=[output_audio, output_json]
+    )
+    
+
+# 启动应用
+if __name__ == "__main__":
+    torch.set_num_threads(1)
+    demo.launch(server_name="0.0.0.0", server_port=8081)

tools/gradio/levo_inference.py

@@ -0,0 +1,126 @@
+import os
+import re
+import sys
+
+import torch
+
+import json
+import numpy as np
+from omegaconf import OmegaConf
+
+from codeclm.trainer.codec_song_pl import CodecLM_PL
+from codeclm.models import CodecLM
+
+from separator import Separator
+
+
+def check_language_by_text(text):
+    chinese_pattern = re.compile(r'[\u4e00-\u9fff]')
+    english_pattern = re.compile(r'[a-zA-Z]')
+    chinese_count = len(re.findall(chinese_pattern, text))
+    english_count = len(re.findall(english_pattern, text))
+    chinese_ratio = chinese_count / len(text)
+    english_ratio = english_count / len(text)
+    if chinese_ratio >= 0.2:
+        return "zh"
+    elif english_ratio >= 0.5:
+        return "en"
+    else:
+        return "en"
+
+
+class LeVoInference(torch.nn.Module):
+    def __init__(self, ckpt_path):
+        super().__init__()
+
+        torch.backends.cudnn.enabled = False 
+        OmegaConf.register_new_resolver("eval", lambda x: eval(x))
+        OmegaConf.register_new_resolver("concat", lambda *x: [xxx for xx in x for xxx in xx])
+        OmegaConf.register_new_resolver("get_fname", lambda: 'default')
+        OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
+
+        cfg_path = os.path.join(ckpt_path, 'config.yaml')
+        pt_path = os.path.join(ckpt_path, 'model.pt')
+
+        self.cfg = OmegaConf.load(cfg_path)
+        self.cfg.mode = 'inference'
+        self.max_duration = self.cfg.max_dur
+
+        # Define model or load pretrained model
+        model_light = CodecLM_PL(self.cfg, pt_path, version="v2")
+
+        model_light = model_light.eval().cuda()
+        model_light.audiolm.cfg = self.cfg
+
+        self.model_lm = model_light.audiolm
+        self.model_audio_tokenizer = model_light.audio_tokenizer
+        self.model_seperate_tokenizer = model_light.seperate_tokenizer
+
+        self.model = CodecLM(name = "tmp",
+            lm = self.model_lm,
+            audiotokenizer = self.model_audio_tokenizer,
+            max_duration = self.max_duration,
+            seperate_tokenizer = self.model_seperate_tokenizer,
+        )
+        self.separator = Separator()
+
+
+        self.default_params = dict(
+            cfg_coef = 1.5,
+            temperature = 1.0,
+            top_k = 50,
+            top_p = 0.0,
+            record_tokens = True,
+            record_window = 50,
+            extend_stride = 5,
+            duration = self.max_duration,
+        )
+
+        self.model.set_generation_params(**self.default_params)
+
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "mixed", params = dict()):
+        params = {**self.default_params, **params}
+        self.model.set_generation_params(**params)
+
+        if prompt_audio_path is not None and os.path.exists(prompt_audio_path):
+            pmt_wav, vocal_wav, bgm_wav = self.separator.run(prompt_audio_path)
+            melody_is_wav = True
+        elif genre is not None and auto_prompt_path is not None:
+            auto_prompt = torch.load(auto_prompt_path)
+            lang = check_language_by_text(lyric)
+            prompt_token = auto_prompt[genre][lang][np.random.randint(0, len(auto_prompt[genre][lang]))]
+            pmt_wav = prompt_token[:,[0],:]
+            vocal_wav = prompt_token[:,[1],:]
+            bgm_wav = prompt_token[:,[2],:]
+            melody_is_wav = False
+        else:
+            pmt_wav = None
+            vocal_wav = None
+            bgm_wav = None
+            melody_is_wav = True
+
+        if gen_type == 'bgm':
+            description = '[Musicality-very-high]' + ', ' + '[Pure-Music]' + ', ' + description.lower() if description else '.'
+        else:
+            description = description.lower() if description else '.'
+            description = '[Musicality-very-high]' + ', ' + description
+
+        generate_inp = {
+            'lyrics': [lyric.replace("  ", " ")] if gen_type != 'bgm' else '.',
+            'descriptions': [description],
+            'melody_wavs': pmt_wav,
+            'vocal_wavs': vocal_wav,
+            'bgm_wavs': bgm_wav,
+            'melody_is_wav': melody_is_wav,
+        }
+
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            tokens = self.model.generate(**generate_inp, return_tokens=True)
+            
+        with torch.no_grad():
+            if melody_is_wav:
+                wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
+            else:
+                wav_seperate = self.model.generate_audio(tokens, gen_type=gen_type)
+
+        return wav_seperate[0]

tools/gradio/levo_inference_lowmem.py

@@ -0,0 +1,187 @@
+import os
+import gc
+import re
+import sys
+
+import torch
+
+import json
+import numpy as np
+from omegaconf import OmegaConf
+
+from codeclm.trainer.codec_song_pl import CodecLM_PL
+from codeclm.models import CodecLM
+from codeclm.models import builders
+
+from separator import Separator
+from codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
+
+
+def check_language_by_text(text):
+    chinese_pattern = re.compile(r'[\u4e00-\u9fff]')
+    english_pattern = re.compile(r'[a-zA-Z]')
+    chinese_count = len(re.findall(chinese_pattern, text))
+    english_count = len(re.findall(english_pattern, text))
+    chinese_ratio = chinese_count / len(text)
+    english_ratio = english_count / len(text)
+    if chinese_ratio >= 0.2:
+        return "zh"
+    elif english_ratio >= 0.5:
+        return "en"
+    else:
+        return "en"
+
+
+class LeVoInference(torch.nn.Module):
+    def __init__(self, ckpt_path):
+        super().__init__()
+
+        torch.backends.cudnn.enabled = False 
+        OmegaConf.register_new_resolver("eval", lambda x: eval(x))
+        OmegaConf.register_new_resolver("concat", lambda *x: [xxx for xx in x for xxx in xx])
+        OmegaConf.register_new_resolver("get_fname", lambda: 'default')
+        OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
+
+        cfg_path = os.path.join(ckpt_path, 'config.yaml')
+        self.pt_path = os.path.join(ckpt_path, 'model.pt')
+
+        self.cfg = OmegaConf.load(cfg_path)
+        self.cfg.mode = 'inference'
+        self.max_duration = self.cfg.max_dur
+
+        self.default_params = dict(
+            top_p = 0.0,
+            record_tokens = True,
+            record_window = 50,
+            extend_stride = 5,
+            duration = self.max_duration,
+        )
+
+
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "mixed", params = dict()):
+        if prompt_audio_path is not None and os.path.exists(prompt_audio_path):
+            separator = Separator()
+            audio_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint, self.cfg)
+            audio_tokenizer = audio_tokenizer.eval().cuda()
+            pmt_wav, vocal_wav, bgm_wav = separator.run(prompt_audio_path)
+            pmt_wav = pmt_wav.cuda()
+            vocal_wav = vocal_wav.cuda()
+            bgm_wav = bgm_wav.cuda()
+            with torch.no_grad():
+                pmt_wav, _ = audio_tokenizer.encode(pmt_wav)
+            del audio_tokenizer
+            del separator
+            torch.cuda.empty_cache()
+
+            seperate_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint_sep, self.cfg)
+            seperate_tokenizer = seperate_tokenizer.eval().cuda()
+            with torch.no_grad():
+                vocal_wav, bgm_wav = seperate_tokenizer.encode(vocal_wav, bgm_wav)
+            del seperate_tokenizer
+            melody_is_wav = False
+            torch.cuda.empty_cache()
+        elif genre is not None and auto_prompt_path is not None:
+            auto_prompt = torch.load(auto_prompt_path)
+            lang = check_language_by_text(lyric)
+            prompt_token = auto_prompt[genre][lang][np.random.randint(0, len(auto_prompt[genre][lang]))]
+            pmt_wav = prompt_token[:,[0],:]
+            vocal_wav = prompt_token[:,[1],:]
+            bgm_wav = prompt_token[:,[2],:]
+            melody_is_wav = False
+        else:
+            pmt_wav = None
+            vocal_wav = None
+            bgm_wav = None
+            melody_is_wav = True
+
+        audiolm = builders.get_lm_model(self.cfg, version="v2")
+        checkpoint = torch.load(self.pt_path, map_location='cpu')
+        audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
+        audiolm.load_state_dict(audiolm_state_dict, strict=False)
+        audiolm = audiolm.eval()
+
+        offload_audiolm = True if 'offload' in self.cfg.keys() and 'audiolm' in self.cfg.offload else False
+        if offload_audiolm:
+            audiolm_offload_param = OffloadParamParse.parse_config(audiolm, self.cfg.offload.audiolm)
+            audiolm_offload_param.show()
+            offload_profiler = OffloadProfiler(device_index=0, **(audiolm_offload_param.init_param_dict()))
+            offload_profiler.offload_layer(**(audiolm_offload_param.offload_layer_param_dict()))
+            offload_profiler.clean_cache_wrapper(**(audiolm_offload_param.clean_cache_param_dict()))
+        else:
+            audiolm = audiolm.cuda().to(torch.float16)
+
+        model = CodecLM(name = "tmp",
+            lm = audiolm,
+            audiotokenizer = None,
+            max_duration = self.max_duration,
+            seperate_tokenizer = None,
+        )
+        params = {**self.default_params, **params}
+        model.set_generation_params(**params)
+
+        if gen_type == 'bgm':
+            description = '[Musicality-very-high]' + ', ' + '[Pure-Music]' + ', ' + description.lower() if description else '.'
+        else:
+            description = description.lower() if description else '.'
+            description = '[Musicality-very-high]' + ', ' + description
+        
+        generate_inp = {
+            'lyrics': [lyric.replace("  ", " ")] if gen_type != 'bgm' else '.',
+            'descriptions': [description],
+            'melody_wavs': pmt_wav,
+            'vocal_wavs': vocal_wav,
+            'bgm_wavs': bgm_wav,
+            'melody_is_wav': melody_is_wav,
+        }
+
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            with torch.no_grad():
+                tokens = model.generate(**generate_inp, return_tokens=True)
+                if offload_audiolm:
+                    offload_profiler.reset_empty_cache_mem_line()
+        offload_profiler.stop()
+        del offload_profiler
+        del audiolm_offload_param
+        del model
+        audiolm = audiolm.cpu()
+        del audiolm
+        del checkpoint
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        seperate_tokenizer = builders.get_audio_tokenizer_model_cpu(self.cfg.audio_tokenizer_checkpoint_sep, self.cfg)
+        device = "cuda:0"
+        seperate_tokenizer.model.device = device
+        seperate_tokenizer.model.vae = seperate_tokenizer.model.vae.to(device)
+        seperate_tokenizer.model.model.device = torch.device(device)
+        seperate_tokenizer = seperate_tokenizer.eval()
+
+        offload_wav_tokenizer_diffusion =  True if 'offload' in self.cfg.keys() and 'wav_tokenizer_diffusion' in self.cfg.offload else False
+        if offload_wav_tokenizer_diffusion:
+            sep_offload_param = OffloadParamParse.parse_config(seperate_tokenizer, self.cfg.offload.wav_tokenizer_diffusion)
+            sep_offload_param.show()
+            sep_offload_profiler = OffloadProfiler(device_index=0, **(sep_offload_param.init_param_dict()))
+            sep_offload_profiler.offload_layer(**(sep_offload_param.offload_layer_param_dict()))
+            sep_offload_profiler.clean_cache_wrapper(**(sep_offload_param.clean_cache_param_dict()))
+        else:
+            seperate_tokenizer.model.model = seperate_tokenizer.model.model.to(device)
+
+        model = CodecLM(name = "tmp",
+            lm = None,
+            audiotokenizer = None,
+            max_duration = self.max_duration,
+            seperate_tokenizer = seperate_tokenizer,
+        )
+
+        with torch.no_grad():
+            if melody_is_wav:
+                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type, chunked=True)
+            else:
+                wav_seperate = model.generate_audio(tokens, gen_type=gen_type, chunked=True)
+
+        if offload_wav_tokenizer_diffusion:
+            sep_offload_profiler.reset_empty_cache_mem_line()
+            sep_offload_profiler.stop()
+        torch.cuda.empty_cache()
+
+        return wav_seperate[0]

tools/gradio/run.sh

@@ -0,0 +1,13 @@
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export CUDA_LAUNCH_BLOCKING=0
+
+export USER=root
+export PYTHONDONTWRITEBYTECODE=1
+export TRANSFORMERS_CACHE="$(pwd)/third_party/hub"
+export NCCL_HOME=/usr/local/tccl
+export PYTHONPATH="$(pwd)/codeclm/tokenizer/":"$(pwd)":"$(pwd)/codeclm/tokenizer/Flow1dVAE/":"$(pwd)/codeclm/tokenizer/":$PYTHONPATH
+
+
+CKPT_PATH=$1
+python3 tools/gradio/app.py $CKPT_PATH

tools/gradio/separator.py

@@ -0,0 +1,50 @@
+import torchaudio
+import os
+import torch
+from third_party.demucs.models.pretrained import get_model_from_yaml
+
+
+class Separator(torch.nn.Module):
+    def __init__(self, dm_model_path='third_party/demucs/ckpt/htdemucs.pth', dm_config_path='third_party/demucs/ckpt/htdemucs.yaml', gpu_id=0) -> None:
+        super().__init__()
+        if torch.cuda.is_available() and gpu_id < torch.cuda.device_count():
+            self.device = torch.device(f"cuda:{gpu_id}")
+        else:
+            self.device = torch.device("cpu")
+        self.demucs_model = self.init_demucs_model(dm_model_path, dm_config_path)
+
+    def init_demucs_model(self, model_path, config_path):
+        model = get_model_from_yaml(config_path, model_path)
+        model.to(self.device)
+        model.eval()
+        return model
+    
+    def load_audio(self, f):
+        a, fs = torchaudio.load(f)
+        if (fs != 48000):
+            a = torchaudio.functional.resample(a, fs, 48000)
+        if a.shape[-1] >= 48000*10:
+            a = a[..., :48000*10]
+        else:
+            a = torch.cat([a, a], -1)
+        return a[:, 0:48000*10]
+    
+    def run(self, audio_path, output_dir='tmp', ext=".flac"):
+        os.makedirs(output_dir, exist_ok=True)
+        name, _ = os.path.splitext(os.path.split(audio_path)[-1])
+        output_paths = []
+
+        for stem in self.demucs_model.sources:
+            output_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
+            if os.path.exists(output_path):
+                output_paths.append(output_path)
+        if len(output_paths) == 1:  # 4
+            vocal_path = output_paths[0]
+        else:
+            drums_path, bass_path, other_path, vocal_path = self.demucs_model.separate(audio_path, output_dir, device=self.device)
+            for path in [drums_path, bass_path, other_path]:
+                os.remove(path)
+        full_audio = self.load_audio(audio_path)
+        vocal_audio = self.load_audio(vocal_path)
+        bgm_audio = full_audio - vocal_audio
+        return full_audio, vocal_audio, bgm_audio

tools/new_auto_prompt.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15d65fd20de92de44fa3b8b867d376d4c1756e56b18c9edf4c235ab66a3dae2e
+size 133