fix 37 (#38)

Co-authored-by: Tuan Tran <{ID}+{username}@users.noreply.github.com>

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to AudioSeal are documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+## [0.1.3] - 2024-06-24
+
+- Update scripts to new training code
+- Fix bugs in loading custom fine-tuned model (https://github.com/facebookresearch/audioseal/issues/37)
+
 ## [0.1.3] - 2024-04-30
 
 - Fix bug in getting the watermark with non-empty message created in CPU, while the model is loaded in CUDA

src/audioseal/builder.py

@@ -4,8 +4,8 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-from dataclasses import asdict, dataclass, is_dataclass
-from typing import Any, Dict, List, Mapping, Optional
+from dataclasses import asdict, dataclass, field, is_dataclass
+from typing import Any, Dict, List, Optional
 
 from omegaconf import DictConfig, OmegaConf
 from torch import device, dtype
@@ -55,7 +55,7 @@ class DecoderConfig:
 
 @dataclass
 class DetectorConfig:
-    output_dim: int
+    output_dim: int = 32
 
 
 @dataclass
@@ -69,7 +69,7 @@ class AudioSealWMConfig:
 class AudioSealDetectorConfig:
     nbits: int
     seanet: SEANetConfig
-    detector: DetectorConfig
+    detector: DetectorConfig = field(default_factory=lambda: DetectorConfig())
 
 
 def as_dict(obj: Any) -> Dict[str, Any]:

src/audioseal/models.py

@@ -17,7 +17,7 @@
 COMPATIBLE_WARNING = """
 AudioSeal is designed to work at a sample rate 16khz.
 Implicit sampling rate usage is deprecated and will be removed in future version.
-To remove this warning please add this argument to the function call: 
+To remove this warning please add this argument to the function call:
 sample_rate = your_sample_rate
 """
 
@@ -111,7 +111,9 @@ def get_watermark(
         if self.msg_processor is not None:
             if message is None:
                 if self.message is None:
-                    message = torch.randint(0, 2, (x.shape[0], self.msg_processor.nbits), device=x.device)
+                    message = torch.randint(
+                        0, 2, (x.shape[0], self.msg_processor.nbits), device=x.device
+                    )
                 else:
                     message = self.message.to(device=x.device)
             else:

src/scripts/checkpoints.py

@@ -10,8 +10,9 @@
 import torch
 
 
-def convert(checkpoint: Path, outdir: Path, suffix: str = "base"):
+def convert(checkpoint: str, outdir: str, suffix: str = "base"):
     """Convert the checkpoint to generator and detector"""
+    outdir_path = Path(outdir)
     ckpt = torch.load(checkpoint)
 
     # keep inference-related params only
@@ -27,16 +28,21 @@ def convert(checkpoint: Path, outdir: Path, suffix: str = "base"):
 
     for layer in ckpt["model"].keys():
         if layer.startswith("detector"):
-            detector_ckpt["model"][layer] = ckpt["model"][layer]
+            new_layer = layer[9:]
+            detector_ckpt["model"][new_layer] = ckpt["model"][layer]  # type: ignore
         elif layer == "msg_processor.msg_processor.0.weight":
-            generator_ckpt["model"]["msg_processor.msg_processor.weight"] = ckpt[
+            generator_ckpt["model"]["msg_processor.msg_processor.weight"] = ckpt[  # type: ignore
                 "model"
-            ][layer]
+            ][
+                layer
+            ]
         else:
-            generator_ckpt["model"][layer] = ckpt["model"][layer]
+            assert layer.startswith("generator"), f"Invalid layer: {layer}"
+            new_layer = layer[10:]
+            generator_ckpt["model"][new_layer] = ckpt["model"][layer]  # type: ignore
 
-    torch.save(generator_ckpt, outdir / (checkpoint.stem + f"_generator_{suffix}.pth"))
-    torch.save(detector_ckpt, outdir / (checkpoint.stem + f"_detector_{suffix}.pth"))
+    torch.save(generator_ckpt, outdir_path / (f"checkpoint_generator_{suffix}.pth"))
+    torch.save(detector_ckpt, outdir_path / (f"checkpoint_detector_{suffix}.pth"))
 
 
 if __name__ == "__main__":