use pth rather than ckpt

ans.py

@@ -7,21 +7,21 @@ import soundfile as sf
 from inference import load_config_and_state_dict, load_generator, process_audio
 
 MSS_MODEL_PATHS = {
-    'vox': ['checkpoints/vox_mss.ckpt'],
-    'gtr': ['checkpoints/gtr_mss.ckpt'],
-    'key': ['checkpoints/key_mss.ckpt'],
-    'syn': ['checkpoints/syn_mss.ckpt', 'checkpoints/syn_mss1.ckpt'],
-    'bass': ['checkpoints/bass_mss.ckpt'],
-    'drums': ['checkpoints/drums_mss.ckpt', 'checkpoints/drums_mss1.ckpt'],
-    'perc': ['checkpoints/perc_mss.ckpt', 'checkpoints/perc_mss1.ckpt'],
-    'orch': ['checkpoints/orch_mss.ckpt', 'checkpoints/orch_mss1.ckpt'],
+    'vox': ['checkpoints/vox_mss.pth'],
+    'gtr': ['checkpoints/gtr_mss.pth'],
+    'key': ['checkpoints/key_mss.pth'],
+    'syn': ['checkpoints/syn_mss.pth', 'checkpoints/syn_mss1.pth'],
+    'bass': ['checkpoints/bass_mss.pth'],
+    'drums': ['checkpoints/drums_mss.pth', 'checkpoints/drums_mss1.pth'],
+    'perc': ['checkpoints/perc_mss.pth', 'checkpoints/perc_mss1.pth'],
+    'orch': ['checkpoints/orch_mss.pth', 'checkpoints/orch_mss1.pth'],
 }
 MSS_MODELS = {key: [] for key in MSS_MODEL_PATHS.keys()}
 
-PRE_MODEL_PATHS = {'vox': ['checkpoints/denoise.ckpt']}
+PRE_MODEL_PATHS = {'vox': ['checkpoints/denoise.pth']}
 PRE_MODELS = {key: [] for key in PRE_MODEL_PATHS.keys()}
     
-POST_MODEL_PATHS = {key: ['checkpoints/dereverb.ckpt'] if key in ['vox'] else [] for key in MSS_MODEL_PATHS.keys()}
+POST_MODEL_PATHS = {key: ['checkpoints/dereverb.pth'] if key in ['vox'] else [] for key in MSS_MODEL_PATHS.keys()}
 POST_MODELS = {key: [] for key in POST_MODEL_PATHS.keys()}
 
 OUTPUT_PATHS = {

app.py

@@ -14,43 +14,43 @@ REPO_ID = os.getenv("MODEL_REPO_ID", "chenxie95/xlance-msr-ckpt")
 # Instrument to checkpoint mapping
 INSTRUMENT_MAP = {
     'vox': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['vox_mss.ckpt'],
-        'post': ['dereverb.ckpt']
+        'pre': ['denoise.pth'],
+        'mss': ['vox_mss.pth'],
+        'post': ['dereverb.pth']
     },
     'gtr': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['gtr_mss.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['gtr_mss.pth'],
         'post': []
     },
     'key': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['key_mss.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['key_mss.pth'],
         'post': []
     },
     'syn': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['syn_mss.ckpt', 'syn_mss1.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['syn_mss.pth', 'syn_mss1.pth'],
         'post': []
     },
     'bass': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['bass_mss.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['bass_mss.pth'],
         'post': []
     },
     'drums': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['drums_mss.ckpt', 'drums_mss1.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['drums_mss.pth', 'drums_mss1.pth'],
         'post': []
     },
     'perc': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['perc_mss.ckpt', 'perc_mss1.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['perc_mss.pth', 'perc_mss1.pth'],
         'post': []
     },
     'orch': {
-        'pre': ['denoise.ckpt'],
-        'mss': ['orch_mss.ckpt', 'orch_mss1.ckpt'],
+        'pre': ['denoise.pth'],
+        'mss': ['orch_mss.pth', 'orch_mss1.pth'],
         'post': []
     }
 }

inference.py

@@ -14,6 +14,7 @@ import soundfile as sf
 import numpy as np
 from tqdm import tqdm
 import librosa
+import yaml
 from models import MelRNN, MelRoFormer, UNet, UFormer
 from models.bs_roformer import bs_roformer as BSRoformer
 from models.bs_roformer import mel_band_roformer as MelBandRoformer
