Update Clotho model

app.py

@@ -23,7 +23,7 @@ def load_model(cfg,
     return model, tokenizer
 
 
-def infer(file, device, model, tokenizer, target_sr):
+def infer(file, runner):
     sr, wav = file
     wav = torch.as_tensor(wav)
     if wav.dtype == torch.short:
@@ -32,9 +32,9 @@ def infer(file, device, model, tokenizer, target_sr):
         wav = wav / 2 ** 31
     if wav.ndim > 1:
         wav = wav.mean(1)
-    wav = resample(wav, sr, target_sr)
+    wav = resample(wav, sr, runner.target_sr)
     wav_len = len(wav)
-    wav = wav.float().unsqueeze(0).to(device)
+    wav = wav.float().unsqueeze(0).to(runner.device)
     input_dict = {
         "mode": "inference",
         "wav": wav,
@@ -44,9 +44,9 @@ def infer(file, device, model, tokenizer, target_sr):
         "beam_size": 3,
     }
     with torch.no_grad():
-        output_dict = model(input_dict)
+        output_dict = runner.model(input_dict)
         seq = output_dict["seq"].cpu().numpy()
-        cap = tokenizer.decode(seq)[0]
+        cap = runner.tokenizer.decode(seq)[0]
     return cap
 
 # def input_toggle(input_type):
@@ -55,43 +55,47 @@ def infer(file, device, model, tokenizer, target_sr):
 #     elif input_type == "mic":
 #         return gr.update(visible=False), gr.update(visible=True)
 
+class InferRunner:
 
-if __name__ == "__main__":
+    def __init__(self, model_name):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        exp_dir = Path(f"./checkpoints/{model_name.lower()}")
+        cfg = train_util.load_config(exp_dir / "config.yaml")
+        self.model, self.tokenizer = load_model(cfg, exp_dir / "ckpt.pth", self.device)
+        self.target_sr = cfg["target_sr"]
     
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--share", action="store_true", default=False)
+    def change_model(self, model_name):
+        exp_dir = Path(f"./checkpoints/{model_name.lower()}")
+        cfg = train_util.load_config(exp_dir / "config.yaml")
+        self.model, self.tokenizer = load_model(cfg, exp_dir / "ckpt.pth", self.device)
+        self.target_sr = cfg["target_sr"]
 
-    args = parser.parse_args()
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    exp_dir = Path("./checkpoints/audiocaps")
-    cfg = train_util.load_config(exp_dir / "config.yaml")
-    target_sr = cfg["target_sr"]
-    model, tokenizer = load_model(cfg, exp_dir / "ckpt.pth", device)
+def change_model(radio):
+    global infer_runner
+    infer_runner.change_model(radio)
 
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                # radio = gr.Radio(
-                #     ["file", "mic"],
-                #     value="file",
-                #     label="Select input type"
-                # )
-                file = gr.Audio(label="Input", visible=True)
-                # mic = gr.Microphone(label="Input", visible=False)
-                # radio.change(fn=input_toggle, inputs=radio, outputs=[file, mic])
-                btn = gr.Button("Run")
-            with gr.Column():
-                output = gr.Textbox(label="Output")
-            btn.click(
-                fn=partial(infer,
-                           device=device,
-                           model=model,
-                           tokenizer=tokenizer,
-                           target_sr=target_sr),
-                inputs=[file,],
-                outputs=output
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            radio = gr.Radio(
+                ["AudioCaps", "Clotho"],
+                value="AudioCaps",
+                label="Select model"
             )
-        
-        demo.launch(share=args.share)
+            infer_runner = InferRunner(radio.value)
+            file = gr.Audio(label="Input", visible=True)
+            radio.change(fn=change_model, inputs=[radio,],)
+            btn = gr.Button("Run")
+        with gr.Column():
+            output = gr.Textbox(label="Output")
+        btn.click(
+            fn=partial(infer,
+                       runner=infer_runner),
+            inputs=[file,],
+            outputs=output
+        )
+    
+    demo.launch()
 

checkpoints/clotho/ckpt.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:694c9e7139be7ec5aff2153d1af980d6bc305403a76be0d8940481579ea51483
+size 54651005

checkpoints/clotho/config.yaml

@@ -0,0 +1,30 @@
+tokenizer:
+    type: text_tokenizer.DictTokenizer
+    args:
+        max_length: 20
+
+target_sr: 16000
+
+model:
+    args:
+        shared_dim: 1024
+        tchr_dim: 768
+    model:
+        args: {}
+        decoder:
+            args:
+                attn_emb_dim: 1408
+                dropout: 0.2
+                emb_dim: 256
+                fc_emb_dim: 1408
+                nlayers: 2
+                tie_weights: true
+                vocab_size: 4368
+            type: models.transformer_decoder.TransformerDecoder
+        encoder:
+            args:
+                freeze: false
+                pretrained: true
+            type: models.cnn_encoder.EfficientNetB2
+        type: models.transformer_model.TransformerModel
+    type: models.kd_wrapper.ContraEncoderKdWrapper

models/eff_latent_encoder.py

@@ -17,7 +17,7 @@ from einops import rearrange, reduce
 from torch.hub import load_state_dict_from_url
 
 
-model_dir = "./"
+model_dir = os.getcwd()
 
 
 class _EffiNet(nn.Module):

utils/model_util.py

@@ -11,7 +11,7 @@ def sort_pack_padded_sequence(input, lengths):
     sorted_lengths, indices = torch.sort(lengths, descending=True)
     tmp = pack_padded_sequence(input[indices], sorted_lengths.cpu(), batch_first=True)
     inv_ix = indices.clone()
-    inv_ix[indices] = torch.arange(0,len(indices)).type_as(inv_ix)
+    inv_ix[indices] = torch.arange(0, len(indices)).type_as(inv_ix)
     return tmp, inv_ix
 
 def pad_unsort_packed_sequence(input, inv_ix):

utils/train_util.py

@@ -80,7 +80,7 @@ def merge_load_state_dict(state_dict,
             pretrained_dict[key] = value
         else:
             mismatch_keys.append(key)
-    output_fn(f"Loading pre-trained model, with mismatched keys {mismatch_keys}")
+    output_fn(f"Loading pre-trained model, with mismatched keys {mismatch_keys}\n")
     model_dict.update(pretrained_dict)
     model.load_state_dict(model_dict, strict=True)
     return pretrained_dict.keys()