@@ -55,9 +56,14 @@ class RoformerSequential(nn.Sequential):
         return self[-1](mixture, target) # also pass target if present
 
 def load_config_and_state_dict(path: str, map_location: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
-    if path.endswith('.pth'):
-        raise ValueError("Use .ckpt files instead of .pth files")
     print(f"Extracting state dict from {path}")
+    if path.endswith('.pth'):
+        model_name = Path(path).stem
+        config_path = f"./configs/{model_name}.yaml" # use config file with same name as model in ./configs
+        print(f"Loading config from {config_path}")
+        with open(config_path, 'r') as f:
+            config = yaml.load(f, Loader=yaml.FullLoader)
+        return config, torch.load(path, map_location=map_location)
     full_checkpoint = torch.load(path, map_location=map_location, weights_only=False)
     full_state_dict = full_checkpoint['state_dict']
     generator_state_dict = OrderedDict()
@@ -119,9 +125,9 @@ def process_audio(config, audio: np.ndarray, generator: nn.Module, device: str =
 
 def main():
     parser = argparse.ArgumentParser(description="Run inference on audio files using trained generator")
-    parser.add_argument("--checkpoint", '-c', type=str, required=True, help="Path to unwrapped generator weights (.ckpt)")
-    parser.add_argument("--checkpoint_pre", '-p', type=str, help="pre-processing model checkpoint (.ckpt)")
-    parser.add_argument("--checkpoint_post", '-P', type=str, help="post-processing model checkpoint (.ckpt)")
+    parser.add_argument("--checkpoint", '-c', type=str, required=True, help="Path to unwrapped generator weights (.ckpt or .pth)")
+    parser.add_argument("--checkpoint_pre", '-p', type=str, help="pre-processing model checkpoint (.ckpt or .pth)")
+    parser.add_argument("--checkpoint_post", '-P', type=str, help="post-processing model checkpoint (.ckpt or .pth)")
     parser.add_argument("--input_dir", '-i', type=str, help="Directory containing input .flac files")
     parser.add_argument("--output_dir", '-o', type=str, help="Directory to save processed audio")
     parser.add_argument("--instrument", type=str, help="Instrument to process (Vox/Gtr/Kbs/Synth/Bass/Rhy_DK/Rhy_PERC/Orch)")

inference_full.py

@@ -171,9 +171,9 @@ def inference_main(args):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="Run inference on audio files using trained generator")
-    parser.add_argument("--checkpoint", '-c', nargs='*', default=[], type=str, help="model checkpoint (.ckpt)")
-    parser.add_argument("--checkpoint_pre", '-p', nargs='*', default=[], type=str, help="pre-processing model checkpoint (.ckpt)")
-    parser.add_argument("--checkpoint_post", '-P', nargs='*', default=[], type=str, help="post-processing model checkpoint (.ckpt)")
+    parser.add_argument("--checkpoint", '-c', nargs='*', default=[], type=str, help="model checkpoint (.ckpt or .pth)")
+    parser.add_argument("--checkpoint_pre", '-p', nargs='*', default=[], type=str, help="pre-processing model checkpoint (.ckpt or .pth)")
+    parser.add_argument("--checkpoint_post", '-P', nargs='*', default=[], type=str, help="post-processing model checkpoint (.ckpt or .pth)")
     parser.add_argument("--input_dir", '-i', type=str, help="Directory containing input files, or a single audio file")
     parser.add_argument("--output_dir", '-o', type=str, help="Directory to save processed audio, or a single audio file name")
     parser.add_argument("--device", type=str, default="cuda", help="Device to run inference on (cuda/cpu)")

start_script.sh

@@ -1,8 +1,8 @@
 # inference_full (for inference on full audio)
 python inference_full.py \
--p checkpoints/denoise.ckpt \
--c checkpoints/vox_mss.ckpt \
--P checkpoints/dereverb.ckpt \ # only vox need it
+-p checkpoints/denoise.pth \
+-c checkpoints/vox_mss.pth \
+-P checkpoints/dereverb.pth \ # only vox need it
 -i test/input \
 -o test/output \
 --device cuda \
@@ -10,8 +10,8 @@ python inference_full.py \
 
 # another example (eight instruments are supported in total)
 python inference_full.py \
--p checkpoints/denoise.ckpt \
--c checkpoints/drums_mss.ckpt checkpoints/drums_mss1.ckpt \ # use multiple checkpoints for some instruments
+-p checkpoints/denoise.pth \
+-c checkpoints/drums_mss.pth checkpoints/drums_mss1.pth \ # use multiple checkpoints for some instruments
 -i test/input \
 -o test/output \
 --device cuda \
@@ -39,4 +39,7 @@ python inference.py \
 -P logs/dereverb/vox/checkpoints/00000010.ckpt \
 -i OrganizersMixture/Vocals \
 -o Answer/Vocals/default \
---no-eval
+--no-eval
+
+# unwrap
+python unwrap.py -i checkpoints/bass_mss.ckpt -o checkpoints/bass_mss.pth