Add stripped inference-only model code mirror

clean/audio/nes2net/SOURCE.md

@@ -0,0 +1,17 @@
+# Source: audio/nes2net
+
+| Field | Value |
+|---|---|
+| Upstream | **UNVERIFIED** -- provenance was lost when this code was vendored |
+| Paper | not recorded |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | no license file present upstream (all rights reserved) |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/audio__nes2net.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.

clean/audio/nes2net/wav2vec2_Nes2Net_X.py

@@ -0,0 +1,317 @@
+import math
+
+import fairseq
+import torch
+import torch.nn as nn
+
+___author__ = "Tianchi Liu"
+__email__ = "tianchi_liu@u.nus.edu"
+# modified from the model script from Hemlata Tak
+
+
+class SSLModel(nn.Module):
+    def __init__(self, device):
+        super(SSLModel, self).__init__()
+        cp_path = (
+            "/app/weights/xlsr2_300m.pt"  # Change the pre-trained XLSR model path.
+        )
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        self.model = model[0]
+        self.device = device
+        self.out_dim = 1024
+        return
+
+    def extract_feat(self, input_data):
+        # put the model to GPU if it not there
+        if (
+            next(self.model.parameters()).device != input_data.device
+            or next(self.model.parameters()).dtype != input_data.dtype
+        ):
+            self.model.to(input_data.device, dtype=input_data.dtype)
+            self.model.train()
+        if True:
+            # input should be in shape (batch, length)
+            if input_data.ndim == 3:
+                input_tmp = input_data[:, :, 0]
+            else:
+                input_tmp = input_data
+            # [batch, length, dim]
+            emb = self.model(input_tmp, mask=False, features_only=True)["x"]
+        return emb
+
+
+class SEModule(nn.Module):
+    def __init__(self, channels, SE_ratio=8):
+        super(SEModule, self).__init__()
+        self.se = nn.Sequential(
+            nn.AdaptiveAvgPool1d(1),
+            nn.Conv1d(channels, channels // SE_ratio, kernel_size=1, padding=0),
+            nn.ReLU(),
+            nn.Conv1d(channels // SE_ratio, channels, kernel_size=1, padding=0),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input):
+        x = self.se(input)
+        return input * x
+
+
+class Bottle2neck(nn.Module):
+
+    def __init__(
+        self, inplanes, planes, kernel_size=None, dilation=None, scale=8, SE_ratio=8
+    ):
+        super(Bottle2neck, self).__init__()
+        width = int(math.floor(planes / scale))
+        self.conv1 = nn.Conv1d(inplanes, width * scale, kernel_size=1)
+        self.bn1 = nn.BatchNorm1d(width * scale)
+        self.nums = scale - 1
+        convs = []
+        bns = []
+        weighted_sum = []
+        num_pad = math.floor(kernel_size / 2) * dilation
+        for i in range(self.nums):
+            convs.append(
+                nn.Conv2d(
+                    width,
+                    width,
+                    kernel_size=(kernel_size, 1),
+                    dilation=(dilation, 1),
+                    padding=(num_pad, 0),
+                )
+            )
+            bns.append(nn.BatchNorm2d(width))
+            initial_value = torch.ones(1, 1, 1, i + 2) * (1 / (i + 2))
+            weighted_sum.append(nn.Parameter(initial_value, requires_grad=True))
+        self.weighted_sum = nn.ParameterList(weighted_sum)
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList(bns)
+        self.conv3 = nn.Conv1d(width * scale, planes, kernel_size=1)
+        self.bn3 = nn.BatchNorm1d(planes)
+        self.relu = nn.ReLU()
+        self.width = width
+        self.se = SEModule(planes, SE_ratio)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.relu(out)
+        out = self.bn1(out).unsqueeze(-1)  # bz c T 1
+
+        spx = torch.split(out, self.width, 1)
+        sp = spx[self.nums]
+        for i in range(self.nums):
+            sp = torch.cat((sp, spx[i]), -1)
+
+            sp = self.bns[i](self.relu(self.convs[i](sp)))
+            sp_s = sp * self.weighted_sum[i]
+            sp_s = torch.sum(sp_s, dim=-1, keepdim=False)
+
+            if i == 0:
+                out = sp_s
+            else:
+                out = torch.cat((out, sp_s), 1)
+        out = torch.cat((out, spx[self.nums].squeeze(-1)), 1)
+        out = self.conv3(out)
+        out = self.relu(out)
+        out = self.bn3(out)
+        out = self.se(out)
+        out += residual
+        return out
+
+
+class ASTP(nn.Module):
+    """Attentive statistics pooling: Channel- and context-dependent
+    statistics pooling, first used in ECAPA_TDNN.
+    """
+
+    def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False):
+        super(ASTP, self).__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't
+        # need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(
+                in_dim * 3, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(
+                in_dim, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(
+            bottleneck_dim, in_dim, kernel_size=1
+        )  # equals V and k in the paper
+
+    def forward(self, x):
+        """
+        x: a 3-dimensional tensor in tdnn-based architecture (B,F,T)
+            or a 4-dimensional tensor in resnet architecture (B,C,F,T)
+            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
+        """
+        if len(x.shape) == 4:
+            x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])
+        assert len(x.shape) == 3
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(
+                torch.var(x, dim=-1, keepdim=True) + 1e-10
+            ).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! ReLU may be hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))  # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        var = torch.sum(alpha * (x**2), dim=2) - mean**2
+        std = torch.sqrt(var.clamp(min=1e-10))
+        return torch.cat([mean, std], dim=1)
+
+
+class Nested_Res2Net_TDNN(nn.Module):
+
+    def __init__(
+        self,
+        Nes_ratio=[8, 8],
+        input_channel=1024,
+        n_output_logits=2,
+        dilation=2,
+        pool_func="mean",
+        SE_ratio=[8],
+    ):
+
+        super(Nested_Res2Net_TDNN, self).__init__()
+        self.Nes_ratio = Nes_ratio[0]
+        assert input_channel % Nes_ratio[0] == 0
+        C = input_channel // Nes_ratio[0]
+        self.C = C
+        Build_in_Res2Nets = []
+        bns = []
+        for i in range(Nes_ratio[0] - 1):
+            Build_in_Res2Nets.append(
+                Bottle2neck(
+                    C,
+                    C,
+                    kernel_size=3,
+                    dilation=dilation,
+                    scale=Nes_ratio[1],
+                    SE_ratio=SE_ratio[0],
+                )
+            )
+            bns.append(nn.BatchNorm1d(C))
+        self.Build_in_Res2Nets = nn.ModuleList(Build_in_Res2Nets)
+        self.bns = nn.ModuleList(bns)
+        self.bn = nn.BatchNorm1d(1024)
+        self.relu = nn.ReLU()
+        self.pool_func = pool_func
+        if pool_func == "mean":
+            self.fc = nn.Linear(1024, n_output_logits)
+        elif pool_func == "ASTP":
+            self.pooling = ASTP(
+                in_dim=input_channel, bottleneck_dim=128, global_context_att=False
+            )
+            self.fc = nn.Linear(2048, n_output_logits)
+
+    def forward(self, x):
+        spx = torch.split(x, self.C, 1)
+        for i in range(self.Nes_ratio - 1):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.Build_in_Res2Nets[i](sp)
+            sp = self.relu(sp)
+            sp = self.bns[i](sp)
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        out = torch.cat((out, spx[-1]), 1)
+        out = self.bn(out)
+        out = self.relu(out)
+        if self.pool_func == "mean":
+            out = torch.mean(out, dim=-1)
+        elif self.pool_func == "ASTP":
+            out = self.pooling(out)
+        out = self.fc(out)
+        return out
+
+
+class wav2vec2_Nes2Net_no_Res_w_allT(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self.device = device
+
+        self.n_output_logits = args.n_output_logits
+
+        ####
+        # create network wav2vec 2.0
+        ####
+        self.ssl_model = SSLModel(self.device)
+        self.Nested_Res2Net_TDNN = Nested_Res2Net_TDNN(
+            Nes_ratio=args.Nes_ratio,
+            input_channel=1024,
+            n_output_logits=self.n_output_logits,
+            dilation=args.dilation,
+            pool_func=args.pool_func,
+            SE_ratio=args.SE_ratio,
+        )
+
+    def forward(self, x):
+        # -------pre-trained Wav2vec model fine tunning ------------------------##
+        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
+        x_ssl_feat = x_ssl_feat.permute(0, 2, 1)
+        output = self.Nested_Res2Net_TDNN(x_ssl_feat)
+
+        return output
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n_output_logits", type=int, default=2)
+    parser.add_argument("--dilation", type=int, default=2)  # not important
+    parser.add_argument(
+        "--pool_func",
+        type=str,
+        default="mean",
+        choices=["mean", "ASTP"],
+        help="pooling function, choose from mean and ASTP",
+    )
+    parser.add_argument(
+        "--Nes_ratio",
+        type=int,
+        nargs="+",
+        default=[8, 8],
+        help="Nes_ratio, from outer to inner",
+    )
+    parser.add_argument(
+        "--SE_ratio",
+        type=int,
+        nargs="+",
+        default=[1],
+        help="SE downsampling ratio in the bottleneck",
+    )
+    args = parser.parse_args()
+
+    model = wav2vec2_Nes2Net_no_Res_w_allT(args=args, device="cpu")
+    x = torch.rand((4, 32000)).to("cpu")
+    model = model.to("cpu")
+    y = model(x)
+    print(y)
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("all:", trainable_params)
+    trainable_params = sum(
+        p.numel() for p in model.ssl_model.parameters() if p.requires_grad
+    )
+    print("SSL:", trainable_params)
+    trainable_params = sum(
+        p.numel() for p in model.Nested_Res2Net_TDNN.parameters() if p.requires_grad
+    )
+    print("Backend:", trainable_params)

clean/audio/safeear/.gitignore

@@ -0,0 +1,167 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+model_zoos/*
+Exps/*
+datas/datasets
+datas/ASVSpoof2019/LA
+datas/ASVSpoof2021/ASVspoof2021_LA_eval
+datas/ASVSpoof2021/keys
+create_tsv.py

clean/audio/safeear/LICENSE

@@ -0,0 +1,23 @@
+# Creative Commons Attribution 4.0 International License
+
+## License
+
+You are free to:
+
+- Share — copy and redistribute the material in any medium or format
+- Adapt — remix, transform, and build upon the material for any purpose, even commercially.
+
+Under the following terms:
+
+1. **Attribution** — You must give appropriate credit, provide a link to the license, and indicate if changes were made. You may do so in any reasonable manner, but not in any way that suggests the licensor endorses you or your use.
+
+2. **No additional restrictions** — You may not apply legal terms or technological measures that legally restrict others from doing anything the license permits.
+
+## Other Terms
+
+- This license applies to all types of works, including but not limited to text, images, audio, video, etc.
+- This license does not apply to any third-party materials included in the work, for which you must obtain permission separately.
+
+## Disclaimer
+
+This work is provided on an "as is" basis, without any warranties or conditions of any kind, either express or implied, including but not limited to implied warranties of merchantability, fitness for a particular purpose, or non-infringement.

clean/audio/safeear/README.md

@@ -0,0 +1,134 @@
+# <font color=E7595C>Safe</font><font color=F6C446>Ear</font><img src="assert/SafeEar_logo.jpg" alt="icon" style="width: 2em; height: 1.5em; vertical-align: middle;">: <font color=E7595C>Content Privacy-Preserving</font> <font color=F6C446>Audio Deepfake Detection</font>
+
+[![arXiv](https://img.shields.io/badge/arXiv-2409.09272-b31b1b.svg)](https://arxiv.org/abs/2409.09272)
+[![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg?style=flat-square)](https://makeapullrequest.com) 
+[![CC BY 4.0](https://img.shields.io/badge/license-CC%20BY%204.0-blue.svg)](https://creativecommons.org/licenses/by/4.0/)
+![GitHub stars](https://img.shields.io/github/stars/LetterLiGo/SafeEar)
+![GitHub forks](https://img.shields.io/github/forks/LetterLiGo/SafeEar)
+![Website](https://img.shields.io/website?url=https://safeearweb.github.io/Project/)
+
+
+By [1] Zhejiang University, [2] Tsinghua University.
+* [Xinfeng Li](https://letterligo.github.io)* [1], [Kai Li](https://cslikai.cn)* [2], Yifan Zheng [1], Chen Yan† [1], Xiaoyu Ji [1], Wenyuan Xu [1].
+
+This repository is an official implementation of the SafeEar accepted to **ACM CCS 2024** (Core-A*, CCF-A, Big4) .
+
+Please also visit our <a href="https://safeearweb.github.io/Project/">(1) Project Website</a>, <a href="https://zenodo.org/records/14062964">(2) Full CVoiceFake Dataset</a>, and <a href="https://zenodo.org/records/11124319">(3) Sampled CVoiceFake Dataset</a>.
+
+## 🔥News
+
+[2025-03-18]: Supported the batch testing for ASVspoof 2019 and 2021, fixed some bugs for datasets and trainer.
+
+[2024-12-10]: Fixed all the bugs for training and test, and uploaded the files for data generation `datas/`.
+
+[2024-12-01]: Uploaded the checkpoint for data generation `datas/`.
+
+## ✨Key Highlights:
+
+In this paper, we propose SafeEar, a novel framework that aims to detect deepfake audios without relying on accessing the speech content within. Our key idea is to devise a neural audio codec into a novel decoupling model that well separates the semantic and acoustic information from audio samples, and only use the acoustic information (e.g., prosody and timbre) for deepfake detection. In this way, no semantic content will be exposed to the detector. To overcome the challenge of identifying diverse deepfake audio without semantic clues, we enhance our deepfake detector with multi-head self-attention and codec augmentation. Extensive experiments conducted on four benchmark datasets demonstrate SafeEar’s effectiveness in detecting various deepfake techniques with an equal error rate (EER) down to 2.02%. Simultaneously, it shields five-language speech content from being deciphered by both machine and human auditory analysis, demonstrated by word error rates (WERs) all above 93.93% and our user study. Furthermore, our benchmark constructed for anti-deepfake and anti-content recovery evaluation helps provide a basis for future research in the realms of audio privacy preservation and deepfake detection.
+
+## 🚀Overall Pipeline
+
+![pipeline](assert/overall.gif)
+
+## 🔧Installation
+
+1. Clone the repository:
+
+```shell
+git clone git@github.com:LetterLiGo/SafeEar.git
+cd SafeEar/
+```
+
+2. Create and activate the conda environment:
+
+```shell
+conda create -n safeear python=3.9 
+conda activate safeear
+```
+
+3. Install PyTorch and torchvision following the [official instructions](https://pytorch.org). The code requires `python=3.9`, `pytorch=1.13`, `torchvision=0.14`.
+
+
+```shell
+pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
+
+```
+4. Install other dependencies:
+
+```shell 
+pip install pip==24.0
+pip install -r requirements.txt
+```
+
+## 📊Model Performance
+### ASVspoof 2019 & 2021
+![](assert/ASVSpoof-results.png)
+### Speech Recognition Performance
+![](assert/exp1.png)
+
+## Data preparation
+
+### AVSpoof 2019 & 2021
+
+Please download the [ASVspoof 2019](https://datashare.is.ed.ac.uk/handle/10283/3336) and [ASVspoof 2021](https://www.asvspoof.org/index2021.html) datasets and extract them to the `datas/datasets` directory.
+
+```shell
+datas/datasets/ASVspoof2019
+datas/datasets/ASVspoof2021
+```
+
+#### Generate the Hubert L9 feature files
+
+```shell
+mkdir model_zoos
+cd model_zoos
+wget https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt
+wget https://cloud.tsinghua.edu.cn/f/413a0cd2e6f749eea956/?dl=1 -O SpeechTokenizer.pt
+cd ../datas
+# Generate the Hubert L9 feature files for ASVspoof 2019
+python dump_hubert_avg_feature.py datasets/ASVSpoof2019 datasets/ASVSpoof2019_Hubert_L9
+# Generate the Hubert L9 feature files for ASVspoof 2021
+python dump_hubert_avg_feature.py datasets/ASVSpoof2021 datasets/ASVSpoof2021_Hubert_L9
+```
+
+## 📚Training
+
+Before starting training, please modify the parameter configurations in [`configs`](configs).
+
+Use the following commands to start training:
+
+```shell
+python train.py --conf_dir config/train19.yaml
+python train.py --conf_dir config/train21.yaml
+```
+
+## 📈Testing/Inference
+
+To evaluate a model on one or more GPUs, specify the `CUDA_VISIBLE_DEVICES`, `dataset`, `model` and `checkpoint`:
+
+```shell
+python test.py --conf_dir Exps/ASVspoof19/config.yaml
+python test.py --conf_dir Exps/ASVspoof21/config.yaml
+```
+
+## Bugs and Issues
+
+If you meet `RuntimeError: Failed to load audio from <_io.BytesIO object at 0x7f45cb978f90>`, please use the following command to fix it:
+
+```shell
+conda install -c anaconda 'ffmpeg<4.4'
+```
+
+## 📜Citation
+
+If you find our work/code/dataset helpful, please consider citing:
+
+```
+@inproceedings{li2024safeear,
+  author       = {Li, Xinfeng and Li, Kai and Zheng, Yifan and Yan, Chen and Ji, Xiaoyu, and Xu, Wenyuan},
+  title        = {{SafeEar: Content Privacy-Preserving Audio Deepfake Detection}},
+  booktitle    = {Proceedings of the 2024 {ACM} {SIGSAC} Conference on Computer and Communications Security (CCS)}
+  year         = {2024},
+}
+```

clean/audio/safeear/SOURCE.md

@@ -0,0 +1,17 @@
+# Source: audio/safeear
+
+| Field | Value |
+|---|---|
+| Upstream | **UNVERIFIED** -- provenance was lost when this code was vendored |
+| Paper | https://arxiv.org/abs/2409.09272 |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/audio__safeear.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.

clean/audio/safeear/config/train19.yaml

@@ -0,0 +1,87 @@
+datamodule:
+  _target_: safeear.datas.asvspoof19.DataModule
+  batch_size: 2
+  num_workers: 8
+  pin_memory: true
+  DataClass_dict:
+    _target_: safeear.datas.asvspoof19.DataClass
+    train_path: ["datas/ASVSpoof2019/train.tsv", "datas/ASVSpoof2019/ASVspoof2019.LA.cm.train.trn.txt", "datas/datasets/ASVSpoof2019_Hubert_L9/ASVspoof2019_LA_train/flac"]
+    val_path: ["datas/ASVSpoof2019/dev.tsv", "datas/ASVSpoof2019/ASVspoof2019.LA.cm.dev.trl.txt", "datas/datasets/ASVSpoof2019_Hubert_L9/ASVspoof2019_LA_dev/flac"]
+    test_path: ["datas/ASVSpoof2019/eval.tsv", "datas/ASVSpoof2019/ASVspoof2019.LA.cm.eval.trl.txt", "datas/datasets/ASVSpoof2019_Hubert_L9/ASVspoof2019_LA_eval/flac"]
+    max_len: 64600
+
+decouple_model:
+  _target_: safeear.models.decouple.SpeechTokenizer
+  n_filters: 64
+  strides: [8,5,4,2]
+  dimension: 1024
+  semantic_dimension: 768
+  bidirectional: true
+  dilation_base: 2
+  residual_kernel_size: 3
+  n_residual_layers: 1
+  lstm_layers: 2
+  activation: ELU
+  codebook_size: 1024
+  n_q: 8
+  sample_rate: 16000
+
+speechtokenizer_path: model_zoos/SpeechTokenizer.pt
+
+detect_model:
+  _target_: safeear.models.safeear.SafeEar1s
+  front:
+    _target_: safeear.models.safeear.SE_Rawformer_front
+  embedding_dim: 1024 
+  dropout_rate: 0.1 
+  attention_dropout: 0.1 
+  stochastic_depth: 0.1 
+  num_layers: 2
+  num_heads: 8 
+  num_classes: 2 
+  positional_embedding: 'sine' 
+  mlp_ratio: 1.0
+
+system:
+  _target_: safeear.trainer.safeear_trainer.SafeEarTrainer
+  lr_raw_former: 3.0e-4
+  save_score_path: ${exp.dir}/${exp.name}
+
+exp: 
+  dir: Exps/ # 修改
+  name: ASVspoof19 # 修改
+
+early_stopping:
+  _target_: pytorch_lightning.callbacks.EarlyStopping
+  monitor: val_eer # 修改
+  mode: min
+  patience: 40
+  verbose: true
+
+checkpoint:
+  _target_: pytorch_lightning.callbacks.ModelCheckpoint
+  dirpath: ${exp.dir}/${exp.name}/checkpoints
+  monitor: val_eer # 修改
+  mode: min
+  verbose: true
+  save_top_k: 1
+  save_last: true
+  filename: '{epoch}-{val_eer:.4f}' # 修改
+
+logger:
+  _target_: pytorch_lightning.loggers.WandbLogger
+  name: ${exp.name}
+  save_dir: ${exp.dir}/${exp.name}/logs
+  offline: true
+  project: SafeEar
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  devices: [0]
+  max_epochs: 500
+  sync_batchnorm: true
+  default_root_dir: ${exp.dir}/${exp.name}/
+  accelerator: gpu
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  fast_dev_run: false

clean/audio/safeear/config/train21.yaml

@@ -0,0 +1,87 @@
+datamodule:
+  _target_: safeear.datas.asvspoof21.DataModule
+  batch_size: 2
+  num_workers: 8
+  pin_memory: true
+  DataClass_dict:
+    _target_: safeear.datas.asvspoof21.DataClass
+    train_path: ["datas/ASVSpoof2019/train.tsv", "datas/ASVSpoof2019/ASVspoof2019.LA.cm.train.trn.txt", "datas/datasets/ASVSpoof2019_Hubert_L9/ASVspoof2019_LA_train/flac"]
+    val_path: ["datas/ASVSpoof2019/dev.tsv", "datas/ASVSpoof2019/ASVspoof2019.LA.cm.dev.trl.txt", "datas/datasets/ASVSpoof2019_Hubert_L9/ASVspoof2019_LA_dev/flac"]
+    test_path: ["datas/ASVSpoof2021/eval.tsv", "datas/ASVSpoof2021/ASVspoof2021.LA.cm.eval.trl.txt", "datas/datasets/ASVSpoof2021_Hubert_L9"]
+    max_len: 64600
+
+decouple_model:
+  _target_: safeear.models.decouple.SpeechTokenizer
+  n_filters: 64
+  strides: [8,5,4,2]
+  dimension: 1024
+  semantic_dimension: 768
+  bidirectional: true
+  dilation_base: 2
+  residual_kernel_size: 3
+  n_residual_layers: 1
+  lstm_layers: 2
+  activation: ELU
+  codebook_size: 1024
+  n_q: 8
+  sample_rate: 16000
+
+speechtokenizer_path: model_zoos/SpeechTokenizer.pt
+
+detect_model:
+  _target_: safeear.models.safeear.SafeEar1s
+  front:
+    _target_: safeear.models.safeear.SE_Rawformer_front
+  embedding_dim: 1024 
+  dropout_rate: 0.1 
+  attention_dropout: 0.1 
+  stochastic_depth: 0.1 
+  num_layers: 2
+  num_heads: 8 
+  num_classes: 2 
+  positional_embedding: 'sine' 
+  mlp_ratio: 1.0
+
+system:
+  _target_: safeear.trainer.safeear_trainer.SafeEarTrainer
+  lr_raw_former: 3.0e-4
+  save_score_path: ${exp.dir}/${exp.name}
+
+exp: 
+  dir: Exps/ # 修改
+  name: ASVspoof21 # 修改
+
+early_stopping:
+  _target_: pytorch_lightning.callbacks.EarlyStopping
+  monitor: val_eer # 修改
+  mode: min
+  patience: 40
+  verbose: true
+
+checkpoint:
+  _target_: pytorch_lightning.callbacks.ModelCheckpoint
+  dirpath: ${exp.dir}/${exp.name}/checkpoints
+  monitor: val_eer # 修改
+  mode: min
+  verbose: true
+  save_top_k: 1
+  save_last: true
+  filename: '{epoch}-{val_eer:.4f}' # 修改
+
+logger:
+  _target_: pytorch_lightning.loggers.WandbLogger
+  name: ${exp.name}
+  save_dir: ${exp.dir}/${exp.name}/logs
+  offline: true
+  project: SafeEar
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  devices: [0]
+  max_epochs: 40
+  sync_batchnorm: true
+  default_root_dir: ${exp.dir}/${exp.name}/
+  accelerator: gpu
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  fast_dev_run: false

clean/audio/safeear/requirements.txt

@@ -0,0 +1,123 @@
+absl-py==2.1.0
+aiohttp==3.9.0
+aiosignal==1.3.1
+antlr4-python3-runtime==4.8
+appdirs==1.4.4
+asttokens==2.4.1
+async-timeout==4.0.3
+attrs==23.1.0
+audioread==3.0.1
+bitarray==2.8.3
+blessed==1.20.0
+certifi==2022.12.7
+cffi==1.16.0
+charset-normalizer==2.1.1
+click==8.1.7
+cmake==3.25.0
+colorama==0.4.6
+contourpy==1.2.0
+cycler==0.12.1
+Cython==3.0.5
+decorator==5.1.1
+docker-pycreds==0.4.0
+einops==0.7.0
+exceptiongroup==1.2.0
+executing==2.0.1
+# Editable install with no version control (fairseq==1.0.0a0)
+-e fairseq_ours
+fast-bss-eval==0.1.4
+filelock==3.9.0
+fonttools==4.45.0
+frozenlist==1.4.0
+fsspec==2023.10.0
+gitdb==4.0.11
+GitPython==3.1.40
+gpustat==1.1.1
+grpcio==1.63.0
+huggingface-hub==0.19.4
+hydra-core==1.0.7
+idna==3.4
+importlib-resources==6.1.1
+importlib_metadata==7.1.0
+ipdb==0.13.13
+ipython==8.18.1
+jedi==0.19.1
+Jinja2==3.1.2
+joblib==1.4.2
+kiwisolver==1.4.5
+lazy_loader==0.4
+librosa==0.10.2
+lightning-utilities==0.10.0
+lit==15.0.7
+llvmlite==0.42.0
+lxml==4.9.3
+Markdown==3.6
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.8.2
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.0.8
+multidict==6.0.4
+networkx==3.0
+numba==0.59.1
+numpy==1.23.5
+nvidia-ml-py==12.535.133
+opencv-python==4.9.0.80
+packaging==23.2
+parso==0.8.3
+pexpect==4.9.0
+Pillow==9.3.0
+platformdirs==4.2.1
+pooch==1.8.1
+portalocker==2.8.2
+prompt-toolkit==3.0.43
+protobuf==4.25.1
+psutil==5.9.6
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pycparser==2.21
+pyDeprecate==0.3.2
+Pygments==2.17.2
+pyparsing==3.1.1
+python-dateutil==2.8.2
+pytorch-lightning==1.6.3
+pytorch-ranger==0.1.1
+PyYAML==6.0.1
+regex==2023.10.3
+requests==2.28.1
+rich==13.7.0
+sacrebleu==2.3.2
+safetensors==0.4.0
+scikit-learn==1.4.2
+scipy==1.11.4
+sentry-sdk==1.36.0
+setproctitle==1.3.3
+six==1.16.0
+smmap==5.0.1
+soundfile>=0.11.0
+soxr==0.3.7
+stack-data==0.6.3
+sympy==1.12
+tabulate==0.9.0
+tensorboard==2.16.2
+tensorboard-data-server==0.7.2
+thop==0.1.1.post2209072238
+threadpoolctl==3.5.0
+timm==0.9.11
+tomli==2.0.1
+torch-mir-eval==0.4
+torch-optimizer==0.3.0
+torchmetrics==1.2.0
+tqdm==4.66.1
+traitlets==5.14.1
+triton==2.0.0
+typing_extensions==4.4.0
+urllib3==1.26.13
+wandb==0.16.0
+wcwidth==0.2.12
+Werkzeug==3.0.2
+yarl==1.9.3
+zipp==3.17.0
+npy_append_array==0.9.16

clean/audio/safeear/safeear/losses/loss.py

@@ -0,0 +1,215 @@
+import torch
+import torch.nn.functional as F
+from torchaudio.transforms import MelSpectrogram
+import numpy as np
+
+def adversarial_g_loss(y_disc_gen):
+    """Hinge loss"""
+    loss = 0.0
+    for i in range(len(y_disc_gen)):
+        stft_loss = F.relu(1 - y_disc_gen[i]).mean().squeeze()
+        loss += stft_loss
+    return loss / len(y_disc_gen)
+
+
+def feature_loss(fmap_r, fmap_gen):
+    loss = 0.0
+    for i in range(len(fmap_r)):
+        for j in range(len(fmap_r[i])):
+            stft_loss = ((fmap_r[i][j] - fmap_gen[i][j]).abs() /
+                         (fmap_r[i][j].abs().mean())).mean()
+            loss += stft_loss
+    return loss / (len(fmap_r) * len(fmap_r[0]))
+
+
+def sim_loss(y_disc_r, y_disc_gen):
+    loss = 0.0
+    for i in range(len(y_disc_r)):
+        loss += F.mse_loss(y_disc_r[i], y_disc_gen[i])
+    return loss / len(y_disc_r)
+
+def reconstruction_loss(x, G_x, lamdba_wav=100, sr=16000, eps=1e-7):
+    # NOTE (lsx): hard-coded now
+    L = lamdba_wav * F.mse_loss(x, G_x)  # wav L1 loss
+    # loss_sisnr = sisnr_loss(G_x, x) # 
+    # L += 0.01*loss_sisnr
+    # 2^6=64 -> 2^10=1024
+    # NOTE (lsx): add 2^11
+    for i in range(6, 12):
+        # for i in range(5, 12): # Encodec setting
+        s = 2**i
+        melspec = MelSpectrogram(
+            sample_rate=sr,
+            n_fft=s,
+            hop_length=s // 4,
+            n_mels=64,
+            wkwargs={"device": x.device}).to(x.device)
+        S_x = melspec(x)
+        S_G_x = melspec(G_x)
+        loss = ((S_x - S_G_x).abs().mean() + (
+            ((torch.log(S_x.abs() + eps) - torch.log(S_G_x.abs() + eps))**2
+             ).mean(dim=-2)**0.5).mean()) / i
+        L += loss
+    return L
+
+
+def criterion_d(y_disc_r, y_disc_gen, fmap_r_det, fmap_gen_det, y_df_hat_r,
+                y_df_hat_g, fmap_f_r, fmap_f_g, y_ds_hat_r, y_ds_hat_g,
+                fmap_s_r, fmap_s_g):
+    """Hinge Loss"""
+    loss = 0.0
+    loss1 = 0.0
+    loss2 = 0.0
+    loss3 = 0.0
+    for i in range(len(y_disc_r)):
+        loss1 += F.relu(1 - y_disc_r[i]).mean() + F.relu(1 + y_disc_gen[
+            i]).mean()
+    for i in range(len(y_df_hat_r)):
+        loss2 += F.relu(1 - y_df_hat_r[i]).mean() + F.relu(1 + y_df_hat_g[
+            i]).mean()
+    for i in range(len(y_ds_hat_r)):
+        loss3 += F.relu(1 - y_ds_hat_r[i]).mean() + F.relu(1 + y_ds_hat_g[
+            i]).mean()
+
+    loss = (loss1 / len(y_disc_gen) + loss2 / len(y_df_hat_r) + loss3 /
+            len(y_ds_hat_r)) / 3.0
+
+    return loss
+
+
+def criterion_g(commit_loss, x, G_x, fmap_r, fmap_gen, y_disc_r, y_disc_gen,
+                y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g, y_ds_hat_r,
+                y_ds_hat_g, fmap_s_r, fmap_s_g, lamdba_wav=100, lamdba_com=1000, lamdba_adv=1, lamdba_feat=1, lamdba_rec=1, sr=16000):
+    adv_g_loss = adversarial_g_loss(y_disc_gen)
+    feat_loss = (feature_loss(fmap_r, fmap_gen) + sim_loss(
+        y_disc_r, y_disc_gen) + feature_loss(fmap_f_r, fmap_f_g) + sim_loss(
+            y_df_hat_r, y_df_hat_g) + feature_loss(fmap_s_r, fmap_s_g) +
+                 sim_loss(y_ds_hat_r, y_ds_hat_g)) / 3.0
+    rec_loss = reconstruction_loss(x.contiguous(), G_x.contiguous(), lamdba_wav, sr)
+    total_loss = lamdba_com * commit_loss + lamdba_adv * adv_g_loss + lamdba_feat * feat_loss + lamdba_rec * rec_loss
+    return total_loss, adv_g_loss, feat_loss, rec_loss
+
+
+def adopt_weight(weight, global_step, threshold=0, value=0.):
+    if global_step < threshold:
+        weight = value
+    return weight
+
+
+def adopt_dis_weight(weight, global_step, threshold=0, value=0.):
+    # 0,3,6,9,13....这些时间步，不更新dis
+    if global_step % 3 == 0:
+        weight = value
+    return weight
+
+
+def calculate_adaptive_weight(nll_loss, g_loss, last_layer, lamdba_adv=1):
+    if last_layer is not None:
+        nll_grads = torch.autograd.grad(
+            nll_loss, last_layer, retain_graph=True)[0]
+        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+    else:
+        print('last_layer cannot be none')
+        assert 1 == 2
+    d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+    d_weight = torch.clamp(d_weight, 1.0, 1.0).detach()
+    d_weight = d_weight * lamdba_adv
+    return d_weight
+
+def loss_g(codebook_loss,
+           inputs,
+           reconstructions,
+           fmap_r,
+           fmap_gen,
+           y_disc_r,
+           y_disc_gen,
+           global_step,
+           y_df_hat_r,
+           y_df_hat_g,
+           y_ds_hat_r,
+           y_ds_hat_g,
+           fmap_f_r,
+           fmap_f_g,
+           fmap_s_r,
+           fmap_s_g,
+           lamdba_wav=100, 
+           lamdba_com=1000, 
+           lamdba_adv=1, 
+           lamdba_feat=1, 
+           sr=16000,
+           discriminator_iter_start=500
+           ):
+    """
+    args:
+        codebook_loss: commit loss.
+        inputs: ground-truth wav.
+        reconstructions: reconstructed wav.
+        fmap_r: real stft-D feature map.
+        fmap_gen: fake stft-D feature map.
+        y_disc_r: real stft-D logits.
+        y_disc_gen: fake stft-D logits.
+        global_step: global training step.
+        y_df_hat_r: real MPD logits.
+        y_df_hat_g: fake MPD logits.
+        y_ds_hat_r: real MSD logits.
+        y_ds_hat_g: fake MSD logits.
+        fmap_f_r: real MPD feature map.
+        fmap_f_g: fake MPD feature map.
+        fmap_s_r: real MSD feature map.
+        fmap_s_g: fake MSD feature map.
+    """
+    rec_loss = reconstruction_loss(inputs.contiguous(),
+                                   reconstructions.contiguous(), lamdba_wav, sr)
+    adv_g_loss = adversarial_g_loss(y_disc_gen)
+    adv_mpd_loss = adversarial_g_loss(y_df_hat_g)
+    adv_msd_loss = adversarial_g_loss(y_ds_hat_g)
+    adv_loss = (adv_g_loss + adv_mpd_loss + adv_msd_loss
+                ) / 3.0  # NOTE(lsx): need to divide by 3?
+    feat_loss = feature_loss(
+        fmap_r,
+        fmap_gen)  #+ sim_loss(y_disc_r, y_disc_gen) # NOTE(lsx): need logits?
+    feat_loss_mpd = feature_loss(fmap_f_r,
+                                 fmap_f_g)  #+ sim_loss(y_df_hat_r, y_df_hat_g)
+    feat_loss_msd = feature_loss(fmap_s_r,
+                                 fmap_s_g)  #+ sim_loss(y_ds_hat_r, y_ds_hat_g)
+    feat_loss_tot = (feat_loss + feat_loss_mpd + feat_loss_msd) / 3.0
+    d_weight = torch.tensor(1.0)
+    disc_factor = adopt_weight(
+        lamdba_adv, global_step, threshold=discriminator_iter_start)
+    if disc_factor == 0.:
+        fm_loss_wt = 0
+    else:
+        fm_loss_wt = lamdba_feat
+    loss = rec_loss + d_weight * disc_factor * adv_loss + \
+           fm_loss_wt * feat_loss_tot + lamdba_com * codebook_loss
+    return loss, rec_loss, adv_loss, feat_loss_tot, d_weight
+
+def compute_det_curve(target_scores, nontarget_scores):
+
+    n_scores = target_scores.size + nontarget_scores.size
+    all_scores = np.concatenate((target_scores, nontarget_scores))
+    labels = np.concatenate((np.ones(target_scores.size), np.zeros(nontarget_scores.size)))
+
+    # Sort labels based on scores
+    indices = np.argsort(all_scores, kind='mergesort')
+    labels = labels[indices]
+
+    # Compute false rejection and false acceptance rates
+    tar_trial_sums = np.cumsum(labels)
+    nontarget_trial_sums = nontarget_scores.size - (np.arange(1, n_scores + 1) - tar_trial_sums)
+
+    frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size))  # false rejection rates
+    far = np.concatenate((np.atleast_1d(1), nontarget_trial_sums / nontarget_scores.size))  # false acceptance rates
+    thresholds = np.concatenate((np.atleast_1d(all_scores[indices[0]] - 0.001), all_scores[indices]))  # Thresholds are the sorted scores
+
+    return frr, far, thresholds
+
+
+def compute_eer(target_scores, nontarget_scores):
+    """ Returns equal error rate (EER) and the corresponding threshold. """
+    frr, far, thresholds = compute_det_curve(target_scores, nontarget_scores)
+    abs_diffs = np.abs(frr - far)
+    min_index = np.argmin(abs_diffs)
+    eer = np.mean((frr[min_index], far[min_index]))
+    print(thresholds[min_index])
+    return eer, thresholds[min_index]

clean/audio/safeear/safeear/models/decouple.py

@@ -0,0 +1,207 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Aug 30 15:47:55 2023
+@author: zhangxin
+"""
+import torch.nn as nn
+from einops import rearrange
+import torch
+
+from .modules.seanet import SEANetEncoder, SEANetDecoder
+from .modules.quantization  import ResidualVectorQuantizer
+
+
+class SpeechTokenizer(nn.Module):
+    def __init__(self, n_filters, dimension, strides, lstm_layers, bidirectional, dilation_base, residual_kernel_size, n_residual_layers, activation, sample_rate, n_q, semantic_dimension, codebook_size):
+        '''
+        
+        Parameters
+        ----------
+        n_filters : int
+            Number of filters in the SEANet encoder/decoder.
+        dimension : int
+            Dimensionality of the encoder/decoder.
+        strides : list
+            List of stride values for the SEANet encoder/decoder.
+        lstm_layers : int
+            Number of LSTM layers in the encoder/decoder.
+        bidirectional : bool
+            Whether to use bidirectional LSTM in the encoder.
+        dilation_base : int
+            Base dilation rate for the residual blocks in the encoder/decoder.
+        residual_kernel_size : int
+            Kernel size for the residual blocks in the encoder/decoder.
+        n_residual_layers : int
+            Number of residual layers in the encoder/decoder.
+        activation : str
+            Activation function to use in the encoder/decoder.
+        sample_rate : int
+            Sample rate of the audio.
+        n_q : int
+            Number of quantization levels.
+        semantic_dimension : int
+            Dimensionality of the semantic representation.
+        codebook_size : int
+            Size of the codebook for vector quantization.
+
+        '''
+        super().__init__()
+        self.encoder = SEANetEncoder(n_filters=n_filters, 
+                                     dimension=dimension, 
+                                     ratios=strides,
+                                     lstm=lstm_layers,
+                                     bidirectional=bidirectional,
+                                     dilation_base=dilation_base,
+                                     residual_kernel_size=residual_kernel_size,
+                                     n_residual_layers=n_residual_layers,
+                                     activation=activation)
+        self.sample_rate = sample_rate
+        self.n_q = n_q
+        if dimension != semantic_dimension:
+            self.transform = nn.Linear(dimension, semantic_dimension)
+        else:
+            self.transform = nn.Identity()
+        self.quantizer = ResidualVectorQuantizer(dimension=dimension, n_q=n_q, bins=codebook_size)
+        self.decoder = SEANetDecoder(n_filters=n_filters, 
+                                     dimension=dimension, 
+                                     ratios=strides,
+                                     lstm=lstm_layers,
+                                     bidirectional=False,
+                                     dilation_base=dilation_base,
+                                     residual_kernel_size=residual_kernel_size,
+                                     n_residual_layers=n_residual_layers,
+                                     activation=activation)
+        
+    @classmethod
+    def load_from_checkpoint(cls, 
+                             config_path: str, 
+                             ckpt_path: str):
+        '''
+
+        Parameters
+        ----------
+        config_path : str
+            Path of model configuration file.
+        ckpt_path : str
+            Path of model  checkpoint.
+
+        Returns
+        -------
+        model : SpeechTokenizer
+            SpeechTokenizer model.
+
+        '''
+        import json
+        with open(config_path) as f:
+            cfg = json.load(f)
+        model = cls(cfg)
+        params = torch.load(ckpt_path, map_location='cpu')
+        model.load_state_dict(params)
+        return model
+    
+    
+    def forward(self, 
+                x: torch.tensor, 
+                n_q: int=None, 
+                layers: list=[0]):
+        '''
+        
+        Parameters
+        ----------
+        x : torch.tensor
+            Input wavs. Shape: (batch, channels, timesteps).
+        n_q : int, optional
+            Number of quantizers in RVQ used to encode. The default is all layers.
+        layers : list[int], optional
+            Layers of RVQ should return quantized result. The default is the first layer.
+
+        Returns
+        -------
+        o : torch.tensor
+            Output wavs. Shape: (batch, channels, timesteps).
+        commit_loss : torch.tensor
+            Commitment loss from residual vector quantizers.
+        feature : torch.tensor
+            Output of RVQ's first layer. Shape: (batch, timesteps, dimension)
+
+        '''
+        n_q = n_q if n_q else self.n_q
+        e = self.encoder(x)
+        quantized, codes, commit_loss, quantized_list = self.quantizer(e, n_q=n_q, layers=layers)
+        feature = rearrange(quantized_list[0], 'b d t -> b t d') # b,t,1024
+        feature = self.transform(feature) #b,t,768      
+        o = self.decoder(quantized)
+        return o, commit_loss, feature, quantized_list[1:]
+    
+    def forward_feature(self, 
+                        x: torch.tensor, 
+                        layers: list=None):
+        '''
+
+        Parameters
+        ----------
+        x : torch.tensor
+            Input wavs. Shape should be (batch, channels, timesteps).
+        layers : list[int], optional
+            Layers of RVQ should return quantized result. The default is all layers.
+
+        Returns
+        -------
+        quantized_list : list[torch.tensor]
+            Quantized of required layers.
+
+        '''
+        e = self.encoder(x)
+        layers = layers if layers else list(range(self.n_q))
+        quantized, codes, commit_loss, quantized_list = self.quantizer(e, layers=layers)
+        return quantized_list
+    
+    def encode(self, 
+               x: torch.tensor, 
+               n_q: int=None, 
+               st: int=None):
+        '''
+
+        Parameters
+        ----------
+        x : torch.tensor
+            Input wavs. Shape: (batch, channels, timesteps).
+        n_q : int, optional
+            Number of quantizers in RVQ used to encode. The default is all layers.
+        st : int, optional
+            Start quantizer index in RVQ. The default is 0.
+
+        Returns
+        -------
+        codes : torch.tensor
+            Output indices for each quantizer. Shape: (n_q, batch, timesteps)
+
+        '''
+        e = self.encoder(x)
+        if st is None:
+            st = 0
+        n_q = n_q if n_q else self.n_q
+        codes = self.quantizer.encode(e, n_q=n_q, st=st)
+        return codes
+    
+    def decode(self, 
+               codes: torch.tensor, 
+               st: int=0):
+        '''
+
+        Parameters
+        ----------
+        codes : torch.tensor
+            Indices for each quantizer. Shape: (n_q, batch, timesteps).
+        st : int, optional
+            Start quantizer index in RVQ. The default is 0.
+
+        Returns
+        -------
+        o : torch.tensor
+            Reconstruct wavs from codes. Shape: (batch, channels, timesteps)
+
+        '''
+        quantized = self.quantizer.decode(codes, st=st)
+        o = self.decoder(quantized)
+        return o

clean/audio/safeear/safeear/models/discriminator.py

@@ -0,0 +1,422 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""MS-STFT discriminator, provided here for reference."""
+import typing as tp
+
+import torch
+import torchaudio
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+import einops
+from torch.nn import AvgPool1d
+from torch.nn.utils import spectral_norm
+from torch.nn.utils import weight_norm
+
+FeatureMapType = tp.List[torch.Tensor]
+LogitsType = torch.Tensor
+DiscriminatorOutput = tp.Tuple[tp.List[LogitsType], tp.List[FeatureMapType]]
+
+CONV_NORMALIZATIONS = frozenset([
+    'none', 'weight_norm', 'spectral_norm', 'time_layer_norm', 'layer_norm',
+    'time_group_norm'
+])
+
+class ConvLayerNorm(nn.LayerNorm):
+    """
+    Convolution-friendly LayerNorm that moves channels to last dimensions
+    before running the normalization and moves them back to original position right after.
+    """
+
+    def __init__(self,
+                 normalized_shape: tp.Union[int, tp.List[int], torch.Size],
+                 **kwargs):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = einops.rearrange(x, 'b ... t -> b t ...')
+        x = super().forward(x)
+        x = einops.rearrange(x, 'b t ... -> b ... t')
+        return
+
+
+def apply_parametrization_norm(module: nn.Module,
+                               norm: str='none') -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return weight_norm(module)
+    elif norm == 'spectral_norm':
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module,
+                    causal: bool=False,
+                    norm: str='none',
+                    **norm_kwargs) -> nn.Module:
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'layer_norm':
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+    
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(self,
+                 *args,
+                 causal: bool=False,
+                 norm: str='none',
+                 norm_kwargs: tp.Dict[str, tp.Any]={},
+                 **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(self,
+                 *args,
+                 norm: str='none',
+                 norm_kwargs: tp.Dict[str, tp.Any]={},
+                 **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(
+            self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+def get_2d_padding(kernel_size: tp.Tuple[int, int],
+                   dilation: tp.Tuple[int, int]=(1, 1)):
+    return (((kernel_size[0] - 1) * dilation[0]) // 2, (
+        (kernel_size[1] - 1) * dilation[1]) // 2)
+
+
+class DiscriminatorSTFT(nn.Module):
+    """STFT sub-discriminator.
+    Args:
+        filters (int): Number of filters in convolutions
+        in_channels (int): Number of input channels. Default: 1
+        out_channels (int): Number of output channels. Default: 1
+        n_fft (int): Size of FFT for each scale. Default: 1024
+        hop_length (int): Length of hop between STFT windows for each scale. Default: 256
+        kernel_size (tuple of int): Inner Conv2d kernel sizes. Default: ``(3, 9)``
+        stride (tuple of int): Inner Conv2d strides. Default: ``(1, 2)``
+        dilations (list of int): Inner Conv2d dilation on the time dimension. Default: ``[1, 2, 4]``
+        win_length (int): Window size for each scale. Default: 1024
+        normalized (bool): Whether to normalize by magnitude after stft. Default: True
+        norm (str): Normalization method. Default: `'weight_norm'`
+        activation (str): Activation function. Default: `'LeakyReLU'`
+        activation_params (dict): Parameters to provide to the activation function.
+        growth (int): Growth factor for the filters. Default: 1
+    """
+
+    def __init__(self,
+                 filters: int,
+                 in_channels: int=1,
+                 out_channels: int=1,
+                 n_fft: int=1024,
+                 hop_length: int=256,
+                 win_length: int=1024,
+                 max_filters: int=1024,
+                 filters_scale: int=1,
+                 kernel_size: tp.Tuple[int, int]=(3, 9),
+                 dilations: tp.List=[1, 2, 4],
+                 stride: tp.Tuple[int, int]=(1, 2),
+                 normalized: bool=True,
+                 norm: str='weight_norm',
+                 activation: str='LeakyReLU',
+                 activation_params: dict={'negative_slope': 0.2}):
+        super().__init__()
+        assert len(kernel_size) == 2
+        assert len(stride) == 2
+        self.filters = filters
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.normalized = normalized
+        self.activation = getattr(torch.nn, activation)(**activation_params)
+        self.spec_transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window_fn=torch.hann_window,
+            normalized=self.normalized,
+            center=False,
+            pad_mode=None,
+            power=None)
+        spec_channels = 2 * self.in_channels
+        self.convs = nn.ModuleList()
+        self.convs.append(
+            NormConv2d(
+                spec_channels,
+                self.filters,
+                kernel_size=kernel_size,
+                padding=get_2d_padding(kernel_size)))
+        in_chs = min(filters_scale * self.filters, max_filters)
+        for i, dilation in enumerate(dilations):
+            out_chs = min((filters_scale**(i + 1)) * self.filters, max_filters)
+            self.convs.append(
+                NormConv2d(
+                    in_chs,
+                    out_chs,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    dilation=(dilation, 1),
+                    padding=get_2d_padding(kernel_size, (dilation, 1)),
+                    norm=norm))
+            in_chs = out_chs
+        out_chs = min((filters_scale**(len(dilations) + 1)) * self.filters,
+                      max_filters)
+        self.convs.append(
+            NormConv2d(
+                in_chs,
+                out_chs,
+                kernel_size=(kernel_size[0], kernel_size[0]),
+                padding=get_2d_padding((kernel_size[0], kernel_size[0])),
+                norm=norm))
+        self.conv_post = NormConv2d(
+            out_chs,
+            self.out_channels,
+            kernel_size=(kernel_size[0], kernel_size[0]),
+            padding=get_2d_padding((kernel_size[0], kernel_size[0])),
+            norm=norm)
+
+    def forward(self, x: torch.Tensor):
+        fmap = []
+        # print('x ', x.shape)
+        z = self.spec_transform(x)  # [B, 2, Freq, Frames, 2]
+        # print('z ', z.shape)
+        z = torch.cat([z.real, z.imag], dim=1)
+        # print('cat_z ', z.shape)
+        z = rearrange(z, 'b c w t -> b c t w')
+        for i, layer in enumerate(self.convs):
+            z = layer(z)
+            z = self.activation(z)
+            # print('z i', i, z.shape)
+            fmap.append(z)
+        z = self.conv_post(z)
+        # print('logit ', z.shape)
+        return z, fmap
+
+
+class MultiScaleSTFTDiscriminator(nn.Module):
+    """Multi-Scale STFT (MS-STFT) discriminator.
+    Args:
+        filters (int): Number of filters in convolutions
+        in_channels (int): Number of input channels. Default: 1
+        out_channels (int): Number of output channels. Default: 1
+        n_ffts (Sequence[int]): Size of FFT for each scale
+        hop_lengths (Sequence[int]): Length of hop between STFT windows for each scale
+        win_lengths (Sequence[int]): Window size for each scale
+        **kwargs: additional args for STFTDiscriminator
+    """
+
+    def __init__(self,
+                 filters: int,
+                 in_channels: int=1,
+                 out_channels: int=1,
+                 n_ffts: tp.List[int]=[1024, 2048, 512, 256, 128],
+                 hop_lengths: tp.List[int]=[256, 512, 128, 64, 32],
+                 win_lengths: tp.List[int]=[1024, 2048, 512, 256, 128],
+                 **kwargs):
+        super().__init__()
+        assert len(n_ffts) == len(hop_lengths) == len(win_lengths)
+        self.discriminators = nn.ModuleList([
+            DiscriminatorSTFT(
+                filters,
+                in_channels=in_channels,
+                out_channels=out_channels,
+                n_fft=n_ffts[i],
+                win_length=win_lengths[i],
+                hop_length=hop_lengths[i],
+                **kwargs) for i in range(len(n_ffts))
+        ])
+        self.num_discriminators = len(self.discriminators)
+
+    def forward(self, x: torch.Tensor) -> DiscriminatorOutput:
+        logits = []
+        fmaps = []
+        for disc in self.discriminators:
+            logit, fmap = disc(x)
+            logits.append(logit)
+            fmaps.append(fmap)
+        return logits, fmaps
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self,
+                 period,
+                 kernel_size=5,
+                 stride=3,
+                 use_spectral_norm=False,
+                 activation: str='LeakyReLU',
+                 activation_params: dict={'negative_slope': 0.2}):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
+        self.activation = getattr(torch.nn, activation)(**activation_params)
+        self.convs = nn.ModuleList([
+            NormConv2d(
+                1,
+                32, (kernel_size, 1), (stride, 1),
+                padding=(get_padding(5, 1), 0)),
+            NormConv2d(
+                32,
+                32, (kernel_size, 1), (stride, 1),
+                padding=(get_padding(5, 1), 0)),
+            NormConv2d(
+                32,
+                32, (kernel_size, 1), (stride, 1),
+                padding=(get_padding(5, 1), 0)),
+            NormConv2d(
+                32,
+                32, (kernel_size, 1), (stride, 1),
+                padding=(get_padding(5, 1), 0)),
+            NormConv2d(32, 32, (kernel_size, 1), 1, padding=(2, 0)),
+        ])
+        self.conv_post = NormConv2d(32, 1, (3, 1), 1, padding=(1, 0))
+
+    def forward(self, x):
+        fmap = []
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = self.activation(x)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super(MultiPeriodDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            DiscriminatorP(2),
+            DiscriminatorP(3),
+            DiscriminatorP(5),
+            DiscriminatorP(7),
+            DiscriminatorP(11),
+        ])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+    
+
+class DiscriminatorS(torch.nn.Module):
+    def __init__(self,
+                 use_spectral_norm=False,
+                 activation: str='LeakyReLU',
+                 activation_params: dict={'negative_slope': 0.2}):
+        super(DiscriminatorS, self).__init__()
+        self.activation = getattr(torch.nn, activation)(**activation_params)
+        self.convs = nn.ModuleList([
+            NormConv1d(1, 32, 15, 1, padding=7),
+            NormConv1d(32, 32, 41, 2, groups=4, padding=20),
+            NormConv1d(32, 32, 41, 2, groups=16, padding=20),
+            NormConv1d(32, 32, 41, 4, groups=16, padding=20),
+            NormConv1d(32, 32, 41, 4, groups=16, padding=20),
+            NormConv1d(32, 32, 41, 1, groups=16, padding=20),
+            NormConv1d(32, 32, 5, 1, padding=2),
+        ])
+        self.conv_post = NormConv1d(32, 1, 3, 1, padding=1)
+
+    def forward(self, x):
+        fmap = []
+        for l in self.convs:
+            x = l(x)
+            x = self.activation(x)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+        return x, fmap
+
+
+class MultiScaleDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super(MultiScaleDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            DiscriminatorS(),
+            DiscriminatorS(),
+            DiscriminatorS(),
+        ])
+        self.meanpools = nn.ModuleList(
+            [AvgPool1d(4, 2, padding=2), AvgPool1d(4, 2, padding=2)])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            if i != 0:
+                y = self.meanpools[i - 1](y)
+                y_hat = self.meanpools[i - 1](y_hat)
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs

clean/audio/safeear/safeear/models/modules/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Torch modules."""
+
+# flake8: noqa
+from .conv import (
+    pad1d,
+    unpad1d,
+    NormConv1d,
+    NormConvTranspose1d,
+    NormConv2d,
+    NormConvTranspose2d,
+    SConv1d,
+    SConvTranspose1d,
+)
+from .lstm import SLSTM
+from .seanet import SEANetEncoder, SEANetDecoder

clean/audio/safeear/safeear/models/modules/conv.py

@@ -0,0 +1,252 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Convolutional layers wrappers and utilities."""
+
+import math
+import typing as tp
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+from .norm import ConvLayerNorm
+
+
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+                                 'time_layer_norm', 'layer_norm', 'time_group_norm'])
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return weight_norm(module)
+    elif norm == 'spectral_norm':
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'layer_norm':
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+                                 padding_total: int = 0) -> int:
+    """See `pad_for_conv1d`.
+    """
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen.
+    """
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose2d(nn.Module):
+    """Wrapper around ConvTranspose2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class SConv1d(nn.Module):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, dilation: int = 1,
+                 groups: int = 1, bias: bool = True, causal: bool = False,
+                 norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 pad_mode: str = 'reflect'):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1'
+                          f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).')
+        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+                               dilation=dilation, groups=groups, bias=bias, causal=causal,
+                               norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        padding_total = (kernel_size - 1) * dilation - (stride - 1)
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        return self.conv(x)
+
+
+class SConvTranspose1d(nn.Module):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, causal: bool = False,
+                 norm: str = 'none', trim_right_ratio: float = 1.,
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+                                          causal=causal, norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        return y

clean/audio/safeear/safeear/models/modules/lstm.py

@@ -0,0 +1,32 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""LSTM layers module."""
+
+from torch import nn
+
+
+class SLSTM(nn.Module):
+    """
+    LSTM without worrying about the hidden state, nor the layout of the data.
+    Expects input as convolutional layout.
+    """
+    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True, bidirectional: bool=False):
+        super().__init__()
+        self.bidirectional = bidirectional
+        self.skip = skip
+        self.lstm = nn.LSTM(dimension, dimension, num_layers, bidirectional=bidirectional)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.bidirectional:
+            x = x.repeat(1, 1, 2)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y
+

clean/audio/safeear/safeear/models/modules/norm.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Normalization modules."""
+
+import typing as tp
+
+import einops
+import torch
+from torch import nn
+
+
+class ConvLayerNorm(nn.LayerNorm):
+    """
+    Convolution-friendly LayerNorm that moves channels to last dimensions
+    before running the normalization and moves them back to original position right after.
+    """
+    def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = einops.rearrange(x, 'b ... t -> b t ...')
+        x = super().forward(x)
+        x = einops.rearrange(x, 'b t ... -> b ... t')
+        return

clean/audio/safeear/safeear/models/modules/quantization/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# flake8: noqa
+from .vq import QuantizedResult, ResidualVectorQuantizer

clean/audio/safeear/safeear/models/modules/quantization/ac.py

@@ -0,0 +1,292 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Arithmetic coder."""
+
+import io
+import math
+import random
+import typing as tp
+import torch
+
+from ..binary import BitPacker, BitUnpacker
+
+
+def build_stable_quantized_cdf(pdf: torch.Tensor, total_range_bits: int,
+                               roundoff: float = 1e-8, min_range: int = 2,
+                               check: bool = True) -> torch.Tensor:
+    """Turn the given PDF into a quantized CDF that splits
+    [0, 2 ** self.total_range_bits - 1] into chunks of size roughly proportional
+    to the PDF.
+
+    Args:
+        pdf (torch.Tensor): probability distribution, shape should be `[N]`.
+        total_range_bits (int): see `ArithmeticCoder`, the typical range we expect
+            during the coding process is `[0, 2 ** total_range_bits - 1]`.
+        roundoff (float): will round the pdf up to that level to remove difference coming
+        from e.g. evaluating the Language Model on different architectures.
+        min_range (int): minimum range width. Should always be at least 2 for numerical
+            stability. Use this to avoid pathological behavior is a value
+            that is expected to be rare actually happens in real life.
+        check (bool): if True, checks that nothing bad happened, can be deactivated for speed.
+    """
+    pdf = pdf.detach()
+    if roundoff:
+        pdf = (pdf / roundoff).floor() * roundoff
+    # interpolate with uniform distribution to achieve desired minimum probability.
+    total_range = 2 ** total_range_bits
+    cardinality = len(pdf)
+    alpha = min_range * cardinality / total_range
+    assert alpha <= 1, "you must reduce min_range"
+    ranges = (((1 - alpha) * total_range) * pdf).floor().long()
+    ranges += min_range
+    quantized_cdf = torch.cumsum(ranges, dim=-1)
+    if min_range < 2:
+        raise ValueError("min_range must be at least 2.")
+    if check:
+        assert quantized_cdf[-1] <= 2 ** total_range_bits, quantized_cdf[-1]
+        if ((quantized_cdf[1:] - quantized_cdf[:-1]) < min_range).any() or quantized_cdf[0] < min_range:
+            raise ValueError("You must increase your total_range_bits.")
+    return quantized_cdf
+
+
+class ArithmeticCoder:
+    """ArithmeticCoder,
+    Let us take a distribution `p` over `N` symbols, and assume we have a stream
+    of random variables `s_t` sampled from `p`. Let us assume that we have a budget
+    of `B` bits that we can afford to write on device. There are `2**B` possible numbers,
+    corresponding to the range `[0, 2 ** B - 1]`. We can map each of those number to a single
+    sequence `(s_t)` by doing the following:
+
+    1) Initialize the current range to` [0 ** 2 B - 1]`.
+    2) For each time step t, split the current range into contiguous chunks,
+        one for each possible outcome, with size roughly proportional to `p`.
+        For instance, if `p = [0.75, 0.25]`, and the range is `[0, 3]`, the chunks
+        would be `{[0, 2], [3, 3]}`.
+    3) Select the chunk corresponding to `s_t`, and replace the current range with this.
+    4) When done encoding all the values, just select any value remaining in the range.
+
+    You will notice that this procedure can fail: for instance if at any point in time
+    the range is smaller than `N`, then we can no longer assign a non-empty chunk to each
+    possible outcome. Intuitively, the more likely a value is, the less the range width
+    will reduce, and the longer we can go on encoding values. This makes sense: for any efficient
+    coding scheme, likely outcomes would take less bits, and more of them can be coded
+    with a fixed budget.
+
+    In practice, we do not know `B` ahead of time, but we have a way to inject new bits
+    when the current range decreases below a given limit (given by `total_range_bits`), without
+    having to redo all the computations. If we encode mostly likely values, we will seldom
+    need to inject new bits, but a single rare value can deplete our stock of entropy!
+
+    In this explanation, we assumed that the distribution `p` was constant. In fact, the present
+    code works for any sequence `(p_t)` possibly different for each timestep.
+    We also assume that `s_t ~ p_t`, but that doesn't need to be true, although the smaller
+    the KL between the true distribution and `p_t`, the most efficient the coding will be.
+
+    Args:
+        fo (IO[bytes]): file-like object to which the bytes will be written to.
+        total_range_bits (int): the range `M` described above is `2 ** total_range_bits.
+            Any time the current range width fall under this limit, new bits will
+            be injected to rescale the initial range.
+    """
+
+    def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24):
+        assert total_range_bits <= 30
+        self.total_range_bits = total_range_bits
+        self.packer = BitPacker(bits=1, fo=fo)  # we push single bits at a time.
+        self.low: int = 0
+        self.high: int = 0
+        self.max_bit: int = -1
+        self._dbg: tp.List[tp.Any] = []
+        self._dbg2: tp.List[tp.Any] = []
+
+    @property
+    def delta(self) -> int:
+        """Return the current range width."""
+        return self.high - self.low + 1
+
+    def _flush_common_prefix(self):
+        # If self.low and self.high start with the sames bits,
+        # those won't change anymore as we always just increase the range
+        # by powers of 2, and we can flush them out to the bit stream.
+        assert self.high >= self.low, (self.low, self.high)
+        assert self.high < 2 ** (self.max_bit + 1)
+        while self.max_bit >= 0:
+            b1 = self.low >> self.max_bit
+            b2 = self.high >> self.max_bit
+            if b1 == b2:
+                self.low -= (b1 << self.max_bit)
+                self.high -= (b1 << self.max_bit)
+                assert self.high >= self.low, (self.high, self.low, self.max_bit)
+                assert self.low >= 0
+                self.max_bit -= 1
+                self.packer.push(b1)
+            else:
+                break
+
+    def push(self, symbol: int, quantized_cdf: torch.Tensor):
+        """Push the given symbol on the stream, flushing out bits
+        if possible.
+
+        Args:
+            symbol (int): symbol to encode with the AC.
+            quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf`
+                to build this from your pdf estimate.
+        """
+        while self.delta < 2 ** self.total_range_bits:
+            self.low *= 2
+            self.high = self.high * 2 + 1
+            self.max_bit += 1
+
+        range_low = 0 if symbol == 0 else quantized_cdf[symbol - 1].item()
+        range_high = quantized_cdf[symbol].item() - 1
+        effective_low = int(math.ceil(range_low * (self.delta / (2 ** self.total_range_bits))))
+        effective_high = int(math.floor(range_high * (self.delta / (2 ** self.total_range_bits))))
+        assert self.low <= self.high
+        self.high = self.low + effective_high
+        self.low = self.low + effective_low
+        assert self.low <= self.high, (effective_low, effective_high, range_low, range_high)
+        self._dbg.append((self.low, self.high))
+        self._dbg2.append((self.low, self.high))
+        outs = self._flush_common_prefix()
+        assert self.low <= self.high
+        assert self.max_bit >= -1
+        assert self.max_bit <= 61, self.max_bit
+        return outs
+
+    def flush(self):
+        """Flush the remaining information to the stream.
+        """
+        while self.max_bit >= 0:
+            b1 = (self.low >> self.max_bit) & 1
+            self.packer.push(b1)
+            self.max_bit -= 1
+        self.packer.flush()
+
+
+class ArithmeticDecoder:
+    """ArithmeticDecoder, see `ArithmeticCoder` for a detailed explanation.
+
+    Note that this must be called with **exactly** the same parameters and sequence
+    of quantized cdf as the arithmetic encoder or the wrong values will be decoded.
+
+    If the AC encoder current range is [L, H], with `L` and `H` having the some common
+    prefix (i.e. the same most significant bits), then this prefix will be flushed to the stream.
+    For instances, having read 3 bits `b1 b2 b3`, we know that `[L, H]` is contained inside
+    `[b1 b2 b3 0 ... 0 b1 b3 b3 1 ... 1]`. Now this specific sub-range can only be obtained
+    for a specific sequence of symbols and a binary-search allows us to decode those symbols.
+    At some point, the prefix `b1 b2 b3` will no longer be sufficient to decode new symbols,
+    and we will need to read new bits from the stream and repeat the process.
+
+    """
+    def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24):
+        self.total_range_bits = total_range_bits
+        self.low: int = 0
+        self.high: int = 0
+        self.current: int = 0
+        self.max_bit: int = -1
+        self.unpacker = BitUnpacker(bits=1, fo=fo)  # we pull single bits at a time.
+        # Following is for debugging
+        self._dbg: tp.List[tp.Any] = []
+        self._dbg2: tp.List[tp.Any] = []
+        self._last: tp.Any = None
+
+    @property
+    def delta(self) -> int:
+        return self.high - self.low + 1
+
+    def _flush_common_prefix(self):
+        # Given the current range [L, H], if both have a common prefix,
+        # we know we can remove it from our representation to avoid handling large numbers.
+        while self.max_bit >= 0:
+            b1 = self.low >> self.max_bit
+            b2 = self.high >> self.max_bit
+            if b1 == b2:
+                self.low -= (b1 << self.max_bit)
+                self.high -= (b1 << self.max_bit)
+                self.current -= (b1 << self.max_bit)
+                assert self.high >= self.low
+                assert self.low >= 0
+                self.max_bit -= 1
+            else:
+                break
+
+    def pull(self, quantized_cdf: torch.Tensor) -> tp.Optional[int]:
+        """Pull a symbol, reading as many bits from the stream as required.
+        This returns `None` when the stream has been exhausted.
+
+        Args:
+            quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf`
+                to build this from your pdf estimate. This must be **exatly**
+                the same cdf as the one used at encoding time.
+        """
+        while self.delta < 2 ** self.total_range_bits:
+            bit = self.unpacker.pull()
+            if bit is None:
+                return None
+            self.low *= 2
+            self.high = self.high * 2 + 1
+            self.current = self.current * 2 + bit
+            self.max_bit += 1
+
+        def bin_search(low_idx: int, high_idx: int):
+            # Binary search is not just for coding interviews :)
+            if high_idx < low_idx:
+                raise RuntimeError("Binary search failed")
+            mid = (low_idx + high_idx) // 2
+            range_low = quantized_cdf[mid - 1].item() if mid > 0 else 0
+            range_high = quantized_cdf[mid].item() - 1
+            effective_low = int(math.ceil(range_low * (self.delta / (2 ** self.total_range_bits))))
+            effective_high = int(math.floor(range_high * (self.delta / (2 ** self.total_range_bits))))
+            low = effective_low + self.low
+            high = effective_high + self.low
+            if self.current >= low:
+                if self.current <= high:
+                    return (mid, low, high, self.current)
+                else:
+                    return bin_search(mid + 1, high_idx)
+            else:
+                return bin_search(low_idx, mid - 1)
+
+        self._last = (self.low, self.high, self.current, self.max_bit)
+        sym, self.low, self.high, self.current = bin_search(0, len(quantized_cdf) - 1)
+        self._dbg.append((self.low, self.high, self.current))
+        self._flush_common_prefix()
+        self._dbg2.append((self.low, self.high, self.current))
+
+        return sym
+
+
+def test():
+    torch.manual_seed(1234)
+    random.seed(1234)
+    for _ in range(4):
+        pdfs = []
+        cardinality = random.randrange(4000)
+        steps = random.randrange(100, 500)
+        fo = io.BytesIO()
+        encoder = ArithmeticCoder(fo)
+        symbols = []
+        for step in range(steps):
+            pdf = torch.softmax(torch.randn(cardinality), dim=0)
+            pdfs.append(pdf)
+            q_cdf = build_stable_quantized_cdf(pdf, encoder.total_range_bits)
+            symbol = torch.multinomial(pdf, 1).item()
+            symbols.append(symbol)
+            encoder.push(symbol, q_cdf)
+        encoder.flush()
+
+        fo.seek(0)
+        decoder = ArithmeticDecoder(fo)
+        for idx, (pdf, symbol) in enumerate(zip(pdfs, symbols)):
+            q_cdf = build_stable_quantized_cdf(pdf, encoder.total_range_bits)
+            decoded_symbol = decoder.pull(q_cdf)
+            assert decoded_symbol == symbol, idx
+        assert decoder.pull(torch.zeros(1)) is None
+
+
+if __name__ == "__main__":
+    test()

clean/audio/safeear/safeear/models/modules/quantization/core_vq.py

@@ -0,0 +1,366 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This implementation is inspired from
+# https://github.com/lucidrains/vector-quantize-pytorch
+# which is released under MIT License. Hereafter, the original license:
+# MIT License
+#
+# Copyright (c) 2020 Phil Wang
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Core vector quantization implementation."""
+import typing as tp
+
+from einops import rearrange, repeat
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .distrib import broadcast_tensors, rank
+
+
+def default(val: tp.Any, d: tp.Any) -> tp.Any:
+    return val if val is not None else d
+
+
+def ema_inplace(moving_avg, new, decay: float):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
+    return (x + epsilon) / (x.sum() + n_categories * epsilon)
+
+
+def uniform_init(*shape: int):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def sample_vectors(samples, num: int):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters: int, num_iters: int = 10):
+    dim, dtype = samples.shape[-1], samples.dtype
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        diffs = rearrange(samples, "n d -> n () d") - rearrange(
+            means, "c d -> () c d"
+        )
+        dists = -(diffs ** 2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+class EuclideanCodebook(nn.Module):
+    """Codebook with Euclidean distance.
+    Args:
+        dim (int): Dimension.
+        codebook_size (int): Codebook size.
+        kmeans_init (bool): Whether to use k-means to initialize the codebooks.
+            If set to true, run the k-means algorithm on the first training batch and use
+            the learned centroids as initialization.
+        kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        kmeans_init: int = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.99,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+    ):
+        super().__init__()
+        self.decay = decay
+        init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        self.codebook_size = codebook_size
+
+        self.kmeans_iters = kmeans_iters
+        self.epsilon = epsilon
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer("inited", torch.Tensor([not kmeans_init]))
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.inited:
+            return
+
+        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed.clone())
+        self.cluster_size.data.copy_(cluster_size)
+        self.inited.data.copy_(torch.Tensor([True]))
+        # Make sure all buffers across workers are in sync after initialization
+        #broadcast_tensors(self.buffers())
+
+    def replace_(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace_(batch_samples, mask=expired_codes)
+        #broadcast_tensors(self.buffers())
+
+    def preprocess(self, x):
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def quantize(self, x):
+        embed = self.embed.t()
+        dist = -(
+            x.pow(2).sum(1, keepdim=True)
+            - 2 * x @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+        embed_ind = dist.max(dim=-1).indices
+        return embed_ind
+
+    def postprocess_emb(self, embed_ind, shape):
+        return embed_ind.view(*shape[:-1])
+
+    def dequantize(self, embed_ind):
+        quantize = F.embedding(embed_ind, self.embed)
+        return quantize
+
+    def encode(self, x):
+        shape = x.shape
+        # pre-process
+        x = self.preprocess(x)
+        # quantize
+        embed_ind = self.quantize(x)
+        # post-process
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        return embed_ind
+
+    def decode(self, embed_ind):
+        quantize = self.dequantize(embed_ind)
+        return quantize
+
+    def forward(self, x):
+        shape, dtype = x.shape, x.dtype
+        x = self.preprocess(x)
+
+        self.init_embed_(x)
+
+        embed_ind = self.quantize(x)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        quantize = self.dequantize(embed_ind)
+
+        if self.training:
+            # We do the expiry of code at that point as buffers are in sync
+            # and all the workers will take the same decision.
+            self.expire_codes_(x)
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = x.t() @ embed_onehot
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+
+        return quantize, embed_ind
+
+
+class VectorQuantization(nn.Module):
+    """Vector quantization implementation.
+    Currently supports only euclidean distance.
+    Args:
+        dim (int): Dimension
+        codebook_size (int): Codebook size
+        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+        commitment_weight (float): Weight for commitment loss.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        codebook_dim: tp.Optional[int] = None,
+        decay: float = 0.99,
+        epsilon: float = 1e-5,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 50,
+        threshold_ema_dead_code: int = 2,
+        commitment_weight: float = 1.,
+    ):
+        super().__init__()
+        _codebook_dim: int = default(codebook_dim, dim)
+
+        requires_projection = _codebook_dim != dim
+        self.project_in = (nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity())
+        self.project_out = (nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity())
+
+        self.epsilon = epsilon
+        self.commitment_weight = commitment_weight
+
+        self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size,
+                                           kmeans_init=kmeans_init, kmeans_iters=kmeans_iters,
+                                           decay=decay, epsilon=epsilon,
+                                           threshold_ema_dead_code=threshold_ema_dead_code)
+        self.codebook_size = codebook_size
+
+    @property
+    def codebook(self):
+        return self._codebook.embed
+
+    def encode(self, x):
+        x = rearrange(x, "b d n -> b n d")
+        x = self.project_in(x)
+        embed_in = self._codebook.encode(x)
+        return embed_in
+
+    def decode(self, embed_ind):
+        quantize = self._codebook.decode(embed_ind)
+        quantize = self.project_out(quantize)
+        quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize
+
+    def forward(self, x):
+        device = x.device
+        x = rearrange(x, "b d n -> b n d")
+        x = self.project_in(x)
+
+        quantize, embed_ind = self._codebook(x)
+
+        if self.training:
+            quantize = x + (quantize - x).detach()
+
+        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
+
+        if self.training:
+            if self.commitment_weight > 0:
+                commit_loss = F.mse_loss(quantize.detach(), x)
+                loss = loss + commit_loss * self.commitment_weight
+
+        quantize = self.project_out(quantize)
+        quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize, embed_ind, loss
+
+
+class ResidualVectorQuantization(nn.Module):
+    """Residual vector quantization implementation.
+    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+    """
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
+        )
+
+    def forward(self, x, n_q: tp.Optional[int] = None, layers: tp.Optional[list] = None):
+        quantized_out = 0.0
+        residual = x
+
+        all_losses = []
+        all_indices = []
+        out_quantized = []
+
+        n_q = n_q or len(self.layers)
+
+        for i, layer in enumerate(self.layers[:n_q]):
+            quantized, indices, loss = layer(residual)
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+
+            all_indices.append(indices)
+            all_losses.append(loss)
+            if layers and i in layers:
+                out_quantized.append(quantized)
+
+        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
+        return quantized_out, out_indices, out_losses, out_quantized
+
+    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int]= None) -> torch.Tensor:
+        residual = x
+        all_indices = []
+        n_q = n_q or len(self.layers)
+        st = st or 0
+        for layer in self.layers[st:n_q]: 
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            residual = residual - quantized
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, q_indices: torch.Tensor, st: int=0) -> torch.Tensor:
+        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        for i, indices in enumerate(q_indices):
+            layer = self.layers[st + i]
+            quantized = layer.decode(indices)
+            quantized_out = quantized_out + quantized
+        return quantized_out

clean/audio/safeear/safeear/models/modules/quantization/distrib.py

@@ -0,0 +1,126 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Torch distributed utilities."""
+
+import typing as tp
+
+import torch
+
+
+def rank():
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    else:
+        return 0
+
+
+def world_size():
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_world_size()
+    else:
+        return 1
+
+
+def is_distributed():
+    return world_size() > 1
+
+
+def all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM):
+    if is_distributed():
+        return torch.distributed.all_reduce(tensor, op)
+
+
+def _is_complex_or_float(tensor):
+    return torch.is_floating_point(tensor) or torch.is_complex(tensor)
+
+
+def _check_number_of_params(params: tp.List[torch.Tensor]):
+    # utility function to check that the number of params in all workers is the same,
+    # and thus avoid a deadlock with distributed all reduce.
+    if not is_distributed() or not params:
+        return
+    #print('params[0].device ', params[0].device)
+    tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long)
+    all_reduce(tensor)
+    if tensor.item() != len(params) * world_size():
+        # If not all the workers have the same number, for at least one of them,
+        # this inequality will be verified.
+        raise RuntimeError(f"Mismatch in number of params: ours is {len(params)}, "
+                           "at least one worker has a different one.")
+
+
+def broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0):
+    """Broadcast the tensors from the given parameters to all workers.
+    This can be used to ensure that all workers have the same model to start with.
+    """
+    if not is_distributed():
+        return
+    tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)]
+    _check_number_of_params(tensors)
+    handles = []
+    for tensor in tensors:
+        # src = int(rank()) # added code
+        handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True)
+        handles.append(handle)
+    for handle in handles:
+        handle.wait()
+
+
+def sync_buffer(buffers, average=True):
+    """
+    Sync grad for buffers. If average is False, broadcast instead of averaging.
+    """
+    if not is_distributed():
+        return
+    handles = []
+    for buffer in buffers:
+        if torch.is_floating_point(buffer.data):
+            if average:
+                handle = torch.distributed.all_reduce(
+                    buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
+            else:
+                handle = torch.distributed.broadcast(
+                    buffer.data, src=0, async_op=True)
+            handles.append((buffer, handle))
+    for buffer, handle in handles:
+        handle.wait()
+        if average:
+            buffer.data /= world_size
+
+
+def sync_grad(params):
+    """
+    Simpler alternative to DistributedDataParallel, that doesn't rely
+    on any black magic. For simple models it can also be as fast.
+    Just call this on your model parameters after the call to backward!
+    """
+    if not is_distributed():
+        return
+    handles = []
+    for p in params:
+        if p.grad is not None:
+            handle = torch.distributed.all_reduce(
+                p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
+            handles.append((p, handle))
+    for p, handle in handles:
+        handle.wait()
+        p.grad.data /= world_size()
+
+
+def average_metrics(metrics: tp.Dict[str, float], count=1.):
+    """Average a dictionary of metrics across all workers, using the optional
+    `count` as unormalized weight.
+    """
+    if not is_distributed():
+        return metrics
+    keys, values = zip(*metrics.items())
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32)
+    tensor *= count
+    all_reduce(tensor)
+    averaged = (tensor[:-1] / tensor[-1]).cpu().tolist()
+    return dict(zip(keys, averaged))

clean/audio/safeear/safeear/models/modules/quantization/vq.py

@@ -0,0 +1,108 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Residual vector quantizer implementation."""
+
+from dataclasses import dataclass, field
+import math
+import typing as tp
+
+import torch
+from torch import nn
+
+from .core_vq import ResidualVectorQuantization
+
+
+@dataclass
+class QuantizedResult:
+    quantized: torch.Tensor
+    codes: torch.Tensor
+    bandwidth: torch.Tensor  # bandwidth in kb/s used, per batch item.
+    penalty: tp.Optional[torch.Tensor] = None
+    metrics: dict = field(default_factory=dict)
+
+
+class ResidualVectorQuantizer(nn.Module):
+    """Residual Vector Quantizer.
+    Args:
+        dimension (int): Dimension of the codebooks.
+        n_q (int): Number of residual vector quantizers used.
+        bins (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dimension: int = 256,
+        n_q: int = 8,
+        bins: int = 1024,
+        decay: float = 0.99,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 50,
+        threshold_ema_dead_code: int = 2,
+    ):
+        super().__init__()
+        self.n_q = n_q
+        self.dimension = dimension
+        self.bins = bins
+        self.decay = decay
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.vq = ResidualVectorQuantization(
+            dim=self.dimension,
+            codebook_size=self.bins,
+            num_quantizers=self.n_q,
+            decay=self.decay,
+            kmeans_init=self.kmeans_init,
+            kmeans_iters=self.kmeans_iters,
+            threshold_ema_dead_code=self.threshold_ema_dead_code,
+        )
+
+    def forward(self, x: torch.Tensor, n_q: tp.Optional[int] = None, layers: tp.Optional[list] = None) -> QuantizedResult:
+        """Residual vector quantization on the given input tensor.
+        Args:
+            x (torch.Tensor): Input tensor.
+            n_q (int): Number of quantizer used to quantize. Default: All quantizers.
+            layers (list): Layer that need to return quantized. Defalt: None.
+        Returns:
+            QuantizedResult:
+                The quantized (or approximately quantized) representation with
+                the associated numbert quantizers and layer quantized required to return.
+        """
+        n_q = n_q if n_q else self.n_q
+        if layers and max(layers) >= n_q:
+            raise ValueError(f'Last layer index in layers: A {max(layers)}. Number of quantizers in RVQ: B {self.n_q}. A must less than B.')
+        quantized, codes, commit_loss, quantized_list = self.vq(x, n_q=n_q, layers=layers)
+        return quantized, codes, torch.mean(commit_loss), quantized_list
+
+
+    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        Args:
+            x (torch.Tensor): Input tensor.
+            n_q (int): Number of quantizer used to quantize. Default: All quantizers.
+            st (int): Start to encode input from which layers. Default: 0.
+        """
+        n_q = n_q if n_q else self.n_q
+        st = st or 0
+        codes = self.vq.encode(x, n_q=n_q, st=st)
+        return codes
+
+    def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor:
+        """Decode the given codes to the quantized representation.
+        Args:
+            codes (torch.Tensor): Input indices for each quantizer.
+            st (int): Start to decode input codes from which layers. Default: 0.
+        """
+        quantized = self.vq.decode(codes, st=st)
+        return quantized

clean/audio/safeear/safeear/models/modules/seanet.py

@@ -0,0 +1,275 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Encodec SEANet-based encoder and decoder implementation."""
+
+import typing as tp
+
+import numpy as np
+import torch.nn as nn
+import torch
+
+from . import (
+    SConv1d,
+    SConvTranspose1d,
+    SLSTM
+)
+
+
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+class SEANetResnetBlock(nn.Module):
+    """Residual block from SEANet model.
+    Args:
+        dim (int): Dimension of the input/output
+        kernel_sizes (list): List of kernel sizes for the convolutions.
+        dilations (list): List of dilations for the convolutions.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3)
+        true_skip (bool): Whether to use true skip connection or a simple convolution as the skip connection.
+    """
+    def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],
+                 activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,
+                 pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):
+        super().__init__()
+        assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'
+        act = getattr(nn, activation) if activation != 'Snake' else Snake1d
+        hidden = dim // compress
+        block = []
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [
+                act(**activation_params) if activation != 'Snake' else act(in_chs),
+                SConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,
+                        norm=norm, norm_kwargs=norm_params,
+                        causal=causal, pad_mode=pad_mode),
+            ]
+        self.block = nn.Sequential(*block)
+        self.shortcut: nn.Module
+        if true_skip:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = SConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,
+                                    causal=causal, pad_mode=pad_mode)
+
+    def forward(self, x):
+        return self.shortcut(x) + self.block(x)
+
+
+
+class SEANetEncoder(nn.Module):
+    """SEANet encoder.
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
+            upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
+            that must match the decoder order
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2, bidirectional:bool = False):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios) # 计算乘积
+
+        act = getattr(nn, activation) if activation != 'Snake' else Snake1d
+        mult = 1
+        model: tp.List[nn.Module] = [
+            SConv1d(channels, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+        # Downsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      norm=norm, norm_params=norm_params,
+                                      activation=activation, activation_params=activation_params,
+                                      causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            # Add downsampling layers
+            model += [
+                act(**activation_params) if activation != 'Snake' else act(mult * n_filters),
+                SConv1d(mult * n_filters, mult * n_filters * 2,
+                        kernel_size=ratio * 2, stride=ratio,
+                        norm=norm, norm_kwargs=norm_params,
+                        causal=causal, pad_mode=pad_mode),
+            ]
+            mult *= 2
+
+        if lstm:
+            model += [SLSTM(mult * n_filters, num_layers=lstm, bidirectional=bidirectional)]
+
+        mult = mult * 2 if bidirectional else mult
+        model += [
+            act(**activation_params) if activation != 'Snake' else act(mult * n_filters),
+            SConv1d(mult * n_filters, dimension, last_kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class SEANetDecoder(nn.Module):
+    """SEANet decoder.
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function
+        final_activation (str): Final activation function after all convolutions.
+        final_activation_params (dict): Parameters to provide to the activation function
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
+            If equal to 1.0, it means that all the trimming is done at the right.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,
+                 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2,
+                 trim_right_ratio: float = 1.0, bidirectional:bool = False):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+
+        act = getattr(nn, activation) if activation != 'Snake' else Snake1d
+        mult = int(2 ** len(self.ratios))
+        model: tp.List[nn.Module] = [
+            SConv1d(dimension, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+
+        if lstm:
+            model += [SLSTM(mult * n_filters, num_layers=lstm, bidirectional=bidirectional)]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            # Add upsampling layers
+            model += [
+                act(**activation_params) if activation != 'Snake' else act(mult * n_filters),
+                SConvTranspose1d(mult * n_filters, mult * n_filters // 2,
+                                 kernel_size=ratio * 2, stride=ratio,
+                                 norm=norm, norm_kwargs=norm_params,
+                                 causal=causal, trim_right_ratio=trim_right_ratio),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      activation=activation, activation_params=activation_params,
+                                      norm=norm, norm_params=norm_params, causal=causal,
+                                      pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            act(**activation_params) if activation != 'Snake' else act(n_filters),
+            SConv1d(n_filters, channels, last_kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+        # Add optional final activation to decoder (eg. tanh)
+        if final_activation is not None:
+            final_act = getattr(nn, final_activation)
+            final_activation_params = final_activation_params or {}
+            model += [
+                final_act(**final_activation_params)
+            ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, z):
+        y = self.model(z)
+        return y
+
+
+def test():
+    import torch
+    encoder = SEANetEncoder()
+    decoder = SEANetDecoder()
+    x = torch.randn(1, 1, 24000)
+    z = encoder(x)
+    print('z ', z.shape)
+    assert 1==2
+    assert list(z.shape) == [1, 128, 75], z.shape
+    y = decoder(z)
+    assert y.shape == x.shape, (x.shape, y.shape)
+
+
+if __name__ == '__main__':
+    test()

clean/audio/safeear/safeear/models/safeear.py

@@ -0,0 +1,959 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList, Linear, Dropout, LayerNorm, Identity, Parameter, init
+from timm.models.layers import trunc_normal_, DropPath
+import random
+from typing import Union
+import numpy as np
+import math
+from torch import Tensor
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+                     
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SELayer, self).__init__()
+        # print('se reduction: ', reduction)
+        # print(channel // reduction)
+        self.avg_pool = nn.AdaptiveAvgPool2d(1) # F_squeeze 
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):   # x: B*C*D*T
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class SEBasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16):
+        super(SEBasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes, 1)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.se = SELayer(planes, reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class Bottleneck(nn.Module):
+    expansion = 2
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class SEBottleneck(nn.Module):
+    expansion = 2
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16):
+        super(SEBottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.se = SELayer(planes * self.expansion, reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+        out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottle2neck(nn.Module):
+    expansion = 2
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 stride=1,
+                 downsample=None,
+                 baseWidth=26,
+                 scale=4,
+                 stype='normal'):
+        """ Constructor
+        Args:
+            inplanes: input channel dimensionality
+            planes: output channel dimensionality
+            stride: conv stride. Replaces pooling layer.
+            downsample: None when stride = 1
+            baseWidth: basic width of conv3x3
+            scale: number of scale.
+            type: 'normal': normal set. 'stage': first block of a new stage.
+        """
+        super(Bottle2neck, self).__init__()
+
+        width = int(math.floor(planes * (baseWidth / 64.0)))
+        self.conv1 = nn.Conv2d(inplanes,
+                               width * scale,
+                               kernel_size=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(width * scale)
+
+        if scale == 1:
+            self.nums = 1
+        else:
+            self.nums = scale - 1
+        if stype == 'stage':
+            self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)
+        convs = []
+        bns = []
+        for i in range(self.nums):
+            convs.append(
+                nn.Conv2d(width,
+                          width,
+                          kernel_size=3,
+                          stride=stride,
+                          padding=1,
+                          bias=False))
+            bns.append(nn.BatchNorm2d(width))
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList(bns)
+
+        self.conv3 = nn.Conv2d(width * scale,
+                               planes * self.expansion,
+                               kernel_size=1,
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        if stride != 1 or inplanes != planes * self.expansion:
+            downsample = nn.Sequential(
+                nn.AvgPool2d(kernel_size=stride,
+                             stride=stride,
+                             ceil_mode=True,
+                             count_include_pad=False),
+                nn.Conv2d(inplanes,
+                          planes * self.expansion,
+                          kernel_size=1,
+                          stride=1,
+                          bias=False),
+                nn.BatchNorm2d(planes * self.expansion),
+            )
+        self.downsample = downsample
+        self.stype = stype
+        self.scale = scale
+        self.width = width
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        spx = torch.split(out, self.width, 1)
+        for i in range(self.nums):
+            if i == 0 or self.stype == 'stage':
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.convs[i](sp)
+            sp = self.relu(self.bns[i](sp))
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        if self.scale != 1 and self.stype == 'normal':
+            out = torch.cat((out, spx[self.nums]), 1)
+        elif self.scale != 1 and self.stype == 'stage':
+            out = torch.cat((out, self.pool(spx[self.nums])), 1)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class SEBottle2neck(nn.Module):
+    expansion = 1
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 stride=1,
+                 kernel_size = 1,
+                 padding=1,
+                 downsample=None,
+                 baseWidth=26,
+                 scale=4,
+                 stype='normal'):
+        """ Constructor
+        Args:
+            inplanes: input channel dimensionality
+            planes: output channel dimensionality
+            stride: conv stride. Replaces pooling layer.
+            downsample: None when stride = 1
+            baseWidth: basic width of conv3x3
+            scale: number of scale.
+            type: 'normal': normal set. 'stage': first block of a new stage.
+        """
+        super(SEBottle2neck, self).__init__()
+
+        width = int(math.floor(planes * (baseWidth / 64.0)))
+        self.conv1 = nn.Conv2d(inplanes,
+                               width * scale,
+                               kernel_size=kernel_size,
+                               padding=padding,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(width * scale)
+
+        if scale == 1:
+            self.nums = 1
+        else:
+            self.nums = scale - 1
+        if stype == 'stage':
+            self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)
+        convs = []
+        bns = []
+        for i in range(self.nums):
+            convs.append(
+                nn.Conv2d(width,
+                          width,
+                          kernel_size=3,
+                          stride=stride,
+                          padding=1,
+                          bias=False))
+            bns.append(nn.BatchNorm2d(width))
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList(bns)
+
+        self.conv3 = nn.Conv2d(width * scale,
+                               planes * self.expansion,
+                               kernel_size=kernel_size,
+                               padding=(0,1),
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.se = SELayer(planes * self.expansion, reduction=16)
+        self.relu = nn.ReLU(inplace=True)
+        if inplanes != planes:
+            self.downsample = True
+            self.conv_downsample = nn.Conv2d(in_channels=inplanes,
+                                             out_channels=planes,
+                                             padding=(0, 0),
+                                             kernel_size=(1, 1),
+                                             stride=1)
+
+        else:
+            self.downsample = False
+        self.stype = stype
+        self.scale = scale
+        self.width = width
+
+    def forward(self, x):
+        residual = x 
+        out = self.conv1(x) 
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        spx = torch.split(out, self.width, 1)
+        for i in range(self.nums):
+            if i == 0 or self.stype == 'stage':
+                sp = spx[i] 
+            else:
+                sp = sp + spx[i]
+            sp = self.convs[i](sp)
+            sp = self.relu(self.bns[i](sp)) 
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1) 
+        if self.scale != 1 and self.stype == 'normal':
+            out = torch.cat((out, spx[self.nums]), 1) 
+        elif self.scale != 1 and self.stype == 'stage':
+            out = torch.cat((out, self.pool(spx[self.nums])), 1)
+
+        out = self.conv3(out) 
+        out = self.bn3(out)
+        out = self.se(out) 
+
+        if self.downsample:
+            residual = self.conv_downsample(residual)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class Res2Net(nn.Module):
+    def __init__(self, block, layers, baseWidth=26, scale=4, m=0.35, num_classes=1000, loss='softmax', **kwargs):
+        self.inplanes = 16
+        super(Res2Net, self).__init__()
+        self.loss = loss
+        self.baseWidth = baseWidth
+        self.scale = scale
+        self.conv1 = nn.Sequential(nn.Conv2d(1, 16, 3, 1, 1, bias=False),
+                                   nn.BatchNorm2d(16), nn.ReLU(inplace=True),
+                                   nn.Conv2d(16, 16, 3, 1, 1, bias=False),
+                                   nn.BatchNorm2d(16), nn.ReLU(inplace=True),
+                                   nn.Conv2d(16, 16, 3, 1, 1, bias=False))
+        self.bn1 = nn.BatchNorm2d(16)
+        self.relu = nn.ReLU()
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 16, layers[0])#64
+        self.layer2 = self._make_layer(block, 32, layers[1], stride=2)#128
+        self.layer3 = self._make_layer(block, 64, layers[2], stride=2)#256
+        self.layer4 = self._make_layer(block, 128, layers[3], stride=2)#512
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        # self.stats_pooling = StatsPooling()
+
+        if self.loss == 'softmax':
+            # self.cls_layer = nn.Linear(2*8*128*block.expansion, num_classes)
+            self.cls_layer = nn.Linear(128*block.expansion, num_classes)
+        else:
+            raise NotImplementedError
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight,
+                                        mode='fan_out',
+                                        nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.AvgPool2d(kernel_size=stride,
+                             stride=stride,
+                             ceil_mode=True,
+                             count_include_pad=False),
+                nn.Conv2d(self.inplanes,
+                          planes * block.expansion,
+                          kernel_size=1,
+                          stride=1,
+                          bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(self.inplanes,
+                  planes,
+                  stride,
+                  downsample=downsample,
+                  stype='stage',
+                  baseWidth=self.baseWidth,
+                  scale=self.scale))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(
+                block(self.inplanes,
+                      planes,
+                      baseWidth=self.baseWidth,
+                      scale=self.scale))
+
+        return nn.Sequential(*layers)
+
+    def _forward(self, x):
+        x = x.unsqueeze(dim=1)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.cls_layer(x)
+
+        return F.log_softmax(x, dim=-1)
+
+    def extract(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return x
+    # Allow for accessing forward method in a inherited class
+    forward = _forward
+    
+def se_res2net50_v1b_14w_8s(**kwargs):
+    """Constructs a Res2Net-50_v1b model.
+    Res2Net-50 refers to the Res2Net-50_v1b_26w_4s.
+    """
+    model = Res2Net(SEBottle2neck, [3, 4, 6, 3], baseWidth=14, scale=8, **kwargs)
+    return model
+
+
+class CONV(nn.Module):
+    @staticmethod
+    def to_mel(hz):
+        return 2595 * np.log10(1 + hz / 700)
+
+    @staticmethod
+    def to_hz(mel):
+        return 700 * (10**(mel / 2595) - 1)
+
+    def __init__(self,
+                 out_channels,
+                 kernel_size,
+                 sample_rate=16000,
+                 in_channels=1,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 bias=False,
+                 groups=1,
+                 mask=False):
+        super().__init__()
+        if in_channels != 1:
+
+            msg = "SincConv only support one input channel (here, in_channels = {%i})" % (
+                in_channels)
+            raise ValueError(msg)
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.sample_rate = sample_rate
+
+        # Forcing the filters to be odd (i.e, perfectly symmetrics)
+        if kernel_size % 2 == 0:
+            self.kernel_size = self.kernel_size + 1
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.mask = mask
+        if bias:
+            raise ValueError('SincConv does not support bias.')
+        if groups > 1:
+            raise ValueError('SincConv does not support groups.')
+
+        NFFT = 512
+        f = int(self.sample_rate / 2) * np.linspace(0, 1, int(NFFT / 2) + 1)
+        fmel = self.to_mel(f)
+        fmelmax = np.max(fmel)
+        fmelmin = np.min(fmel)
+        filbandwidthsmel = np.linspace(fmelmin, fmelmax, self.out_channels + 1)
+        filbandwidthsf = self.to_hz(filbandwidthsmel)
+
+        self.mel = filbandwidthsf
+        self.hsupp = torch.arange(-(self.kernel_size - 1) / 2,
+                                  (self.kernel_size - 1) / 2 + 1)
+        self.band_pass = torch.zeros(self.out_channels, self.kernel_size)
+        for i in range(len(self.mel) - 1):
+            fmin = self.mel[i]
+            fmax = self.mel[i + 1]
+            hHigh = (2*fmax/self.sample_rate) * \
+                np.sinc(2*fmax*self.hsupp/self.sample_rate)
+            hLow = (2*fmin/self.sample_rate) * \
+                np.sinc(2*fmin*self.hsupp/self.sample_rate)
+            hideal = hHigh - hLow
+
+            self.band_pass[i, :] = Tensor(np.hamming(
+                self.kernel_size)) * Tensor(hideal)
+
+    def forward(self, x, mask=False):
+        band_pass_filter = self.band_pass.clone().to(x.device)
+        if mask:
+            A = np.random.uniform(0, 20)
+            A = int(A)
+            A0 = random.randint(0, band_pass_filter.shape[0] - A)
+            band_pass_filter[A0:A0 + A, :] = 0
+        else:
+            band_pass_filter = band_pass_filter
+
+        self.filters = (band_pass_filter).view(self.out_channels, 1,
+                                               self.kernel_size)
+
+        return F.conv1d(x,
+                        self.filters,
+                        stride=self.stride,
+                        padding=self.padding,
+                        dilation=self.dilation,
+                        bias=None,
+                        groups=1)
+
+
+class My_Residual_block(nn.Module):
+    def __init__(self, nb_filts, first=False, conv1=[2, 3, 1, 1, 1, 1], conv2=[2, 3, 0, 1, 1, 3], conv3=[1, 3, 0, 1, 1, 3], pool=(1, 3)):
+        super().__init__()
+        self.first = first
+
+        if not self.first:
+            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
+        self.conv1 = nn.Conv2d(in_channels=nb_filts[0],
+                               out_channels=nb_filts[1],
+                               kernel_size=(conv1[0], conv1[1]),
+                               padding=(conv1[2], conv1[3]),
+                               stride=(conv1[4], conv1[5]))
+        self.selu = nn.SELU(inplace=True)
+
+        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
+        self.conv2 = nn.Conv2d(in_channels=nb_filts[1],
+                               out_channels=nb_filts[1],
+                               kernel_size=(conv2[0], conv2[1]),
+                               padding=(conv2[2], conv2[3]),
+                               stride=(conv2[4], conv2[5]))
+
+        self.downsample = True
+        self.conv_downsample = nn.Conv2d(in_channels=nb_filts[0],
+                                         out_channels=nb_filts[1],
+                                         kernel_size=(conv3[0], conv3[1]),
+                                         padding=(conv3[2], conv3[3]),
+                                         stride=(conv3[4], conv3[5]))
+
+        # self.mp = nn.MaxPool2d((1,4))
+        self.mp = nn.MaxPool2d((pool[0], pool[1]))
+
+    def forward(self, x):
+        identity = x
+        if not self.first:
+            out = self.bn1(x) 
+            out = self.selu(out)  
+        else:
+            out = x
+        out = self.conv1(x)
+
+        out = self.bn2(out)  
+        out = self.selu(out)  
+        out = self.conv2(out)
+
+        if self.downsample:
+            identity = self.conv_downsample(identity)
+
+        out += identity
+        out = self.mp(out)
+        return out
+
+
+class My_SERes2Net_block(nn.Module):
+    def __init__(self, nb_filts, first=False, conv1=[2, 3, 1, 1, 1, 1], conv2=[3, 3, 1, 1, 1, 3], conv3=[1, 3, 0, 1, 1, 3], pool=(1, 3), radix=2, groups=2):
+        super().__init__()
+        self.first = first
+
+        if not self.first:
+            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
+
+        self.conv1 = SEBottle2neck(inplanes=nb_filts[0],
+                                   planes=nb_filts[1], kernel_size=(conv1[0], conv1[1]))
+        self.selu = nn.SELU(inplace=True)
+
+        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
+        self.conv2 = nn.Conv2d(in_channels=nb_filts[1],
+                               out_channels=nb_filts[1],
+                               kernel_size=(conv2[0], conv2[1]),
+                               padding=(conv2[2], conv2[3]),
+                               stride=(conv2[4], conv2[5]))
+
+        self.downsample = True
+        self.conv_downsample = nn.Conv2d(in_channels=nb_filts[0],
+                                         out_channels=nb_filts[1],
+                                         kernel_size=(conv3[0], conv3[1]),
+                                         padding=(conv3[2], conv3[3]),
+                                         stride=(conv3[4], conv3[5]))
+
+        self.mp = nn.MaxPool2d((pool[0], pool[1]))
+
+    def forward(self, x): 
+        identity = x
+        if not self.first:
+            out = self.bn1(x)  
+            out = self.selu(out)  
+        else:
+            out = x
+
+        out = self.conv1(x)
+        out = self.bn2(out)
+        out = self.selu(out)
+        out = self.conv2(out)
+        if self.downsample:
+            identity = self.conv_downsample(identity)
+
+        out += identity
+        out = self.mp(out)
+        return out
+
+
+class Attention(Module):
+    """
+    Obtained from timm: github.com:rwightman/pytorch-image-models
+    """
+
+    def __init__(self, dim, num_heads=8, attention_dropout=0.1, projection_dropout=0.1):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // self.num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = Linear(dim, dim * 3, bias=False)
+        self.attn_drop = Dropout(attention_dropout)
+        self.proj = Linear(dim, dim)
+        self.proj_drop = Dropout(projection_dropout)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C //
+                                  self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SimpleRelativeAttention(nn.Module):
+    # we implement this relative position embedding here., this is not used in our experiments
+    def __init__(self, dim, seq_length, num_heads=8, qkv_bias=True, qk_scale=None, attn_drop=0.1, proj_drop=0.1):
+        super().__init__()
+        self.dim = dim
+        self.length = seq_length
+        self.num_heads = num_heads
+        head_dim = dim//num_heads
+        self.scale = qk_scale or head_dim**-0.5
+        self.relative_position_table = nn.Parameter(
+            torch.zeros(size=(seq_length*2-1, num_heads)))
+        coords = torch.arange(seq_length)
+        relative_coords = coords[:, None]-coords[None, :]
+        relative_coords = relative_coords+seq_length-1
+        self.register_buffer('relative_index', relative_coords)
+        self.qkv = nn.Linear(dim, dim*3)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_table, std=0.02)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C //
+                                  self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q = q*self.scale
+        attn = torch.einsum('bhqe,bhke->bhqk', q, k)
+        relative_position_bias = self.relative_position_table[self.relative_index.reshape(-1)].reshape(
+            self.length, self.length, self.num_heads
+        )
+        relative_position_bias = relative_position_bias.permute(
+            2, 0, 1).contiguous()
+        attn = attn+relative_position_bias.unsqueeze(0)
+        attn = attn.softmax(-1)
+        attn = self.attn_drop(attn)
+        x = torch.einsum('bnqk,bnqe->bnqe', attn,
+                         v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj_drop(self.proj(x))
+        return x
+
+
+class TransformerEncoderLayer(Module):
+    """
+    Inspired by torch.nn.TransformerEncoderLayer and timm.
+    """
+
+    def __init__(self, d_model, nhead, atten=Attention, dim_feedforward=2048, dropout=0.1,
+                 attention_dropout=0.1, drop_path_rate=0.1):
+        super(TransformerEncoderLayer, self).__init__()
+        self.pre_norm = LayerNorm(d_model)
+        self.self_attn = atten(dim=d_model, num_heads=nhead,
+                               attention_dropout=attention_dropout, projection_dropout=dropout)
+
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.dropout1 = Dropout(dropout)
+        self.norm1 = LayerNorm(d_model)
+        self.linear2 = Linear(dim_feedforward, d_model)
+        self.dropout2 = Dropout(dropout)
+
+        self.drop_path = DropPath(
+            drop_path_rate) if drop_path_rate > 0 else Identity()
+
+        self.activation = F.gelu
+
+    def forward(self, src: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        src = src + self.drop_path(self.self_attn(self.pre_norm(src)))
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout1(self.activation(self.linear1(src))))
+        src = src + self.drop_path(self.dropout2(src2))
+        return src
+
+
+class TransformerClassifier(Module):
+    """
+    Adopted from https://github.com/SHI-Labs/Compact-Transformers.git
+    """
+
+    def __init__(self,
+                 embedding_dim=768,
+                 num_classes=1000,
+                 num_layers=12,
+                 num_heads=12,
+                 mlp_ratio=4.0,
+                 dropout_rate=0.1,
+                 attention_dropout=0.1,
+                 stochastic_depth_rate=0.1,
+                 positional_embedding='sine',
+                 sequence_length=10000,
+                 *args, **kwargs):
+        super().__init__()
+        positional_embedding = positional_embedding if \
+            positional_embedding in ['sine', 'learnable', 'none'] else 'sine'
+        dim_feedforward = int(embedding_dim * mlp_ratio)
+        self.embedding_dim = embedding_dim
+        self.sequence_length = sequence_length
+
+        assert sequence_length is not None or positional_embedding == 'none'
+
+        if positional_embedding != 'none':
+            if positional_embedding == 'learnable':
+                self.positional_emb = Parameter(torch.zeros(1, sequence_length, embedding_dim),
+                                                requires_grad=True)
+                init.trunc_normal_(self.positional_emb, std=0.2)
+            else:
+                print('here!!! sinusoidal_embedding')
+                self.positional_emb = Parameter(self.sinusoidal_embedding(sequence_length, embedding_dim),
+                                                requires_grad=False)
+        else:
+            self.positional_emb = None
+
+        self.dropout = Dropout(p=dropout_rate)
+        dpr = [x.item() for x in torch.linspace(
+            0, stochastic_depth_rate, num_layers)]
+        self.blocks = ModuleList([
+            TransformerEncoderLayer(d_model=embedding_dim, nhead=num_heads,
+                                    dim_feedforward=dim_feedforward, dropout=dropout_rate,
+                                    attention_dropout=attention_dropout, drop_path_rate=dpr[i])
+            for i in range(num_layers)])
+        self.norm = LayerNorm(embedding_dim)
+        self.flattener = nn.Flatten(2, 3)
+        self.attention_pool = Linear(self.embedding_dim, 1)
+        self.fc = Linear(embedding_dim, num_classes)
+        self.apply(self.init_weight)
+
+    def forward(self, x):
+        x = torch.transpose(x,-1,-2)
+        seq_len = x.size(1)
+        x += self.positional_emb[:, :seq_len, :]
+
+        x = self.dropout(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        
+        feature = torch.matmul(F.softmax(self.attention_pool(
+            x), dim=1).transpose(-1, -2), x).squeeze(-2)
+        logits = self.fc(feature)
+        
+        return logits, feature
+
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, Linear):
+            init.trunc_normal_(m.weight, std=.02)
+            if isinstance(m, Linear) and m.bias is not None:
+                init.constant_(m.bias, 0)
+        elif isinstance(m, LayerNorm):
+            init.constant_(m.bias, 0)
+            init.constant_(m.weight, 1.0)
+
+    @staticmethod
+    def sinusoidal_embedding(n_channels, dim):
+        pe = torch.FloatTensor([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
+                                for p in range(n_channels)])
+        pe[:, 0::2] = torch.sin(pe[:, 0::2])
+        pe[:, 1::2] = torch.cos(pe[:, 1::2])
+        return pe.unsqueeze(0)
+    
+class SE_Rawformer_front(nn.Module):
+    def __init__(self, conv1 = [2,3,1,1,1,1],conv2 = [3,3,1,1,1,2],conv3 = [1,3,0,1,1,2]):
+        super().__init__()
+        filts = [70, [1, 32], [32, 32], [32, 64], [64, 64]]
+        self.conv_time = CONV(out_channels=filts[0],
+                              kernel_size=128,
+                              in_channels=1)  # 70 129
+        self.first_bn = nn.BatchNorm2d(num_features=1)
+        self.drop = nn.Dropout(0.5, inplace=True)
+        self.drop_way = nn.Dropout(0.2, inplace=True)
+        self.selu = nn.SELU(inplace=True)
+
+        self.encoder = nn.Sequential(
+            nn.Sequential(My_Residual_block(nb_filts=filts[1], conv1 = conv1,conv2 = [2,3,0,1,1,2],conv3 = conv3,first=True)),
+            nn.Sequential(My_SERes2Net_block(nb_filts=filts[2], conv1 = conv1,conv2 = conv2,conv3 = conv3)),
+            nn.Sequential(My_SERes2Net_block(nb_filts=filts[3], conv1 = conv1,conv2 = conv2,conv3 = conv3)),
+            nn.Sequential(My_SERes2Net_block(nb_filts=filts[4], conv1 = conv1,conv2 = conv2,conv3 = conv3)))
+
+    def forward(self, x, Freq_aug=False):
+        x = self.conv_time(x, mask=Freq_aug)  
+        x = x.unsqueeze(dim=1) 
+        x = F.max_pool2d(torch.abs(x), (3, 3)) 
+        x = self.first_bn(x) 
+        x = self.selu(x) 
+
+        encoder = self.encoder(x) 
+        return encoder
+    
+
+class SafeEar(nn.Module):
+    def __init__(self,front, *args, **kwargs):
+        super().__init__()
+        # self.front = front
+        self.bottleneck = nn.Sequential(
+            nn.Conv1d(kwargs["embedding_dim"]*7, kwargs["embedding_dim"], kernel_size=1),
+            nn.BatchNorm1d(kwargs["embedding_dim"])
+        )
+        self.classifier = TransformerClassifier(*args, **kwargs)
+        
+    def forward(self, encoder):
+        encoder = self.bottleneck(torch.cat(encoder, dim=1))
+        batch_size, feature_dim, frame_num = encoder.size()
+
+        for i in range(0, frame_num, 50):
+            encoder[:, :, i:i+50] = torch.flip(encoder[:, :, i:i+50], dims=[2])
+
+        logits, feature = self.classifier(encoder) 
+
+        return logits, feature
+ 
+
+class SafeEar1s(nn.Module):
+    def __init__(self,front, *args, **kwargs):
+        super().__init__()
+        # self.front = front
+        self.bottleneck = nn.Sequential(
+            nn.Conv1d(kwargs["embedding_dim"]*7, kwargs["embedding_dim"], kernel_size=1),
+            nn.BatchNorm1d(kwargs["embedding_dim"])
+        )
+        self.classifier = TransformerClassifier(*args, **kwargs)
+        
+    def forward(self, encoder):
+        encoder = self.bottleneck(torch.cat(encoder, dim=1))
+        
+        batch_size, feature_dim, frame_num = encoder.size()
+        for i in range(0, frame_num, 50):
+            end = i+50 if i+50 <= frame_num else frame_num
+            indices = torch.randperm(end - i)
+            encoder[:, :, i:end] = encoder[:, :, i+indices]
+
+        logits, feature = self.classifier(encoder) 
+
+        return logits, feature

clean/audio/safeear/safeear/trainer/safeear_trainer.py

@@ -0,0 +1,189 @@
+import math
+import torch
+import pytorch_lightning as pl
+from ..losses.loss import compute_eer
+import numpy as np
+import warnings
+warnings.filterwarnings("ignore")
+
+def get_input(x):
+    x = x.to(memory_format=torch.contiguous_format)
+    return x.float()
+
+class SafeEarTrainer(pl.LightningModule):
+    def __init__(
+            self,
+            decouple_model,
+            detect_model,
+            lr_raw_former,
+            save_score_path
+        ) -> None:
+        super().__init__()
+        
+        self.decouple_model = decouple_model
+        self.detect_model = detect_model
+        self.lr_raw_former = lr_raw_former
+        self.save_score_path = save_score_path 
+            
+        self.detect_loss = torch.nn.BCELoss()
+        
+        self.automatic_optimization = False
+        
+        self.val_index_loader = []
+        self.val_score_loader = []
+        self.eval_index_loader = []
+        self.eval_score_loader = []
+        self.eval_filename_loader = []
+        self.default_monitor = "val_eer"
+
+    def forward(self, batch, is_train=True):
+        if is_train:
+            x, feat, target = batch
+        else:
+            if len(batch) == 4:
+                x, feat, target, audio_path = batch
+            else:
+                x, feat, target = batch
+                audio_path = None
+        x_wav = get_input(x)
+        with torch.no_grad():
+            self.decouple_model.eval()
+            G_x, commit_loss, last_layer, acoustic_tokens = self.decouple_model(x_wav, layers=[0,1,2,3,4,5,6,7])
+        raw_logits, raw_feature = self.detect_model(acoustic_tokens)
+        
+        if is_train:
+            onehot_target = torch.eye(2).to(self.device)[target, :]
+            raw_logits = torch.softmax(raw_logits, dim=-1)
+            raw_former_loss_ = self.detect_loss(raw_logits,onehot_target)
+            return raw_former_loss_, raw_logits, target
+        else:
+            raw_logits = torch.softmax(raw_logits, dim=-1)[:, 0]
+            raw_former_loss_ = 0
+            return audio_path, raw_former_loss_, raw_logits, target
+        
+    def training_step(self, batch, batch_idx):
+        raw_opt = self.optimizers()
+        
+        raw_former_loss_, raw_logits, target = self(batch, is_train=True)
+        raw_opt.zero_grad()
+        self.manual_backward(raw_former_loss_)
+        raw_opt.step()
+        
+        self.log_dict(
+            {   
+                'train_loss': raw_former_loss_
+            },
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+            logger=True)
+    
+    def validation_step(self, batch, batch_idx):
+        _, raw_former_loss_, raw_logits, target = self(batch, is_train=False)
+        
+        self.val_index_loader.append(target)
+        self.val_score_loader.append(raw_logits)
+        
+        self.log_dict(
+            {
+                'val_loss': raw_former_loss_,
+            },
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+            logger=True)
+        
+    def on_validation_epoch_end(self):
+        all_index = self.all_gather(torch.cat(self.val_index_loader, dim=0)).view(-1).cpu().numpy()
+        all_score = self.all_gather(torch.cat(self.val_score_loader, dim=0)).view(-1).cpu().numpy()
+        val_eer = compute_eer(all_score[all_index == 0], all_score[all_index == 1])[0]
+        other_val_eer = compute_eer(-all_score[all_index == 0], -all_score[all_index == 1])[0]
+        val_eer = min(val_eer, other_val_eer)
+        self.log_dict(
+            {
+                "val_eer": val_eer,
+            },
+            sync_dist=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True)
+        
+        self.val_index_loader.clear()  # free memory
+        self.val_score_loader.clear()  # free memory
+        
+        self.log_dict(
+            {
+                "lr": self.optimizers().param_groups[0]['lr'],
+            },
+            sync_dist=True,
+            on_epoch=True,
+            prog_bar=False,
+            logger=True
+        )
+        
+        adjust_learning_rate(self.optimizers(), self.current_epoch, self.lr_raw_former, self.trainer.max_epochs*0.1, self.trainer.max_epochs)
+        
+    def test_step(self, batch, batch_idx):
+          
+        audio_path, raw_former_loss_, raw_logits, target = self(batch, is_train=False)
+        
+        self.eval_index_loader.append(target)
+        self.eval_score_loader.append(raw_logits)
+        self.eval_filename_loader.append(audio_path)
+        self.log_dict(
+            {
+                'val_loss_rawformer': raw_former_loss_,
+            },
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+            logger=True)
+        
+    def on_test_epoch_end(self):
+        
+        string_list = [list(item) for item in self.eval_filename_loader]
+    
+        all_filename = np.array(string_list)
+        all_filename = all_filename.reshape(-1, 1)
+        
+       
+        all_index = self.all_gather(torch.cat(self.eval_index_loader, dim=0)).view(-1).cpu().numpy()
+        all_score = self.all_gather(torch.cat(self.eval_score_loader, dim=0)).view(-1).cpu().numpy()
+        
+        # gpu_id = torch.cuda.current_device()
+        
+        data_to_write = zip(all_filename, all_score,all_index)
+        csv_filename = self.save_score_path + '/score.csv'
+        eval_eer = compute_eer(all_score[all_index == 0], all_score[all_index == 1])[0]
+        other_eval_eer = compute_eer(-all_score[all_index == 0], -all_score[all_index == 1])[0]   
+        eval_eer = min(eval_eer, other_eval_eer)
+        
+        self.log_dict(
+            {
+                "test_eer": eval_eer,
+            },
+            sync_dist=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True)
+        
+        self.eval_index_loader.clear()  # free memory
+        self.eval_score_loader.clear()  # free memory
+        self.eval_filename_loader.clear()  # free memory
+        
+    def configure_optimizers(self):
+        optimizer_rawformer = torch.optim.AdamW(self.detect_model.parameters(), lr=self.lr_raw_former, weight_decay=1e-4)
+        
+        return [optimizer_rawformer]
+        
+def adjust_learning_rate(optimizer, epoch, lr, warmup, epochs=100):
+    lr = lr
+    if epoch < warmup:
+        lr = lr / (warmup - epoch)
+    else:
+        lr *= 0.5 * (1. + math.cos(math.pi *
+                     (epoch - warmup) / (epochs - warmup)))
+
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr

clean/audio/safeear/safeear/utils/dump_hubert_feature.py

@@ -0,0 +1,108 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+import sys
+import tqdm
+sys.path.append('../../fairseq_ours/') # we recommend an abosulte path here.
+import fairseq
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+from npy_append_array import NpyAppendArray
+from pathlib import Path
+# from feature_utils import get_path_iterator, dump_feature
+
+
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+    stream=sys.stdout,
+)
+logger = logging.getLogger("dump_hubert_feature")
+
+
+class HubertFeatureReader(object):
+    def __init__(self, ckpt_path, layer, max_chunk=1600000):
+        (
+            model,
+            cfg,
+            task,
+        ) = fairseq.checkpoint_utils.load_model_ensemble_and_task([ckpt_path])
+        self.model = model[0].eval().cuda()
+        self.task = task
+        self.layer = layer
+        self.max_chunk = max_chunk
+        logger.info(f"TASK CONFIG:\n{self.task.cfg}")
+        logger.info(f" max_chunk = {self.max_chunk}")
+
+    def read_audio(self, path, ref_len=None):
+        wav, sr = sf.read(path)
+        assert sr == self.task.cfg.sample_rate, sr
+        if wav.ndim == 2:
+            wav = wav.mean(-1)
+        assert wav.ndim == 1, wav.ndim
+        if ref_len is not None and abs(ref_len - len(wav)) > 160:
+            logging.warning(f"ref {ref_len} != read {len(wav)} ({path})")
+        return wav
+
+    def get_feats(self, path, ref_len=None):
+        x = self.read_audio(path, ref_len)
+        with torch.no_grad():
+            x = torch.from_numpy(x).float().cuda()
+            if self.task.cfg.normalize:
+                x = F.layer_norm(x, x.shape)
+            x = x.view(1, -1)
+
+            feat = []
+            for start in range(0, x.size(1), self.max_chunk):
+                x_chunk = x[:, start: start + self.max_chunk]
+                feat_chunk, _, _ = self.model.extract_features(
+                    source=x_chunk,
+                    padding_mask=None,
+                    mask=False,
+                    output_layer=self.layer,
+                )
+                feat.append(feat_chunk)
+        return torch.cat(feat, 1).squeeze(0)
+
+def dump_feature(reader,audio_dir,save_dir):
+    save_dir = Path(save_dir)
+    audio_dir = Path(audio_dir)
+    
+    audio_files = list(audio_dir.glob("**/*.flac"))
+    
+    for audio_file in tqdm.tqdm(audio_files):
+        releative_path = audio_file.relative_to(audio_dir).with_suffix(".npy")
+        save_path = save_dir / releative_path
+        # import pdb; pdb.set_trace()
+        if not save_path.parent.exists():
+            save_path.parent.mkdir(parents=True)
+        
+        feat_f = NpyAppendArray(save_path)
+        feat = reader.get_feats(audio_file)
+        feat_f.append(feat.cpu().numpy())
+    logger.info("finished successfully")
+
+def main(audio_dir, save_dir, ckpt_path, layer, max_chunk):
+    reader = HubertFeatureReader(ckpt_path, layer, max_chunk)
+    dump_feature(reader, audio_dir, save_dir)
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--audio_dir", default="./datasets/ASVSpoof2021/ASVspoof2021_LA_eval/flac")
+    parser.add_argument("--save_dir", default="./datasets/ASVSpoof2021/ASVspoof2021_LA_eval/Hubert_L9")
+   
+    parser.add_argument("--ckpt_path", default="./model_zoo/hubert/hubert_base_ls960.pt")
+    parser.add_argument("--layer", type=int, default=9)
+    parser.add_argument("--max_chunk", type=int, default=1600000)
+    args = parser.parse_args()
+    logger.info(args)
+
+    main(**vars(args))

clean/audio/safeear/test.py

@@ -0,0 +1,78 @@
+import importlib
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple
+import argparse
+import pytorch_lightning as pl
+import torch
+import hydra
+
+torch.set_float32_matmul_precision("high")
+
+from pytorch_lightning import Callback, LightningDataModule, LightningModule, Trainer
+from pytorch_lightning.strategies.ddp import DDPStrategy
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+from pytorch_lightning.utilities import rank_zero_only
+
+@rank_zero_only
+def print_only(message: str):
+    """Prints a message only on rank 0."""
+    print(message)
+    
+def train(cfg: DictConfig, args) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    
+    # instantiate datamodule
+    print_only(f"Instantiating datamodule <{cfg.datamodule._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.datamodule)
+    
+    # instantiate decouple model
+    print_only(f"Instantiating decouple model <{cfg.decouple_model._target_}>")
+    decouple_model: torch.nn.Module = hydra.utils.instantiate(cfg.decouple_model)
+    decouple_model.load_state_dict(torch.load(cfg.speechtokenizer_path))
+    # import pdb; pdb.set_trace()
+    
+    # instantiate detect model
+    print(f"Instantiating detect model <{cfg.detect_model._target_}>")
+    detect_model: torch.nn.Module = hydra.utils.instantiate(cfg.detect_model)
+    # import pdb; pdb.set_trace()
+    
+    # instantiate system
+    print_only(f"Instantiating system <{cfg.system._target_}>")
+    system: LightningModule = hydra.utils.instantiate(
+        cfg.system,
+        decouple_model=decouple_model,
+        detect_model=detect_model,
+    )
+    
+    # instantiate trainer
+    print_only(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(
+        cfg.trainer,
+        strategy=DDPStrategy(find_unused_parameters=True),
+    )
+    
+    trainer.test(system, datamodule=datamodule, ckpt_path=args.ckpt_path)
+
+if __name__ == "__main__":
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--conf_dir",
+        default="local/conf.yml",
+        help="Full path to save best validation model",
+    )
+    parser.add_argument(
+        "--ckpt_path",
+        help="Full path to save best validation model",
+    )
+    
+    args = parser.parse_args()
+    cfg = OmegaConf.load(args.conf_dir)
+    
+    os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name), exist_ok=True)
+    # 保存配置到新的文件
+    OmegaConf.save(cfg, os.path.join(cfg.exp.dir, cfg.exp.name, "config.yaml"))
+    
+    train(cfg, args)
+    

clean/audio/safeear/train.py

@@ -0,0 +1,102 @@
+import importlib
+import json
+import os
+import warnings
+warnings.filterwarnings("ignore")
+from typing import Any, Dict, List, Optional, Tuple
+import argparse
+import pytorch_lightning as pl
+import torch
+import hydra
+
+torch.set_float32_matmul_precision("high")
+
+from pytorch_lightning import Callback, LightningDataModule, LightningModule, Trainer
+from pytorch_lightning.strategies.ddp import DDPStrategy
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+from pytorch_lightning.utilities import rank_zero_only
+
+@rank_zero_only
+def print_only(message: str):
+    """Prints a message only on rank 0."""
+    print(message)
+    
+def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    
+    # instantiate datamodule
+    print_only(f"Instantiating datamodule <{cfg.datamodule._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.datamodule)
+    datamodule.setup()
+    
+    # instantiate decouple model
+    print_only(f"Instantiating decouple model <{cfg.decouple_model._target_}>")
+    decouple_model: torch.nn.Module = hydra.utils.instantiate(cfg.decouple_model)
+    decouple_model.load_state_dict(torch.load(cfg.speechtokenizer_path))
+    # import pdb; pdb.set_trace()
+    
+    # instantiate detect model
+    print(f"Instantiating detect model <{cfg.detect_model._target_}>")
+    detect_model: torch.nn.Module = hydra.utils.instantiate(cfg.detect_model)
+    # import pdb; pdb.set_trace()
+    
+    # instantiate system
+    print_only(f"Instantiating system <{cfg.system._target_}>")
+    system: LightningModule = hydra.utils.instantiate(
+        cfg.system,
+        decouple_model=decouple_model,
+        detect_model=detect_model,
+    )
+    # instantiate callbacks
+    callbacks: List[Callback] = []
+    if cfg.get("early_stopping"):
+        print_only(f"Instantiating early_stopping <{cfg.early_stopping._target_}>")
+        callbacks.append(hydra.utils.instantiate(cfg.early_stopping))
+    if cfg.get("checkpoint"):
+        print_only(f"Instantiating checkpoint <{cfg.checkpoint._target_}>")
+        checkpoint: pl.callbacks.ModelCheckpoint = hydra.utils.instantiate(cfg.checkpoint)
+        callbacks.append(checkpoint)
+        
+    # instantiate logger
+    print_only(f"Instantiating logger <{cfg.logger._target_}>")
+    os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name, "logs"), exist_ok=True)
+    logger = hydra.utils.instantiate(cfg.logger)
+    
+    # instantiate trainer
+    print_only(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(
+        cfg.trainer,
+        callbacks=callbacks,
+        logger=logger,
+        strategy=DDPStrategy(find_unused_parameters=True),
+    )
+    
+    trainer.fit(system, datamodule=datamodule)
+    print_only("Training finished!")
+    best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
+    with open(os.path.join(cfg.exp.dir, cfg.exp.name, "best_k_models.json"), "w") as f:
+        json.dump(best_k, f, indent=0)
+
+    import wandb
+    if wandb.run:
+        print_only("Closing wandb!")
+        wandb.finish()
+
+if __name__ == "__main__":
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--conf_dir",
+        default="local/conf.yml",
+        help="Full path to save best validation model",
+    )
+    
+    args = parser.parse_args()
+    cfg = OmegaConf.load(args.conf_dir)
+    
+    os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name), exist_ok=True)
+    # 保存配置到新的文件
+    OmegaConf.save(cfg, os.path.join(cfg.exp.dir, cfg.exp.name, "config.yaml"))
+    
+    train(cfg)
+    

clean/audio/shiftyspeech/.env

@@ -0,0 +1,2 @@
+WANDB_API_KEY="<wandb api-key>"
+WANDB_PROJECT_NAME="SSL-AASIST"

clean/audio/shiftyspeech/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Hemlata
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

clean/audio/shiftyspeech/RawBoost.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import copy
+
+import numpy as np
+from scipy import signal
+
+"""
+   Hemlata Tak, Madhu Kamble, Jose Patino, Massimiliano Todisco, Nicholas Evans.
+   RawBoost: A Raw Data Boosting and Augmentation Method applied to Automatic Speaker Verification Anti-Spoofing.
+   In Proc. ICASSP 2022, pp:6382--6386.
+"""
+
+
+def randRange(x1, x2, integer):
+    y = np.random.uniform(low=x1, high=x2, size=(1,))
+    if integer:
+        y = int(y)
+    return y
+
+
+def normWav(x, always):
+    if always:
+        x = x / np.amax(abs(x))
+    elif np.amax(abs(x)) > 1:
+        x = x / np.amax(abs(x))
+    return x
+
+
+def genNotchCoeffs(
+    nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs
+):
+    b = 1
+    for i in range(0, nBands):
+        fc = randRange(minF, maxF, 0)
+        bw = randRange(minBW, maxBW, 0)
+        c = randRange(minCoeff, maxCoeff, 1)
+
+        if c / 2 == int(c / 2):
+            c = c + 1
+        f1 = fc - bw / 2
+        f2 = fc + bw / 2
+        if f1 <= 0:
+            f1 = 1 / 1000
+        if f2 >= fs / 2:
+            f2 = fs / 2 - 1 / 1000
+        b = np.convolve(
+            signal.firwin(c, [float(f1), float(f2)], window="hamming", fs=fs), b
+        )
+
+    G = randRange(minG, maxG, 0)
+    _, h = signal.freqz(b, 1, fs=fs)
+    b = pow(10, G / 20) * b / np.amax(abs(h))
+    return b
+
+
+def filterFIR(x, b):
+    N = b.shape[0] + 1
+    xpad = np.pad(x, (0, N), "constant")
+    y = signal.lfilter(b, 1, xpad)
+    y = y[int(N / 2) : int(y.shape[0] - N / 2)]
+    return y
+
+
+# Linear and non-linear convolutive noise
+def LnL_convolutive_noise(
+    x,
+    N_f,
+    nBands,
+    minF,
+    maxF,
+    minBW,
+    maxBW,
+    minCoeff,
+    maxCoeff,
+    minG,
+    maxG,
+    minBiasLinNonLin,
+    maxBiasLinNonLin,
+    fs,
+):
+    y = [0] * x.shape[0]
+    for i in range(0, N_f):
+        if i == 1:
+            minG = minG - minBiasLinNonLin
+            maxG = maxG - maxBiasLinNonLin
+        b = genNotchCoeffs(
+            nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs
+        )
+        y = y + filterFIR(np.power(x, (i + 1)), b)
+    y = y - np.mean(y)
+    y = normWav(y, 0)
+    return y
+
+
+# Impulsive signal dependent noise
+def ISD_additive_noise(x, P, g_sd):
+    beta = randRange(0, P, 0)
+
+    y = copy.deepcopy(x)
+    x_len = x.shape[0]
+    n = int(x_len * (beta / 100))
+    p = np.random.permutation(x_len)[:n]
+    f_r = np.multiply(
+        ((2 * np.random.rand(p.shape[0])) - 1), ((2 * np.random.rand(p.shape[0])) - 1)
+    )
+    r = g_sd * x[p] * f_r
+    y[p] = x[p] + r
+    y = normWav(y, 0)
+    return y
+
+
+# Stationary signal independent noise
+
+
+def SSI_additive_noise(
+    x,
+    SNRmin,
+    SNRmax,
+    nBands,
+    minF,
+    maxF,
+    minBW,
+    maxBW,
+    minCoeff,
+    maxCoeff,
+    minG,
+    maxG,
+    fs,
+):
+    noise = np.random.normal(0, 1, x.shape[0])
+    b = genNotchCoeffs(
+        nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs
+    )
+    noise = filterFIR(noise, b)
+    noise = normWav(noise, 1)
+    SNR = randRange(SNRmin, SNRmax, 0)
+    noise = (
+        noise / np.linalg.norm(noise, 2) * np.linalg.norm(x, 2) / 10.0 ** (0.05 * SNR)
+    )
+    x = x + noise
+    return x

clean/audio/shiftyspeech/SOURCE.md

@@ -0,0 +1,17 @@
+# Source: audio/shiftyspeech
+
+| Field | Value |
+|---|---|
+| Upstream | **UNVERIFIED** -- provenance was lost when this code was vendored |
+| Paper | not recorded |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/audio__shiftyspeech.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.

clean/audio/shiftyspeech/Simplified_CM_solution.py

@@ -0,0 +1,227 @@
+import math
+from collections import OrderedDict
+
+import fairseq
+import numpy as np
+import scipy.io as sio
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+from torch.utils import data
+
+___author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+# from losses_anti_spoofing import AMSoftmax
+
+############################
+## FOR fine-tuning SSL MODEL
+############################
+
+
+class SSLModel(nn.Module):
+    def __init__(self, device):
+        super(SSLModel, self).__init__()
+
+        cp_path = "/change_to_path_to_pre_trained_model_XLR_300M/xlsr2_300m.pt"
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        self.model = model[0]
+        self.device = device
+        self.out_dim = 1024
+        return
+
+    def extract_feat(self, input_data):
+
+        # put the model to GPU if it not there
+        if (
+            next(self.model.parameters()).device != input_data.device
+            or next(self.model.parameters()).dtype != input_data.dtype
+        ):
+            self.model.to(input_data.device, dtype=input_data.dtype)
+            self.model.train()
+
+        if True:
+            # input should be in shape (batch, length)
+            if input_data.ndim == 3:
+                input_tmp = input_data[:, :, 0]
+            else:
+                input_tmp = input_data
+
+            # [batch, length, dim]
+            emb = self.model(input_tmp, mask=False, features_only=True)["x"]
+        return emb
+
+
+# ---------Graph attention simple back-end------------------------#
+""" 
+    Hemlata Tak, Jee-weon Jung, Jose Patino, Madhu Kamble, Massimiliano Todisco, Nicholas Evans.
+    End-to-end spectro-temporal graph attention networks for speaker verification anti-spoofing and speech deepfake detection.
+    In Proc. Automatic Speaker Verification and Spoofing Countermeasures Challenge 2021 Interspeech 2021 satellite workshop.
+"""
+
+
+class GraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, **kwargs):
+        super(GraphAttentionLayer, self).__init__()
+
+        # attention map
+        self.att_proj = nn.Linear(in_dim, out_dim)
+        self.att_weight = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = nn.Linear(in_dim, out_dim)
+        self.proj_without_att = nn.Linear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        self.act = nn.SELU(inplace=True)
+
+    def forward(self, x):
+        """
+        x   :(#bs, #node, #dim)
+        """
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+
+        return x
+
+    def _pairwise_mul_nodes(self, x):
+        """
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        """
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map(self, x):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = self._pairwise_mul_nodes(x)
+        att_map = torch.tanh(
+            self.att_proj(att_map)
+        )  # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.matmul(att_map, self.att_weight)  # size: (#bs, #node, #node, 1)
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+
+class GraphPool(nn.Module):
+    def __init__(self, k: float, in_dim: int, p):
+        super().__init__()
+        self.k = k
+        self.sigmoid = nn.Sigmoid()
+        self.proj = nn.Linear(in_dim, 1)
+        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
+        self.in_dim = in_dim
+
+    def forward(self, h):
+        Z = self.drop(h)
+        weights = self.proj(Z)
+        scores = self.sigmoid(weights)
+        new_h = self.top_k_graph(scores, h, self.k)
+
+        return new_h
+
+    def top_k_graph(self, scores, h, k):
+        """
+        args
+        =====
+        scores: attention-based weights (#bs, #node, 1)
+        h: graph data (#bs, #node, #dim)
+        k: ratio of remaining nodes, (float)
+        returns
+        =====
+        h: graph pool applied data (#bs, #node', #dim)
+        """
+        _, n_nodes, n_feat = h.size()
+        n_nodes = max(int(n_nodes * k), 1)
+        _, idx = torch.topk(scores, n_nodes, dim=1)
+        idx = idx.expand(-1, -1, n_feat)
+
+        h = h * scores
+        h = torch.gather(h, 1, idx)
+
+        return h
+
+
+class Model(nn.Module):
+    def __init__(self, d_args, device):
+        super(Model, self).__init__()
+
+        # SSL model
+        self.device = device
+        self.ssl_model = SSLModel(self.device)
+        self.LL = nn.Linear(self.ssl_model.out_dim, 128)
+        self.first_bn = nn.BatchNorm1d(num_features=128)
+        self.selu = nn.SELU(inplace=True)
+
+        # graph module layer
+        self.GAT_layer = GraphAttentionLayer(128, 64)
+        self.proj = nn.Linear(64, 1)
+        self.pool = GraphPool(0.8, 64, 0.3)
+
+        # classifier head
+        self.proj_node = nn.Linear(53, 2)
+
+    def forward(self, x_inp, Freq_aug=False):
+        # SSL wav2vec 2.0 model
+        x_ssl_feat = self.ssl_model.extract_feat(x_inp.squeeze(-1))
+        x_SSL = self.LL(x_ssl_feat)  # (bs,frame_number,feat_out_dim)
+        x_SSL = x_SSL.transpose(1, 2)  # (bs,feat_out_dim,frame_number)
+
+        x = F.max_pool1d(x_SSL, (3))
+        x = self.first_bn(x)
+        x = self.selu(x)
+
+        x = self.GAT_layer(x.transpose(1, 2))
+        x = self.pool(x)
+        x = self.proj(x).flatten(1)
+        output = self.proj_node(x)
+        return output

clean/audio/shiftyspeech/data_utils.py

@@ -0,0 +1,292 @@
+import os
+import random
+from random import randrange
+
+import librosa
+import numpy as np
+import torch
+import torch.nn as nn
+from RawBoost import (
+    ISD_additive_noise,
+    LnL_convolutive_noise,
+    SSI_additive_noise,
+    normWav,
+)
+from torch import Tensor
+from torch.utils.data import Dataset
+
+__author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+
+def genSpoof_list(dir_meta, is_train=False, is_eval=False):
+    d_meta = {}
+    file_list = []
+    with open(dir_meta, "r") as f:
+        l_meta = f.readlines()
+
+    if is_train:
+        for line in l_meta:
+            key, label = line.strip().split()
+            file_list.append(key)
+            d_meta[key] = 1 if label == "bonafide" else 0
+        return d_meta, file_list
+
+    elif is_eval:
+        for line in l_meta:
+            key, _ = line.strip().split(" ")
+            file_list.append(key)
+        return file_list
+    else:
+        for line in l_meta:
+            key, label = line.strip().split()
+            file_list.append(key)
+            d_meta[key] = 1 if label == "bonafide" else 0
+        return d_meta, file_list
+
+
+def pad(x, max_len=64600):
+    x_len = x.shape[0]
+    if x_len >= max_len:
+        return x[:max_len]
+    # need to pad
+    num_repeats = int(max_len / x_len) + 1
+    padded_x = np.tile(x, (1, num_repeats))[:, :max_len][0]
+    return padded_x
+
+
+class Dataset_ASVspoof2019_train(Dataset):
+    def __init__(self, args, metafile, algo):
+        """self.list_IDs : list of strings (each string: utt key),
+        self.labels: dictionary (key: utt key, value: label integer)"""
+
+        self.uttpath_labels = []
+        with open(metafile, "r") as f:
+            for line in f:
+                items = line.strip().split()
+                lb = 1 if items[-1] == "bonafide" else 0
+                self.uttpath_labels.append((items[0], lb))
+
+        self.algo = algo
+        self.args = args
+        self.cut = 64600  # take ~4 sec audio (64600 samples)
+
+    def __len__(self):
+        return len(self.uttpath_labels)
+
+    def __getitem__(self, index):
+        path, target = self.uttpath_labels[index]
+        X, fs = librosa.load(path, sr=16000)
+        Y = process_Rawboost_feature(X, fs, self.args, self.algo)
+        X_pad = pad(Y, self.cut)
+        x_inp = Tensor(X_pad)
+        return x_inp, target
+
+
+class Dataset_ASVspoof2021_eval(Dataset):
+    def __init__(self, list_IDs):
+        """self.list_IDs	: list of strings (each string: utt key),"""
+
+        self.list_IDs = list_IDs
+        self.cut = 64600  # take ~4 sec audio (64600 samples)
+
+    def __len__(self):
+        return len(self.list_IDs)
+
+    def __getitem__(self, index):
+        utt_id = self.list_IDs[index]
+        X, fs = librosa.load(utt_id, sr=16000)
+        X_pad = pad(X, self.cut)
+        x_inp = Tensor(X_pad)
+        return x_inp, utt_id
+
+
+# --------------RawBoost data augmentation algorithms---------------------------##
+def process_Rawboost_feature(feature, sr, args, algo):
+
+    # Data process by Convolutive noise (1st algo)
+    if algo == 1:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+
+    # Data process by Impulsive noise (2nd algo)
+    elif algo == 2:
+
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+
+    # Data process by coloured additive noise (3rd algo)
+    elif algo == 3:
+
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by all 3 algo. together in series (1+2+3)
+    elif algo == 4:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by 1st two algo. together in series (1+2)
+    elif algo == 5:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+
+    # Data process by 1st and 3rd algo. together in series (1+3)
+    elif algo == 6:
+
+        feature = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by 2nd and 3rd algo. together in series (2+3)
+    elif algo == 7:
+
+        feature = ISD_additive_noise(feature, args.P, args.g_sd)
+        feature = SSI_additive_noise(
+            feature,
+            args.SNRmin,
+            args.SNRmax,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            sr,
+        )
+
+    # Data process by 1st two algo. together in Parallel (1||2)
+    elif algo == 8:
+
+        feature1 = LnL_convolutive_noise(
+            feature,
+            args.N_f,
+            args.nBands,
+            args.minF,
+            args.maxF,
+            args.minBW,
+            args.maxBW,
+            args.minCoeff,
+            args.maxCoeff,
+            args.minG,
+            args.maxG,
+            args.minBiasLinNonLin,
+            args.maxBiasLinNonLin,
+            sr,
+        )
+        feature2 = ISD_additive_noise(feature, args.P, args.g_sd)
+
+        feature_para = feature1 + feature2
+        feature = normWav(feature_para, 0)  # normalized resultant waveform
+
+    # original data without Rawboost processing
+    else:
+
+        feature = feature
+
+    return feature

clean/audio/shiftyspeech/model.py

@@ -0,0 +1,603 @@
+import random
+from typing import Union
+
+import fairseq
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+___author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+############################
+## FOR fine-tuned SSL MODEL
+############################
+
+
+class SSLModel(nn.Module):
+    def __init__(self, device):
+        super(SSLModel, self).__init__()
+
+        cp_path = "models/xlsr2_300m.pt"  # Change the pre-trained XLSR model path.
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        self.model = model[0]
+        self.device = device
+        self.out_dim = 1024
+        return
+
+    def extract_feat(self, input_data):
+
+        # put the model to GPU if it not there
+        if (
+            next(self.model.parameters()).device != input_data.device
+            or next(self.model.parameters()).dtype != input_data.dtype
+        ):
+            self.model.to(input_data.device, dtype=input_data.dtype)
+            self.model.train()
+
+        if True:
+            # input should be in shape (batch, length)
+            if input_data.ndim == 3:
+                input_tmp = input_data[:, :, 0]
+            else:
+                input_tmp = input_data
+
+            # [batch, length, dim]
+            emb = self.model(input_tmp, mask=False, features_only=True)["x"]
+        return emb
+
+
+# ---------AASIST back-end------------------------#
+""" Jee-weon Jung, Hee-Soo Heo, Hemlata Tak, Hye-jin Shim, Joon Son Chung, Bong-Jin Lee, Ha-Jin Yu and Nicholas Evans. 
+    AASIST: Audio Anti-Spoofing Using Integrated Spectro-Temporal Graph Attention Networks. 
+    In Proc. ICASSP 2022, pp: 6367--6371."""
+
+
+class GraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, **kwargs):
+        super().__init__()
+
+        # attention map
+        self.att_proj = nn.Linear(in_dim, out_dim)
+        self.att_weight = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = nn.Linear(in_dim, out_dim)
+        self.proj_without_att = nn.Linear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        # activate
+        self.act = nn.SELU(inplace=True)
+
+        # temperature
+        self.temp = 1.0
+        if "temperature" in kwargs:
+            self.temp = kwargs["temperature"]
+
+    def forward(self, x):
+        """
+        x   :(#bs, #node, #dim)
+        """
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+        return x
+
+    def _pairwise_mul_nodes(self, x):
+        """
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        """
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map(self, x):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = self._pairwise_mul_nodes(x)
+        # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.tanh(self.att_proj(att_map))
+        # size: (#bs, #node, #node, 1)
+        att_map = torch.matmul(att_map, self.att_weight)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+
+class HtrgGraphAttentionLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, **kwargs):
+        super().__init__()
+
+        self.proj_type1 = nn.Linear(in_dim, in_dim)
+        self.proj_type2 = nn.Linear(in_dim, in_dim)
+
+        # attention map
+        self.att_proj = nn.Linear(in_dim, out_dim)
+        self.att_projM = nn.Linear(in_dim, out_dim)
+
+        self.att_weight11 = self._init_new_params(out_dim, 1)
+        self.att_weight22 = self._init_new_params(out_dim, 1)
+        self.att_weight12 = self._init_new_params(out_dim, 1)
+        self.att_weightM = self._init_new_params(out_dim, 1)
+
+        # project
+        self.proj_with_att = nn.Linear(in_dim, out_dim)
+        self.proj_without_att = nn.Linear(in_dim, out_dim)
+
+        self.proj_with_attM = nn.Linear(in_dim, out_dim)
+        self.proj_without_attM = nn.Linear(in_dim, out_dim)
+
+        # batch norm
+        self.bn = nn.BatchNorm1d(out_dim)
+
+        # dropout for inputs
+        self.input_drop = nn.Dropout(p=0.2)
+
+        # activate
+        self.act = nn.SELU(inplace=True)
+
+        # temperature
+        self.temp = 1.0
+        if "temperature" in kwargs:
+            self.temp = kwargs["temperature"]
+
+    def forward(self, x1, x2, master=None):
+        """
+        x1  :(#bs, #node, #dim)
+        x2  :(#bs, #node, #dim)
+        """
+
+        num_type1 = x1.size(1)
+        num_type2 = x2.size(1)
+
+        x1 = self.proj_type1(x1)
+
+        x2 = self.proj_type2(x2)
+
+        x = torch.cat([x1, x2], dim=1)
+
+        if master is None:
+            master = torch.mean(x, dim=1, keepdim=True)
+
+        # apply input dropout
+        x = self.input_drop(x)
+
+        # derive attention map
+        att_map = self._derive_att_map(x, num_type1, num_type2)
+
+        # directional edge for master node
+        master = self._update_master(x, master)
+
+        # projection
+        x = self._project(x, att_map)
+
+        # apply batch norm
+        x = self._apply_BN(x)
+        x = self.act(x)
+
+        x1 = x.narrow(1, 0, num_type1)
+
+        x2 = x.narrow(1, num_type1, num_type2)
+
+        return x1, x2, master
+
+    def _update_master(self, x, master):
+
+        att_map = self._derive_att_map_master(x, master)
+        master = self._project_master(x, master, att_map)
+
+        return master
+
+    def _pairwise_mul_nodes(self, x):
+        """
+        Calculates pairwise multiplication of nodes.
+        - for attention map
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, #dim)
+        """
+
+        nb_nodes = x.size(1)
+        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
+        x_mirror = x.transpose(1, 2)
+
+        return x * x_mirror
+
+    def _derive_att_map_master(self, x, master):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = x * master
+        att_map = torch.tanh(self.att_projM(att_map))
+
+        att_map = torch.matmul(att_map, self.att_weightM)
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _derive_att_map(self, x, num_type1, num_type2):
+        """
+        x           :(#bs, #node, #dim)
+        out_shape   :(#bs, #node, #node, 1)
+        """
+        att_map = self._pairwise_mul_nodes(x)
+        # size: (#bs, #node, #node, #dim_out)
+        att_map = torch.tanh(self.att_proj(att_map))
+        # size: (#bs, #node, #node, 1)
+
+        att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1)
+
+        att_board[:, :num_type1, :num_type1, :] = torch.matmul(
+            att_map[:, :num_type1, :num_type1, :], self.att_weight11
+        )
+        att_board[:, num_type1:, num_type1:, :] = torch.matmul(
+            att_map[:, num_type1:, num_type1:, :], self.att_weight22
+        )
+        att_board[:, :num_type1, num_type1:, :] = torch.matmul(
+            att_map[:, :num_type1, num_type1:, :], self.att_weight12
+        )
+        att_board[:, num_type1:, :num_type1, :] = torch.matmul(
+            att_map[:, num_type1:, :num_type1, :], self.att_weight12
+        )
+
+        att_map = att_board
+
+        # apply temperature
+        att_map = att_map / self.temp
+
+        att_map = F.softmax(att_map, dim=-2)
+
+        return att_map
+
+    def _project(self, x, att_map):
+        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
+        x2 = self.proj_without_att(x)
+
+        return x1 + x2
+
+    def _project_master(self, x, master, att_map):
+
+        x1 = self.proj_with_attM(torch.matmul(att_map.squeeze(-1).unsqueeze(1), x))
+        x2 = self.proj_without_attM(master)
+
+        return x1 + x2
+
+    def _apply_BN(self, x):
+        org_size = x.size()
+        x = x.view(-1, org_size[-1])
+        x = self.bn(x)
+        x = x.view(org_size)
+
+        return x
+
+    def _init_new_params(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+
+class GraphPool(nn.Module):
+    def __init__(self, k: float, in_dim: int, p: Union[float, int]):
+        super().__init__()
+        self.k = k
+        self.sigmoid = nn.Sigmoid()
+        self.proj = nn.Linear(in_dim, 1)
+        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
+        self.in_dim = in_dim
+
+    def forward(self, h):
+        Z = self.drop(h)
+        weights = self.proj(Z)
+        scores = self.sigmoid(weights)
+        new_h = self.top_k_graph(scores, h, self.k)
+
+        return new_h
+
+    def top_k_graph(self, scores, h, k):
+        """
+        args
+        =====
+        scores: attention-based weights (#bs, #node, 1)
+        h: graph data (#bs, #node, #dim)
+        k: ratio of remaining nodes, (float)
+        returns
+        =====
+        h: graph pool applied data (#bs, #node', #dim)
+        """
+        _, n_nodes, n_feat = h.size()
+        n_nodes = max(int(n_nodes * k), 1)
+        _, idx = torch.topk(scores, n_nodes, dim=1)
+        idx = idx.expand(-1, -1, n_feat)
+
+        h = h * scores
+        h = torch.gather(h, 1, idx)
+
+        return h
+
+
+class Residual_block(nn.Module):
+    def __init__(self, nb_filts, first=False):
+        super().__init__()
+        self.first = first
+
+        if not self.first:
+            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
+        self.conv1 = nn.Conv2d(
+            in_channels=nb_filts[0],
+            out_channels=nb_filts[1],
+            kernel_size=(2, 3),
+            padding=(1, 1),
+            stride=1,
+        )
+        self.selu = nn.SELU(inplace=True)
+
+        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
+        self.conv2 = nn.Conv2d(
+            in_channels=nb_filts[1],
+            out_channels=nb_filts[1],
+            kernel_size=(2, 3),
+            padding=(0, 1),
+            stride=1,
+        )
+
+        if nb_filts[0] != nb_filts[1]:
+            self.downsample = True
+            self.conv_downsample = nn.Conv2d(
+                in_channels=nb_filts[0],
+                out_channels=nb_filts[1],
+                padding=(0, 1),
+                kernel_size=(1, 3),
+                stride=1,
+            )
+
+        else:
+            self.downsample = False
+
+    def forward(self, x):
+        identity = x
+        if not self.first:
+            out = self.bn1(x)
+            out = self.selu(out)
+        else:
+            out = x
+
+        out = self.conv1(x)
+
+        out = self.bn2(out)
+        out = self.selu(out)
+
+        out = self.conv2(out)
+
+        if self.downsample:
+            identity = self.conv_downsample(identity)
+
+        out += identity
+
+        return out
+
+
+class Model(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self.device = device
+
+        # AASIST parameters
+        filts = [128, [1, 32], [32, 32], [32, 64], [64, 64]]
+        gat_dims = [64, 32]
+        pool_ratios = [0.5, 0.5, 0.5, 0.5]
+        temperatures = [2.0, 2.0, 100.0, 100.0]
+
+        ####
+        # create network wav2vec 2.0
+        ####
+        self.ssl_model = SSLModel(self.device)
+        self.LL = nn.Linear(self.ssl_model.out_dim, 128)
+
+        self.first_bn = nn.BatchNorm2d(num_features=1)
+        self.first_bn1 = nn.BatchNorm2d(num_features=64)
+        self.drop = nn.Dropout(0.5, inplace=True)
+        self.drop_way = nn.Dropout(0.2, inplace=True)
+        self.selu = nn.SELU(inplace=True)
+
+        # RawNet2 encoder
+        self.encoder = nn.Sequential(
+            nn.Sequential(Residual_block(nb_filts=filts[1], first=True)),
+            nn.Sequential(Residual_block(nb_filts=filts[2])),
+            nn.Sequential(Residual_block(nb_filts=filts[3])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+            nn.Sequential(Residual_block(nb_filts=filts[4])),
+        )
+
+        self.attention = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=(1, 1)),
+            nn.SELU(inplace=True),
+            nn.BatchNorm2d(128),
+            nn.Conv2d(128, 64, kernel_size=(1, 1)),
+        )
+        # position encoding
+        self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1]))
+
+        self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+        self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
+
+        # Graph module
+        self.GAT_layer_S = GraphAttentionLayer(
+            filts[-1][-1], gat_dims[0], temperature=temperatures[0]
+        )
+        self.GAT_layer_T = GraphAttentionLayer(
+            filts[-1][-1], gat_dims[0], temperature=temperatures[1]
+        )
+        # HS-GAL layer
+        self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer(
+            gat_dims[0], gat_dims[1], temperature=temperatures[2]
+        )
+        self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer(
+            gat_dims[1], gat_dims[1], temperature=temperatures[2]
+        )
+        self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer(
+            gat_dims[0], gat_dims[1], temperature=temperatures[2]
+        )
+        self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer(
+            gat_dims[1], gat_dims[1], temperature=temperatures[2]
+        )
+
+        # Graph pooling layers
+        self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3)
+        self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3)
+        self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+        self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+
+        self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+        self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
+
+        self.out_layer = nn.Linear(5 * gat_dims[1], 2)
+
+    def forward(self, x):
+        # -------pre-trained Wav2vec model fine tunning ------------------------##
+        x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
+        x = self.LL(x_ssl_feat)  # (bs,frame_number,feat_out_dim)
+
+        # post-processing on front-end features
+        x = x.transpose(1, 2)  # (bs,feat_out_dim,frame_number)
+        x = x.unsqueeze(dim=1)  # add channel
+        x = F.max_pool2d(x, (3, 3))
+        x = self.first_bn(x)
+        x = self.selu(x)
+
+        # RawNet2-based encoder
+        x = self.encoder(x)
+        x = self.first_bn1(x)
+        x = self.selu(x)
+
+        w = self.attention(x)
+
+        # ------------SA for spectral feature-------------#
+        w1 = F.softmax(w, dim=-1)
+        m = torch.sum(x * w1, dim=-1)
+        e_S = m.transpose(1, 2) + self.pos_S
+
+        # graph module layer
+        gat_S = self.GAT_layer_S(e_S)
+        out_S = self.pool_S(gat_S)  # (#bs, #node, #dim)
+
+        # ------------SA for temporal feature-------------#
+        w2 = F.softmax(w, dim=-2)
+        m1 = torch.sum(x * w2, dim=-2)
+
+        e_T = m1.transpose(1, 2)
+
+        # graph module layer
+        gat_T = self.GAT_layer_T(e_T)
+        out_T = self.pool_T(gat_T)
+
+        # learnable master node
+        master1 = self.master1.expand(x.size(0), -1, -1)
+        master2 = self.master2.expand(x.size(0), -1, -1)
+
+        # inference 1
+        out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11(
+            out_T, out_S, master=self.master1
+        )
+
+        out_S1 = self.pool_hS1(out_S1)
+        out_T1 = self.pool_hT1(out_T1)
+
+        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12(
+            out_T1, out_S1, master=master1
+        )
+        out_T1 = out_T1 + out_T_aug
+        out_S1 = out_S1 + out_S_aug
+        master1 = master1 + master_aug
+
+        # inference 2
+        out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21(
+            out_T, out_S, master=self.master2
+        )
+        out_S2 = self.pool_hS2(out_S2)
+        out_T2 = self.pool_hT2(out_T2)
+
+        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22(
+            out_T2, out_S2, master=master2
+        )
+        out_T2 = out_T2 + out_T_aug
+        out_S2 = out_S2 + out_S_aug
+        master2 = master2 + master_aug
+
+        out_T1 = self.drop_way(out_T1)
+        out_T2 = self.drop_way(out_T2)
+        out_S1 = self.drop_way(out_S1)
+        out_S2 = self.drop_way(out_S2)
+        master1 = self.drop_way(master1)
+        master2 = self.drop_way(master2)
+
+        out_T = torch.max(out_T1, out_T2)
+        out_S = torch.max(out_S1, out_S2)
+        master = torch.max(master1, master2)
+
+        # Readout operation
+        T_max, _ = torch.max(torch.abs(out_T), dim=1)
+        T_avg = torch.mean(out_T, dim=1)
+
+        S_max, _ = torch.max(torch.abs(out_S), dim=1)
+        S_avg = torch.mean(out_S, dim=1)
+
+        last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1)
+
+        last_hidden = self.drop(last_hidden)
+        output = self.out_layer(last_hidden)
+
+        return output

clean/audio/shiftyspeech/startup_config.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+"""
+startup_config
+
+Startup configuration utilities
+
+"""
+
+from __future__ import absolute_import
+
+import importlib
+import os
+import random
+import sys
+
+import numpy as np
+import torch
+
+__author__ = "Xin Wang"
+__email__ = "wangxin@nii.ac.jp"
+__copyright__ = "Copyright 2020, Xin Wang"
+
+
+def set_random_seed(random_seed, args=None):
+    """set_random_seed(random_seed, args=None)
+
+    Set the random_seed for numpy, python, and cudnn
+
+    input
+    -----
+      random_seed: integer random seed
+      args: argue parser
+    """
+
+    # initialization
+    torch.manual_seed(random_seed)
+    random.seed(random_seed)
+    np.random.seed(random_seed)
+    os.environ["PYTHONHASHSEED"] = str(random_seed)
+
+    # For torch.backends.cudnn.deterministic
+    # Note: this default configuration may result in RuntimeError
+    # see https://pytorch.org/docs/stable/notes/randomness.html
+    if args is None:
+        cudnn_deterministic = True
+        cudnn_benchmark = False
+    else:
+        cudnn_deterministic = args.cudnn_deterministic_toggle
+        cudnn_benchmark = args.cudnn_benchmark_toggle
+
+        if not cudnn_deterministic:
+            print("cudnn_deterministic set to False")
+        if cudnn_benchmark:
+            print("cudnn_benchmark set to True")
+
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(random_seed)
+        torch.backends.cudnn.deterministic = cudnn_deterministic
+        torch.backends.cudnn.benchmark = cudnn_benchmark
+    return

clean/audio/shiftyspeech/train.py

@@ -0,0 +1,446 @@
+import argparse
+import os
+import sys
+
+import librosa
+import numpy as np
+import torch
+import wandb
+import yaml
+from data_utils import (
+    Dataset_ASVspoof2019_train,
+    Dataset_ASVspoof2021_eval,
+    genSpoof_list,
+    pad,
+    process_Rawboost_feature,
+)
+from dotenv import load_dotenv
+from model import Model
+from sklearn.metrics import roc_auc_score
+from startup_config import set_random_seed
+from tensorboardX import SummaryWriter
+from torch import Tensor, nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+__author__ = "Hemlata Tak"
+__email__ = "tak@eurecom.fr"
+
+
+def compute_det_curve(target_scores, nontarget_scores):
+
+    n_scores = target_scores.size + nontarget_scores.size
+    all_scores = np.concatenate((target_scores, nontarget_scores))
+    labels = np.concatenate(
+        (np.ones(target_scores.size), np.zeros(nontarget_scores.size))
+    )
+
+    indices = np.argsort(all_scores, kind="mergesort")
+    labels = labels[indices]
+    tar_trial_sums = np.cumsum(labels)
+    nontarget_trial_sums = nontarget_scores.size - (
+        np.arange(1, n_scores + 1) - tar_trial_sums
+    )
+
+    frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size))
+    far = np.concatenate(
+        (np.atleast_1d(1), nontarget_trial_sums / nontarget_scores.size)
+    )
+    # Thresholds are the sorted scores
+    thresholds = np.concatenate(
+        (np.atleast_1d(all_scores[indices[0]] - 0.001), all_scores[indices])
+    )
+
+    return frr, far, thresholds
+
+
+def compute_eer(target_scores, nontarget_scores):
+    """Returns equal error rate (EER) and the corresponding threshold."""
+    frr, far, thresholds = compute_det_curve(target_scores, nontarget_scores)
+    abs_diffs = np.abs(frr - far)
+    min_index = np.argmin(abs_diffs)
+    eer = np.mean((frr[min_index], far[min_index]))
+    return eer, thresholds[min_index], frr, far
+
+
+def calculate_tDCF_EER(cm_scores_file, output_file, printout=True):
+    # Load CM scores
+    cm_data = np.genfromtxt(cm_scores_file, dtype=str)
+    cm_utt_id = cm_data[:, 0]
+    cm_keys = cm_data[:, 1]
+    cm_scores = cm_data[:, 2].astype(float)
+    # Extract bona fide (real human) and spoof scores from the CM scores
+    bona_cm = cm_scores[cm_keys == "bonafide"]
+    spoof_cm = cm_scores[cm_keys == "spoof"]
+    all_scores = np.concatenate([bona_cm, spoof_cm])
+    all_true_labels = np.concatenate([np.ones_like(bona_cm), np.zeros_like(spoof_cm)])
+
+    auc = roc_auc_score(all_true_labels, all_scores, max_fpr=0.05)
+    eer_cm, eer_threshold, frr, far = compute_eer(bona_cm, spoof_cm)
+
+    if printout:
+        with open(output_file, "w") as f_res:
+            f_res.write("\nCM SYSTEM\n")
+            f_res.write(
+                "\tEER\t\t= {:8.9f} % "
+                "(Equal error rate for countermeasure)\n".format(eer_cm * 100)
+            )
+            f_res.write("\t pAUC with max fpr - 0.05 is :{}".format(auc))
+
+
+def evaluate_accuracy(dev_loader, model, device, args):
+    val_loss = 0.0
+    num_total = 0.0
+    algo = args.algo
+    cut = 64600
+    model.eval()
+
+    weight = torch.FloatTensor([0.1, 0.9]).to(device)
+    criterion = nn.CrossEntropyLoss(weight=weight)
+    progress_bar = tqdm(dev_loader, desc=f"Epoch {epoch+1}/{args.num_epochs}")
+    for current_step, (batch_pths, batch_y) in enumerate(progress_bar):
+        batch_x = batch_pths
+        batch_size = batch_x.size(0)
+        num_total += batch_size
+        batch_x = batch_x.to(device)
+        batch_y = batch_y.view(-1).type(torch.int64).to(device)
+        batch_out = model(batch_x)
+
+        batch_loss = criterion(batch_out, batch_y)
+        val_loss += batch_loss.item() * batch_size
+
+    val_loss /= num_total
+
+    return val_loss
+
+
+def produce_evaluation_file(dataset, model, device, save_path, trial_path):
+    data_loader = DataLoader(dataset, batch_size=10, shuffle=False, drop_last=False)
+    num_correct = 0.0
+    num_total = 0.0
+    model.eval()
+    with open(trial_path, "r") as f_trl:
+        trial_lines = f_trl.readlines()
+
+    fname_list = []
+    score_list = []
+
+    for batch_x, utt_id in data_loader:
+
+        batch_size = batch_x.size(0)
+        batch_x = batch_x.to(device)
+
+        batch_out = model(batch_x)
+
+        batch_score = (batch_out[:, 1]).data.cpu().numpy().ravel()
+        # add outputs
+        fname_list.extend(utt_id)
+        score_list.extend(batch_score.tolist())
+    assert len(trial_lines) == len(fname_list) == len(score_list)
+
+    with open(save_path, "a+") as fh:
+        for fname, cm, trl in zip(fname_list, score_list, trial_lines):
+            utt_id, key = trl.strip().split(" ")
+            assert fname == utt_id
+            fh.write("{} {} {}\n".format(fname, key, cm))
+    fh.close()
+    print("Scores saved to {}".format(save_path))
+
+
+def train_epoch(train_loader, model, lr, optim, device, args):
+    running_loss = 0
+
+    num_total = 0.0
+    algo = args.algo
+    model.train()
+    cut = 64600
+    # set objective (Loss) functions
+    weight = torch.FloatTensor([0.1, 0.9]).to(device)
+    criterion = nn.CrossEntropyLoss(weight=weight)
+    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.num_epochs}")
+    for current_step, (batch_pths, batch_y) in enumerate(progress_bar):
+        batch_x = batch_pths
+        batch_size = batch_x.size(0)
+        num_total += batch_size
+
+        batch_x = batch_x.to(device)
+        batch_y = batch_y.view(-1).type(torch.int64).to(device)
+        batch_out = model(batch_x)
+
+        batch_loss = criterion(batch_out, batch_y)
+
+        running_loss += batch_loss.item() * batch_size
+
+        optimizer.zero_grad()
+        batch_loss.backward()
+        optimizer.step()
+
+    running_loss /= num_total
+
+    return running_loss
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="SSL-AASIST baseline system")
+
+    # Hyperparameters
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--num_epochs", type=int, default=100)
+    parser.add_argument("--lr", type=float, default=0.000001)
+    parser.add_argument("--weight_decay", type=float, default=0.0001)
+    parser.add_argument("--model_name", type=str, default="SSL-AASIST")
+    parser.add_argument("--loss", type=str, default="weighted_CCE")
+    parser.add_argument("--trn_list_path", default=None, help="path to train file")
+    parser.add_argument("--dev_list_path", default=None, help="path to validation file")
+    parser.add_argument("--test_list_path", default=None, help="path to test file")
+    parser.add_argument(
+        "--test_score_dir", default=None, help="path to save test scores"
+    )
+    # model
+    parser.add_argument(
+        "--seed", type=int, default=1234, help="random seed (default: 1234)"
+    )
+    parser.add_argument("--save_path", type=str, default=".", help="Model save path")
+    parser.add_argument("--model_path", type=str, default=None, help="Model checkpoint")
+    parser.add_argument(
+        "--comment", type=str, default=None, help="Comment to describe the saved model"
+    )
+    # Auxiliary arguments
+
+    parser.add_argument("--eval", action="store_true", default=False, help="eval mode")
+    parser.add_argument("--eval_part", type=int, default=0)
+    # backend options
+    parser.add_argument(
+        "--cudnn-deterministic-toggle",
+        action="store_false",
+        default=True,
+        help="use cudnn-deterministic? (default true)",
+    )
+
+    parser.add_argument(
+        "--cudnn-benchmark-toggle",
+        action="store_true",
+        default=False,
+        help="use cudnn-benchmark? (default false)",
+    )
+
+    ##===================================================Rawboost data augmentation ======================================================================#
+
+    parser.add_argument(
+        "--algo",
+        type=int,
+        default=5,
+        help="Rawboost algos discriptions. 0: No augmentation 1: LnL_convolutive_noise, 2: ISD_additive_noise, 3: SSI_additive_noise, 4: series algo (1+2+3), \
+                          5: series algo (1+2), 6: series algo (1+3), 7: series algo(2+3), 8: parallel algo(1,2) .[default=0]",
+    )
+
+    # LnL_convolutive_noise parameters
+    parser.add_argument(
+        "--nBands",
+        type=int,
+        default=5,
+        help="number of notch filters.The higher the number of bands, the more aggresive the distortions is.[default=5]",
+    )
+    parser.add_argument(
+        "--minF",
+        type=int,
+        default=20,
+        help="minimum centre frequency [Hz] of notch filter.[default=20] ",
+    )
+    parser.add_argument(
+        "--maxF",
+        type=int,
+        default=8000,
+        help="maximum centre frequency [Hz] (<sr/2)  of notch filter.[default=8000]",
+    )
+    parser.add_argument(
+        "--minBW",
+        type=int,
+        default=100,
+        help="minimum width [Hz] of filter.[default=100] ",
+    )
+    parser.add_argument(
+        "--maxBW",
+        type=int,
+        default=1000,
+        help="maximum width [Hz] of filter.[default=1000] ",
+    )
+    parser.add_argument(
+        "--minCoeff",
+        type=int,
+        default=10,
+        help="minimum filter coefficients. More the filter coefficients more ideal the filter slope.[default=10]",
+    )
+    parser.add_argument(
+        "--maxCoeff",
+        type=int,
+        default=100,
+        help="maximum filter coefficients. More the filter coefficients more ideal the filter slope.[default=100]",
+    )
+    parser.add_argument(
+        "--minG",
+        type=int,
+        default=0,
+        help="minimum gain factor of linear component.[default=0]",
+    )
+    parser.add_argument(
+        "--maxG",
+        type=int,
+        default=0,
+        help="maximum gain factor of linear component.[default=0]",
+    )
+    parser.add_argument(
+        "--minBiasLinNonLin",
+        type=int,
+        default=5,
+        help=" minimum gain difference between linear and non-linear components.[default=5]",
+    )
+    parser.add_argument(
+        "--maxBiasLinNonLin",
+        type=int,
+        default=20,
+        help=" maximum gain difference between linear and non-linear components.[default=20]",
+    )
+    parser.add_argument(
+        "--N_f",
+        type=int,
+        default=5,
+        help="order of the (non-)linearity where N_f=1 refers only to linear components.[default=5]",
+    )
+
+    # ISD_additive_noise parameters
+    parser.add_argument(
+        "--P",
+        type=int,
+        default=10,
+        help="Maximum number of uniformly distributed samples in [%].[defaul=10]",
+    )
+    parser.add_argument(
+        "--g_sd", type=int, default=2, help="gain parameters > 0. [default=2]"
+    )
+
+    # SSI_additive_noise parameters
+    parser.add_argument(
+        "--SNRmin",
+        type=int,
+        default=10,
+        help="Minimum SNR value for coloured additive noise.[defaul=10]",
+    )
+    parser.add_argument(
+        "--SNRmax",
+        type=int,
+        default=40,
+        help="Maximum SNR value for coloured additive noise.[defaul=40]",
+    )
+
+    ##===================================================Rawboost data augmentation ======================================================================#
+
+    load_dotenv()
+    wandb_api_key = os.getenv("WANDB_API_KEY")
+    wandb_project_name = os.getenv("WANDB_PROJECT_NAME")
+
+    if not os.path.exists("models"):
+        os.mkdir("models")
+    args = parser.parse_args()
+    wandb.login(key=wandb_api_key)
+    wandb.init(
+        project=wandb_project_name,
+        config={
+            "learning_rate": args.lr,
+            "epochs": args.num_epochs,
+            "batch_size": args.batch_size,
+            "weight_decay": args.weight_decay,
+        },
+    )
+
+    # make experiment reproducible
+    set_random_seed(args.seed, args)
+
+    # define model saving path
+    model_tag = "model_{}_{}_{}_{}".format(
+        args.loss, args.num_epochs, args.batch_size, args.lr
+    )
+    if args.comment:
+        model_tag = model_tag + "_{}".format(args.comment)
+    model_save_path = os.path.join(args.save_path, model_tag)
+
+    # set model save directory
+    if not os.path.exists(model_save_path):
+        os.mkdir(model_save_path)
+
+    # GPU device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("Device: {}".format(device))
+
+    model = Model(args, device)
+    nb_params = sum([param.view(-1).size()[0] for param in model.parameters()])
+    model = model.to(device)
+    print("nb_params:", nb_params)
+
+    # set Adam optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
+
+    if args.model_path:
+        model.load_state_dict(torch.load(args.model_path, map_location=device))
+        print("Model loaded : {}".format(args.model_path))
+
+    # evaluation
+
+    if args.eval:
+        file_eval = genSpoof_list(
+            dir_meta=args.test_list_path, is_train=False, is_eval=True
+        )
+        print("no. of eval trials", len(file_eval))
+        eval_set = Dataset_ASVspoof2021_eval(list_IDs=file_eval)
+        eval_output = os.path.join(
+            args.test_score_dir, f"{args.model_name}_model_score.txt"
+        )
+        produce_evaluation_file(
+            eval_set, model, device, eval_output, args.test_list_path
+        )
+        output_file = os.path.join(
+            args.test_score_dir, f"{args.model_name}_model_eer.txt"
+        )
+        eval_eer = calculate_tDCF_EER(
+            cm_scores_file=eval_output, output_file=output_file
+        )
+
+        sys.exit(0)
+
+    trn_list_path = args.trn_list_path
+    dev_trial_path = args.dev_list_path
+    train_set = Dataset_ASVspoof2019_train(args, metafile=trn_list_path, algo=args.algo)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        num_workers=16,
+        shuffle=True,
+        drop_last=True,
+    )
+    del train_set
+
+    dev_set = Dataset_ASVspoof2019_train(args, metafile=dev_trial_path, algo=args.algo)
+    dev_loader = DataLoader(
+        dev_set, batch_size=args.batch_size, num_workers=16, shuffle=False
+    )
+    del dev_set
+    # Training and validation
+    num_epochs = args.num_epochs
+    writer = SummaryWriter("logs/{}".format(model_tag))
+
+    for epoch in range(num_epochs):
+
+        running_loss = train_epoch(
+            train_loader, model, args.lr, optimizer, device, args
+        )
+        val_loss = evaluate_accuracy(dev_loader, model, device, args)
+        wandb.log({"epoch": epoch, "train_loss": running_loss, "val_loss": val_loss})
+        writer.add_scalar("val_loss", val_loss, epoch)
+        writer.add_scalar("loss", running_loss, epoch)
+        print("\n{} - {} - {} ".format(epoch, running_loss, val_loss))
+        torch.save(
+            model.state_dict(),
+            os.path.join(model_save_path, "epoch_{}.pth".format(epoch)),
+        )

clean/image/aide/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Shilin Yan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

clean/image/aide/README.md

@@ -0,0 +1,168 @@
+<div align="center">
+<br>
+<h3>A Sanity Check for AI-generated Image Detection</h3>
+
+[Shilin Yan](https://scholar.google.com/citations?user=2VhjOykAAAAJ&hl=zh-CN&oi=ao)<sup>1†</sup>, Ouxiang Li<sup>1,2†</sup>, Jiayin Cai<sup>1†</sup>, [Yanbin Hao](https://scholar.google.com/citations?user=vhPSOkEAAAAJ&hl=en&oi=ao)<sup>2</sup>, [Xiaolong Jiang](https://scholar.google.com/citations?user=G0Ow8j8AAAAJ&hl=en&oi=ao)<sup>1</sup>, [Yao Hu](https://scholar.google.com/citations?user=LIu7k7wAAAAJ&hl=en)<sup>1</sup>, [Weidi Xie](https://scholar.google.com/citations?user=Vtrqj4gAAAAJ&hl=en)<sup>3‡</sup>
+
+<div class="is-size-6 publication-authors">
+  <p class="footnote">
+    <span class="footnote-symbol"><sup>†</sup></span>Equal contribution
+    <span class="footnote-symbol"><sup>‡</sup></span>Corresponding author
+  </p>
+</div>
+
+<sup>1</sup>Xiaohongshu Inc. <sup>2</sup>University of Science and Technology of China <sup>3</sup>Shanghai Jiao Tong University
+
+
+<p align="center">
+  <a href='https://shilinyan99.github.io/AIDE'>
+    <img src='https://img.shields.io/badge/Project-Page-pink?style=flat&logo=Google%20chrome&logoColor=pink'>
+  </a>
+  <a href='https://arxiv.org/abs/2406.19435'>
+    <img src='https://img.shields.io/badge/Arxiv-2406.19435-A42C25?style=flat&logo=arXiv&logoColor=A42C25'>
+  </a>
+  <a href='https://arxiv.org/pdf/2406.19435'>
+    <img src='https://img.shields.io/badge/Paper-PDF-yellow?style=flat&logo=arXiv&logoColor=yellow'>
+  </a>
+<!--   <img src="https://visitor-badge.laobi.icu/badge?page_id=shilinyan99/AIDE" alt="visitors"> -->
+</p>
+</div>
+
+
+<!-- <div align="center">
+<h1>
+<b>
+A Sanity Check for AI-generated Image Detection
+</b>
+</h1>
+</div> -->
+## 🔥 News
+* [2025-01-23]🎉🎉🎉 AIDE is accepted by ICLR 2025.
+* [2024-12-29]🔥🔥🔥 We release the Chamelon dataset.
+* [2024-06-20]🔥🔥🔥 We release the code and checkpoints of AIDE.
+
+
+## 🔍 Chameleon 
+
+
+**License**:
+```
+Chameleon is only used for academic research. Commercial use in any form is prohibited.
+```
+
+🌟🌟🌟 If you need the Chameleon dataset, please send an email to **tattoo.ysl@gmail.com**. 🔥🔥🔥
+
+
+
+**Comparison of `Chameleon` with existing benchmarks.**
+
+<p align="center"><img src="docs/Chameleon.jpg" width="800"/></p>
+
+We visualize two contemporary AI-generated image benchmarks, namely:
+
+- **(a) AIGCDetect Benchmark** 
+- **(b) GenImage Benchmark** 
+
+where all images are generated from publicly available generators, such as ProGAN (GAN-based), SD v1.4 (DM-based), and Midjourney (commercial API). These images are generated by unconditional situations or conditioned on simple prompts (e.g., *photo of a plane*) without delicate manual adjustments, thereby inclined to generate obvious artifacts in consistency and semantics (marked with <span style="color:red">red boxes</span>).
+
+In contrast, our **`Chameleon`** dataset in **(c)** aims to simulate real-world scenarios by collecting diverse images from online websites, where these online images are carefully adjusted by photographers and AI artists.
+
+
+
+## 👀 Method
+
+We conduct a sanity check on **"whether the task of AI-generated image detection has been solved"**. To start with, we present **Chameleon** dataset, consisting AI-generated images that are genuinely challenging for human perception. To quantify the generalization of existing methods, we evaluate 9 off-the-shelf AI-generated image detectors on **Chameleon** dataset. Upon analysis, almost all models classify AI-generated images as real ones. Later, we propose **AIDE**~(**A**I-generated **I**mage **DE**tector with Hybrid Features), which leverages multiple experts to simultaneously extract visual artifacts and noise patterns. 
+
+<p align="center"><img src="docs/network.png" width="800"/></p>
+
+
+
+
+## Requirements
+
+We test the codes in the following environments, other versions may also be compatible:
+
+- CUDA 11.8
+- Python 3.10
+- Pytorch 2.0.1
+
+
+## Setup
+
+First, clone the repository locally.
+
+```
+https://github.com/shilinyan99/AIDE
+```
+
+Then, install Pytorch 2.0.1 using the conda environment.
+```
+conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 -c pytorch
+```
+
+Lastly, install the necessary packages and pycocotools.
+
+```
+pip install -r requirements.txt 
+```
+
+
+## Get Started
+
+### Training
+
+```
+./scripts/train.sh  --data_path [/path/to/train_data] --eval_data_path [/path/to/eval_data] --resnet_path [/path/to/pretrained_resnet_path] --convnext_path [/path/to/pretrained_convnext_path] --output_dir [/path/to/output_dir] [other args]
+```
+
+For example, training on ProGAN, run the following command:
+
+```
+./scripts/train.sh --data_path dataset/progan/train --eval_data_path dataset/progan/eval --resnet_path pretrained_ckpts/resnet50.pth --convnext_path pretrained_ckpts/open_clip_pytorch_model.bin --output_dir results/progan_train
+```
+
+### Inference
+
+Inference using the trained model.
+```
+./scripts/eval.sh --data_path [/path/to/train_data] --eval_data_path [/path/to/eval_data] --resume [/path/to/progan_train] --eval True --output_dir [/path/to/output_dir]
+```
+
+For example, evaluating the progan_train model, run the following command:
+
+```
+./scripts/eval.sh --data_path dataset/progan/train --eval_data_path dataset/progan/eval --resume results/progan_train/progan_train.pth --eval True --output_dir results/progan_train
+```
+
+
+
+## Dataset
+
+### Training Set
+We adopt the training set in [CNNSpot](https://github.com/peterwang512/CNNDetection) and [GenImage](https://github.com/Andrew-Zhu/GenImage).
+
+### Test Set
+The whole test set we used in our experiments can be downloaded from [AIGCDetectBenchmark](https://github.com/Ekko-zn/AIGCDetectBenchmark?tab=readme-ov-file) and [GenImage](https://github.com/Andrew-Zhu/GenImage).
+
+
+## Model Zoo
+
+Our training checkpoints can be downloaded from [link](https://drive.google.com/drive/folders/1qx76UFvDpgCxaPLBCmsA2WY-SSzeJrd4?usp=sharing).
+
+## Acknowledgement
+
+This repo is based on [ConvNeXt](https://github.com/facebookresearch/ConvNeXt-V2). We also refer to the repositories [CNNSpot](https://github.com/peterwang512/CNNDetection)、[AIGCDetectBenchmark](https://github.com/Ekko-zn/AIGCDetectBenchmark?tab=readme-ov-file)、[GenImage](https://github.com/Andrew-Zhu/GenImage) and [DNF](https://github.com/YichiCS/DNF). Thanks for their wonderful works.
+
+## Citation
+
+```
+@article{yan2024sanity,
+  title={A Sanity Check for AI-generated Image Detection},
+  author={Yan, Shilin and Li, Ouxiang and Cai, Jiayin and Hao, Yanbin and Jiang, Xiaolong and Hu, Yao and Xie, Weidi},
+  journal={arXiv preprint arXiv:2406.19435},
+  year={2024}
+}
+```
+
+## Contact
+If you have any question about this project, please feel free to contact tattoo.ysl@gmail.com.

clean/image/aide/SOURCE.md

@@ -0,0 +1,17 @@
+# Source: image/aide
+
+| Field | Value |
+|---|---|
+| Upstream | https://github.com/shilinyan99/AIDE |
+| Paper | https://arxiv.org/abs/2406.19435 |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/image__aide.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.

clean/image/aide/data/dct.py

@@ -0,0 +1,107 @@
+"""DCT frequency decomposition module for AIDE.
+
+Source: https://github.com/shilinyan99/AIDE (ICLR 2025)
+Selects the most and least frequency-active image patches for
+multi-view input to the AIDE deepfake detector.
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def DCT_mat(size):
+    m = [[(np.sqrt(1./size) if i == 0 else np.sqrt(2./size)) * np.cos((j + 0.5) * np.pi * i / size) for j in range(size)] for i in range(size)]
+    return m
+
+def generate_filter(start, end, size):
+    return [[0. if i + j > end or i + j < start else 1. for j in range(size)] for i in range(size)]
+
+def norm_sigma(x):
+    return 2. * torch.sigmoid(x) - 1.
+
+class Filter(nn.Module):
+    def __init__(self, size, band_start, band_end, use_learnable=False, norm=False):
+        super(Filter, self).__init__()
+        self.use_learnable = use_learnable
+        self.base = nn.Parameter(torch.tensor(generate_filter(band_start, band_end, size)), requires_grad=False)
+        if self.use_learnable:
+            self.learnable = nn.Parameter(torch.randn(size, size), requires_grad=True)
+            self.learnable.data.normal_(0., 0.1)
+        self.norm = norm
+        if norm:
+            self.ft_num = nn.Parameter(torch.sum(torch.tensor(generate_filter(band_start, band_end, size))), requires_grad=False)
+
+    def forward(self, x):
+        if self.use_learnable:
+            filt = self.base + norm_sigma(self.learnable)
+        else:
+            filt = self.base
+        if self.norm:
+            y = x * filt / self.ft_num
+        else:
+            y = x * filt
+        return y
+
+class DCT_base_Rec_Module(nn.Module):
+    def __init__(self, window_size=32, stride=16, output=256, grade_N=6, level_fliter=[0]):
+        super().__init__()
+        assert output % window_size == 0
+        assert len(level_fliter) > 0
+        self.window_size = window_size
+        self.grade_N = grade_N
+        self.level_N = len(level_fliter)
+        self.N = (output // window_size) * (output // window_size)
+        self._DCT_patch = nn.Parameter(torch.tensor(DCT_mat(window_size)).float(), requires_grad=False)
+        self._DCT_patch_T = nn.Parameter(torch.transpose(torch.tensor(DCT_mat(window_size)).float(), 0, 1), requires_grad=False)
+        self.unfold = nn.Unfold(kernel_size=(window_size, window_size), stride=stride)
+        self.fold0 = nn.Fold(output_size=(window_size, window_size), kernel_size=(window_size, window_size), stride=window_size)
+        level_f = [Filter(window_size, 0, window_size * 2)]
+        self.level_filters = nn.ModuleList([level_f[i] for i in level_fliter])
+        self.grade_filters = nn.ModuleList([Filter(window_size, window_size * 2. / grade_N * i, window_size * 2. / grade_N * (i+1), norm=True) for i in range(grade_N)])
+
+    def forward(self, x):
+        N = self.N
+        grade_N = self.grade_N
+        level_N = self.level_N
+        window_size = self.window_size
+        C, W, H = x.shape
+        x_unfold = self.unfold(x.unsqueeze(0)).squeeze(0)
+        _, L = x_unfold.shape
+        x_unfold = x_unfold.transpose(0, 1).reshape(L, C, window_size, window_size)
+        x_dct = self._DCT_patch @ x_unfold @ self._DCT_patch_T
+        y_list = []
+        for i in range(self.level_N):
+            x_pass = self.level_filters[i](x_dct)
+            y = self._DCT_patch_T @ x_pass @ self._DCT_patch
+            y_list.append(y)
+        level_x_unfold = torch.cat(y_list, dim=1)
+        grade = torch.zeros(L).to(x.device)
+        w, k = 1, 2
+        for _ in range(grade_N):
+            _x = torch.abs(x_dct)
+            _x = torch.log(_x + 1)
+            _x = self.grade_filters[_](_x)
+            _x = torch.sum(_x, dim=[1,2,3])
+            grade += w * _x
+            w *= k
+        _, idx = torch.sort(grade)
+        max_idx = torch.flip(idx, dims=[0])[:N]
+        maxmax_idx = max_idx[0]
+        maxmax_idx1 = max_idx[1] if len(max_idx) > 1 else max_idx[0]
+        min_idx = idx[:N]
+        minmin_idx = idx[0]
+        minmin_idx1 = idx[1] if len(min_idx) > 1 else idx[0]
+        x_minmin = torch.index_select(level_x_unfold, 0, minmin_idx)
+        x_maxmax = torch.index_select(level_x_unfold, 0, maxmax_idx)
+        x_minmin1 = torch.index_select(level_x_unfold, 0, minmin_idx1)
+        x_maxmax1 = torch.index_select(level_x_unfold, 0, maxmax_idx1)
+        x_minmin = x_minmin.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+        x_maxmax = x_maxmax.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+        x_minmin1 = x_minmin1.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+        x_maxmax1 = x_maxmax1.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+        x_minmin = self.fold0(x_minmin)
+        x_maxmax = self.fold0(x_maxmax)
+        x_minmin1 = self.fold0(x_minmin1)
+        x_maxmax1 = self.fold0(x_maxmax1)
+        return x_minmin, x_maxmax, x_minmin1, x_maxmax1

clean/image/aide/engine_finetune.py

@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import math
+from typing import Iterable, Optional
+
+import torch
+import torch.distributed as dist
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+
+import utils
+from utils import adjust_learning_rate
+from scipy.special import softmax
+from sklearn.metrics import (
+    average_precision_score, 
+    accuracy_score
+)
+import numpy as np
+
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None, 
+                    log_writer=None, args=None):
+    model.train(True)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    update_freq = args.update_freq
+    use_amp = args.use_amp
+    optimizer.zero_grad()
+
+    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % update_freq == 0:
+            adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        if use_amp:
+            with torch.cuda.amp.autocast():
+                output = model(samples)
+                loss = criterion(output, targets)
+        else: # full precision
+            output = model(samples)
+            loss = criterion(output, targets)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            assert math.isfinite(loss_value)
+
+        if use_amp:
+            # this attribute is added by timm on one optimizer (adahessian)
+            is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
+            loss /= update_freq
+            grad_norm = loss_scaler(loss, optimizer, clip_grad=max_norm,
+                                    parameters=model.parameters(), create_graph=is_second_order,
+                                    update_grad=(data_iter_step + 1) % update_freq == 0)
+            if (data_iter_step + 1) % update_freq == 0:
+                optimizer.zero_grad()
+                if model_ema is not None:
+                    model_ema.update(model)
+        else: # full precision
+            loss /= update_freq
+            loss.backward()
+            if (data_iter_step + 1) % update_freq == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+                if model_ema is not None:
+                    model_ema.update(model)
+        
+        torch.cuda.synchronize()
+
+        if mixup_fn is None:
+            class_acc = (output.max(-1)[-1] == targets).float().mean()
+        else:
+            class_acc = None
+
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(class_acc=class_acc)
+        min_lr = 10.
+        max_lr = 0.
+        for group in optimizer.param_groups:
+            min_lr = min(min_lr, group["lr"])
+            max_lr = max(max_lr, group["lr"])
+        
+        metric_logger.update(lr=max_lr)
+        metric_logger.update(min_lr=min_lr)
+        weight_decay_value = None
+        for group in optimizer.param_groups:
+            if group["weight_decay"] > 0:
+                weight_decay_value = group["weight_decay"]
+        metric_logger.update(weight_decay=weight_decay_value)
+        if use_amp:
+            metric_logger.update(grad_norm=grad_norm)
+        if log_writer is not None:
+            log_writer.update(loss=loss_value, head="loss")
+            log_writer.update(class_acc=class_acc, head="loss")
+            log_writer.update(lr=max_lr, head="opt")
+            log_writer.update(min_lr=min_lr, head="opt")
+            log_writer.update(weight_decay=weight_decay_value, head="opt")
+            if use_amp:
+                log_writer.update(grad_norm=grad_norm, head="opt")
+            log_writer.set_step()
+    
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+@torch.no_grad()
+def evaluate(data_loader, model, device, use_amp=False):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    for index, batch in enumerate(metric_logger.log_every(data_loader, 10, header)):
+        images = batch[0]
+        target = batch[-1]
+
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        if use_amp:
+            with torch.cuda.amp.autocast(dytpe=torch.bfloat16):
+                output = model(images)
+                if isinstance(output, dict):
+                    output = output['logits']
+                loss = criterion(output, target)
+        else:
+            output = model(images) #[bs, num_cls]
+            if isinstance(output, dict):
+                output = output['logits']
+            
+            loss = criterion(output, target)
+        
+        if index == 0:
+            predictions = output
+            labels = target
+        else:
+            predictions = torch.cat((predictions, output), 0)
+            labels = torch.cat((labels, target), 0)
+
+        torch.cuda.synchronize()
+
+        acc1, acc5 = accuracy(output, target, topk=(1, 2))
+
+        batch_size = images.shape[0]
+        metric_logger.update(loss=loss.item())
+        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+          .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+
+
+    output_ddp = [torch.zeros_like(predictions) for _ in range(utils.get_world_size())]
+    dist.all_gather(output_ddp, predictions)
+    labels_ddp = [torch.zeros_like(labels) for _ in range(utils.get_world_size())]
+    dist.all_gather(labels_ddp, labels)
+
+    output_all = torch.concat(output_ddp, dim=0)
+    labels_all = torch.concat(labels_ddp, dim=0)
+
+
+    y_pred = softmax(output_all.detach().cpu().numpy(), axis=1)[:, 1]
+    y_true = labels_all.detach().cpu().numpy()
+    y_true = y_true.astype(int)
+    
+  
+    acc = accuracy_score(y_true, y_pred > 0.5)
+    ap = average_precision_score(y_true, y_pred)
+    
+   
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, acc, ap

clean/image/aide/main_finetune.py

@@ -0,0 +1,449 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import argparse
+import datetime
+import numpy as np
+import time
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+
+from timm.models.layers import trunc_normal_
+from timm.data.mixup import Mixup
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.utils import ModelEma
+from optim_factory import create_optimizer, LayerDecayValueAssigner
+
+from data.datasets import TrainDataset, TestDataset
+from engine_finetune import train_one_epoch, evaluate
+
+import utils
+from utils import NativeScalerWithGradNormCount as NativeScaler
+from utils import str2bool, remap_checkpoint_keys
+import models.AIDE as AIDE
+import csv
+import warnings
+
+warnings.filterwarnings('ignore')
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('Resnet fine-tuning', add_help=False)
+    parser.add_argument('--batch_size', default=64, type=int,
+                        help='Per GPU batch size')
+    parser.add_argument('--epochs', default=100, type=int)
+    parser.add_argument('--update_freq', default=1, type=int,
+                        help='gradient accumulation steps')
+
+    # Model parameters
+    parser.add_argument('--model', default='AIDE', type=str, metavar='MODEL',
+                        help='Name of model to train')
+    parser.add_argument('--resnet_path', default=None, type=str, metavar='MODEL',
+                        help='Path of resnet model')
+    parser.add_argument('--convnext_path', default=None, type=str, metavar='MODEL',
+                        help='Path of ConvNeXt of model ')
+    
+    # EMA related parameters
+    parser.add_argument('--model_ema', type=str2bool, default=False)
+    parser.add_argument('--model_ema_decay', type=float, default=0.9999, help='')
+    parser.add_argument('--model_ema_force_cpu', type=str2bool, default=False, help='')
+    parser.add_argument('--model_ema_eval', type=str2bool, default=False, help='Using ema to eval during training.')
+
+    # Optimization parameters
+    parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
+                        help='Clip gradient norm (default: None, no clipping)')
+    parser.add_argument('--weight_decay', type=float, default=0.,
+                        help='weight decay (default: 0.05)')
+    parser.add_argument('--lr', type=float, default=None, metavar='LR',
+                        help='learning rate (absolute lr)')
+    parser.add_argument('--blr', type=float, default=5e-4, metavar='LR',
+                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
+    parser.add_argument('--layer_decay', type=float, default=1.0)
+    parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR',
+                        help='lower lr bound for cyclic schedulers that hit 0 (1e-6)')
+    parser.add_argument('--warmup_epochs', type=int, default=0, metavar='N',
+                        help='epochs to warmup LR, if scheduler supports')
+    
+    parser.add_argument('--warmup_steps', type=int, default=-1, metavar='N',
+                        help='num of steps to warmup LR, will overload warmup_epochs if set > 0')    
+    parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
+                        help='Optimizer (default: "adamw"')
+    parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON',
+                        help='Optimizer Epsilon (default: 1e-8)')
+    parser.add_argument('--opt_betas', default=None, type=float, nargs='+', metavar='BETA',
+                        help='Optimizer Betas (default: None, use opt default)')
+    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
+                        help='SGD momentum (default: 0.9)')
+    parser.add_argument('--weight_decay_end', type=float, default=None, help="""Final value of the
+        weight decay. We use a cosine schedule for WD and using a larger decay by
+        the end of training improves performance for ViTs.""")
+
+    # Augmentation parameters
+    parser.add_argument('--color_jitter', type=float, default=None, metavar='PCT',
+                       help='Color jitter factor (enabled only when not using Auto/RandAug)')
+    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
+                        help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m9-mstd0.5-inc1)')
+    parser.add_argument('--smoothing', type=float, default=0.1,
+                        help='Label smoothing (default: 0.1)')
+    
+    parser.add_argument('--train_interpolation', type=str, default='bicubic',
+                        help='Training interpolation (random, bilinear, bicubic default: "bicubic")')
+
+    # * Random Erase params
+    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
+                        help='Random erase prob (default: 0.25)')
+    parser.add_argument('--remode', type=str, default='pixel',
+                        help='Random erase mode (default: "pixel")')
+    parser.add_argument('--recount', type=int, default=1,
+                        help='Random erase count (default: 1)')
+    parser.add_argument('--resplit', type=str2bool, default=False,
+                        help='Do not random erase first (clean) augmentation split')
+
+    # * Mixup params
+    parser.add_argument('--mixup', type=float, default=0.,
+                        help='mixup alpha, mixup enabled if > 0.')
+    parser.add_argument('--cutmix', type=float, default=0.,
+                        help='cutmix alpha, cutmix enabled if > 0.')
+    parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None,
+                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
+    parser.add_argument('--mixup_prob', type=float, default=1.0,
+                        help='Probability of performing mixup or cutmix when either/both is enabled')
+    parser.add_argument('--mixup_switch_prob', type=float, default=0.5,
+                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
+    parser.add_argument('--mixup_mode', type=str, default='batch',
+                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
+
+    # * Finetuning params
+    parser.add_argument('--finetune', default='',
+                        help='finetune from checkpoint')
+    parser.add_argument('--head_init_scale', default=0.001, type=float,
+                        help='classifier head initial scale, typically adjusted in fine-tuning')
+    parser.add_argument('--model_key', default='model|module', type=str,
+                        help='which key to load from saved state dict, usually model or model_ema')
+    parser.add_argument('--model_prefix', default='', type=str)
+
+    # Dataset parameters
+    parser.add_argument('--data_path', default='path/dataset', type=str,
+                        help='dataset path')
+    parser.add_argument('--nb_classes', default=2, type=int,
+                        help='number of the classification types')
+    parser.add_argument('--output_dir', default='',
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--log_dir', default=None,
+                        help='path where to tensorboard log')
+    parser.add_argument('--device', default='cuda',
+                        help='device to use for training / testing')
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--resume', default='',
+                        help='resume from checkpoint')
+    
+    parser.add_argument('--eval_data_path', default=None, type=str,
+                        help='dataset path for evaluation')
+    parser.add_argument('--imagenet_default_mean_and_std', type=str2bool, default=True)
+    parser.add_argument('--data_set', default='IMNET', choices=['CIFAR', 'IMNET', 'image_folder'],
+                        type=str, help='ImageNet dataset path')
+    parser.add_argument('--auto_resume', type=str2bool, default=True)
+    parser.add_argument('--save_ckpt', type=str2bool, default=True)
+    parser.add_argument('--save_ckpt_freq', default=1, type=int)
+    parser.add_argument('--save_ckpt_num', default=100, type=int)
+
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
+    parser.add_argument('--eval', type=str2bool, default=False,
+                        help='Perform evaluation only')
+    parser.add_argument('--dist_eval', type=str2bool, default=True,
+                        help='Enabling distributed evaluation')
+    parser.add_argument('--disable_eval', type=str2bool, default=False,
+                        help='Disabling evaluation during training')
+    parser.add_argument('--num_workers', default=16, type=int)
+    parser.add_argument('--pin_mem', type=str2bool, default=True,
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+
+    # Evaluation parameters
+    parser.add_argument('--crop_pct', type=float, default=None)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', type=str2bool, default=False)
+    parser.add_argument('--dist_url', default='env://',
+                        help='url used to set up distributed training')
+    
+    parser.add_argument('--use_amp', type=str2bool, default=False, 
+                        help="Use apex AMP (Automatic Mixed Precision) or not")
+    return parser
+
+def main(args):
+    utils.init_distributed_mode(args)
+    print(args)
+    device = torch.device(args.device)
+
+    # fix the seed for reproducibility
+    seed = args.seed + utils.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = True
+
+    dataset_train = TrainDataset(is_train=True, args=args)
+
+    if args.disable_eval:
+        args.dist_eval = False
+        dataset_val = None
+    else:
+        dataset_val = TrainDataset(is_train=False, args=args)
+
+    num_tasks = utils.get_world_size()
+    global_rank = utils.get_rank()
+    
+    sampler_train = torch.utils.data.DistributedSampler(
+        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True, seed=args.seed,
+    )
+    print("Sampler_train = %s" % str(sampler_train))
+    if args.dist_eval:
+        if len(dataset_val) % num_tasks != 0:
+            print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
+                    'This will slightly alter validation results as extra duplicate entries are added to achieve '
+                    'equal num of samples per-process.')
+        sampler_val = torch.utils.data.DistributedSampler(
+            dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+    else:
+        sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+    
+    if global_rank == 0 and args.log_dir is not None:
+        os.makedirs(args.log_dir, exist_ok=True)
+        log_writer = utils.TensorboardLogger(log_dir=args.log_dir)
+    else:
+        log_writer = None
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+    )
+    if dataset_val is not None:
+        data_loader_val = torch.utils.data.DataLoader(
+            dataset_val, sampler=sampler_val,
+            batch_size=args.batch_size,
+            num_workers=args.num_workers,
+            pin_memory=args.pin_mem,
+            drop_last=False
+        )
+    else:
+        data_loader_val = None
+
+    mixup_fn = None
+    mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
+    if mixup_active:
+        print("Mixup is activated!")
+        mixup_fn = Mixup(
+            mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
+            prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
+            label_smoothing=args.smoothing, num_classes=args.nb_classes)
+
+
+    model = AIDE.__dict__[args.model](
+        resnet_path=args.resnet_path, 
+        convnext_path=args.convnext_path
+    )
+        
+    model.to(device)
+
+    model_ema = None
+    if args.model_ema:
+        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+        model_ema = ModelEma(
+            model,
+            decay=args.model_ema_decay,
+            device='cpu' if args.model_ema_force_cpu else '',
+            resume='')
+        print("Using EMA with decay = %.8f" % args.model_ema_decay)
+
+    model_without_ddp = model
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print("Model = %s" % str(model_without_ddp))
+    print('number of params:', n_parameters)
+
+    eff_batch_size = args.batch_size * args.update_freq * utils.get_world_size()
+    num_training_steps_per_epoch = len(dataset_train) // eff_batch_size
+    
+    if args.lr is None:
+        args.lr = args.blr * eff_batch_size / 256
+
+    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
+    print("actual lr: %.2e" % args.lr)
+
+    print("accumulate grad iterations: %d" % args.update_freq)
+    print("effective batch size: %d" % eff_batch_size)
+
+    assigner = None
+    
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    optimizer = create_optimizer(
+        args, model_without_ddp, skip_list=None,
+        get_num_layer=assigner.get_layer_id if assigner is not None else None, 
+        get_layer_scale=assigner.get_scale if assigner is not None else None)
+    loss_scaler = NativeScaler()
+
+    if mixup_fn is not None:
+        # smoothing is handled with mixup label transform
+        criterion = SoftTargetCrossEntropy()
+    elif args.smoothing > 0.:
+        criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
+    else:
+        criterion = torch.nn.CrossEntropyLoss()
+    
+    print("criterion = %s" % str(criterion))
+
+    utils.auto_load_model(
+        args=args, model=model, model_without_ddp=model_without_ddp,
+        optimizer=optimizer, loss_scaler=loss_scaler, model_ema=model_ema)
+
+    if args.eval:
+        print(f"Eval only mode")
+        
+        vals = os.listdir(args.eval_data_path)
+        if len(vals) == 16:
+            vals = ["progan", "stylegan", "biggan", "cyclegan", "stargan", "gaugan", "stylegan2", "whichfaceisreal", "ADM", "Glide", "Midjourney", "stable_diffusion_v_1_4", "stable_diffusion_v_1_5", "VQDM", "wukong", "DALLE2"]
+        if len(vals) == 8:
+            vals = ["Midjourney", "stable_diffusion_v_1_4", "stable_diffusion_v_1_5", "ADM", "glide", "wukong", "VQDM", "BigGAN"]
+        eval_data_path = args.eval_data_path
+
+        rows = [["{} model testing on...".format(args.resume)],
+            ['testset', 'accuracy', 'avg precision']]
+
+        for v_id, val in enumerate(vals):
+            
+            args.eval_data_path = os.path.join(args.eval_data_path, val)
+            dataset_val = TestDataset(is_train=False, args=args)
+            args.eval_data_path = eval_data_path
+
+            if args.dist_eval:
+                if len(dataset_val) % num_tasks != 0:
+                    print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
+                            'This will slightly alter validation results as extra duplicate entries are added to achieve '
+                            'equal num of samples per-process.')
+                sampler_val = torch.utils.data.DistributedSampler(
+                    dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+            else:
+                sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+
+            data_loader_val = torch.utils.data.DataLoader(
+                dataset_val, sampler=sampler_val,
+                batch_size=args.batch_size,
+                num_workers=args.num_workers,
+                pin_memory=args.pin_mem,
+                drop_last=False
+            )
+
+
+            test_stats, acc, ap = evaluate(data_loader_val, model, device)
+            print(f"Accuracy of the network on {len(dataset_val)} test images: {test_stats['acc1']:.5f}%")
+        
+            print(f"test dataset is {val} acc: {acc}, ap: {ap}")
+            print("***********************************")
+
+            rows.append([val, acc, ap])
+
+
+        test_dataset_name  = args.eval_data_path.split('/')[-2]
+
+        csv_name = os.path.join(args.output_dir, f'{os.path.basename(args.resume)}_{test_dataset_name}.csv')
+        with open(csv_name, 'w') as f:
+            csv_writer = csv.writer(f, delimiter=',')
+            csv_writer.writerows(rows)
+        return
+    
+    max_accuracy = 0.0
+    if args.model_ema and args.model_ema_eval:
+        max_accuracy_ema = 0.0
+
+    print("Start training for %d epochs" % args.epochs)
+    start_time = time.time()
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+        if log_writer is not None:
+            log_writer.set_step(epoch * num_training_steps_per_epoch * args.update_freq)
+        train_stats = train_one_epoch(
+            model, criterion, data_loader_train,
+            optimizer, device, epoch, loss_scaler, 
+            args.clip_grad, model_ema, mixup_fn,
+            log_writer=log_writer,
+            args=args
+        )
+        if args.output_dir and args.save_ckpt:
+            if (epoch + 1) % args.save_ckpt_freq == 0 or epoch + 1 == args.epochs:
+                utils.save_model(
+                    args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+                    loss_scaler=loss_scaler, epoch=epoch, model_ema=model_ema)
+        if data_loader_val is not None:
+            test_stats, acc, ap = evaluate(data_loader_val, model, device, use_amp=args.use_amp)
+            print(f"Accuracy of the model on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%, ap: {ap}.")
+            if max_accuracy < test_stats["acc1"]:
+                max_accuracy = test_stats["acc1"]
+                if args.output_dir and args.save_ckpt:
+                    utils.save_model(
+                        args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+                        loss_scaler=loss_scaler, epoch="best", model_ema=model_ema)
+            print(f'Max accuracy: {max_accuracy:.2f}%')
+
+            if log_writer is not None:
+                log_writer.update(test_acc1=test_stats['acc1'], head="perf", step=epoch)
+                log_writer.update(test_acc5=test_stats['acc5'], head="perf", step=epoch)
+                log_writer.update(test_loss=test_stats['loss'], head="perf", step=epoch)
+
+            log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                         **{f'test_{k}': v for k, v in test_stats.items()},
+                         'epoch': epoch,
+                         'n_parameters': n_parameters}
+            
+            # repeat testing routines for EMA, if ema eval is turned on
+            if args.model_ema and args.model_ema_eval:
+                test_stats_ema, acc, ap = evaluate(data_loader_val, model_ema.ema, device, use_amp=args.use_amp)
+                print(f"Accuracy of the model EMA on {len(dataset_val)} test images: {test_stats_ema['acc1']:.1f}%, ap: {ap}")
+                if max_accuracy_ema < test_stats_ema["acc1"]:
+                    max_accuracy_ema = test_stats_ema["acc1"]
+                    if args.output_dir and args.save_ckpt:
+                        utils.save_model(
+                            args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+                            loss_scaler=loss_scaler, epoch="best-ema", model_ema=model_ema)
+                    print(f'Max EMA accuracy: {max_accuracy_ema:.2f}%')
+                if log_writer is not None:
+                    log_writer.update(test_acc1_ema=test_stats_ema['acc1'], head="perf", step=epoch)
+                log_stats.update({**{f'test_{k}_ema': v for k, v in test_stats_ema.items()}})
+        else:
+            log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                        'epoch': epoch,
+                        'n_parameters': n_parameters}
+        
+        if args.output_dir and utils.is_main_process():
+            if log_writer is not None:
+                log_writer.flush()
+            with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('AIDE traning', parents=[get_args_parser()])
+    args = parser.parse_args()
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

clean/image/aide/models/AIDE.py

@@ -0,0 +1,298 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch
+import clip
+import open_clip
+from .srm_filter_kernel import all_normalized_hpf_list
+import numpy as np
+
+class HPF(nn.Module):
+  def __init__(self):
+    super(HPF, self).__init__()
+
+    #Load 30 SRM Filters
+    all_hpf_list_5x5 = []
+
+    for hpf_item in all_normalized_hpf_list:
+      if hpf_item.shape[0] == 3:
+        hpf_item = np.pad(hpf_item, pad_width=((1, 1), (1, 1)), mode='constant')
+
+      all_hpf_list_5x5.append(hpf_item)
+
+    hpf_weight = torch.Tensor(all_hpf_list_5x5).view(30, 1, 5, 5).contiguous()
+    hpf_weight = torch.nn.Parameter(hpf_weight.repeat(1, 3, 1, 1), requires_grad=False)
+   
+
+    self.hpf = nn.Conv2d(3, 30, kernel_size=5, padding=2, bias=False)
+    self.hpf.weight = hpf_weight
+
+
+  def forward(self, input):
+
+    output = self.hpf(input)
+
+    return output
+
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=True):
+        super(ResNet, self).__init__()
+
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(30, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+
+
+        return x
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+class AIDE_Model(nn.Module):
+
+    def __init__(self, resnet_path, convnext_path):
+        super(AIDE_Model, self).__init__()
+        self.hpf = HPF()
+        self.model_min = ResNet(Bottleneck, [3, 4, 6, 3])
+        self.model_max = ResNet(Bottleneck, [3, 4, 6, 3])
+       
+        if resnet_path is not None:
+            pretrained_dict = torch.load(resnet_path, map_location='cpu')
+        
+            model_min_dict = self.model_min.state_dict()
+            model_max_dict = self.model_max.state_dict()
+    
+            for k in pretrained_dict.keys():
+                if k in model_min_dict and pretrained_dict[k].size() == model_min_dict[k].size():
+                    model_min_dict[k] = pretrained_dict[k]
+                    model_max_dict[k] = pretrained_dict[k]
+                else:
+                    print(f"Skipping layer {k} because of size mismatch")
+        
+        self.fc = Mlp(2048 + 256 , 1024, 2)
+
+        print("build model with convnext_xxl")
+        self.openclip_convnext_xxl, _, _ = open_clip.create_model_and_transforms(
+            "convnext_xxlarge", pretrained=convnext_path
+        )
+
+        self.openclip_convnext_xxl = self.openclip_convnext_xxl.visual.trunk
+        self.openclip_convnext_xxl.head.global_pool = nn.Identity()
+        self.openclip_convnext_xxl.head.flatten = nn.Identity()
+
+        self.openclip_convnext_xxl.eval()
+        
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.convnext_proj = nn.Sequential(
+            nn.Linear(3072, 256),
+
+        )
+        for param in self.openclip_convnext_xxl.parameters():
+            param.requires_grad = False
+
+    
+
+    def forward(self, x):
+
+        b, t, c, h, w = x.shape
+
+        x_minmin = x[:, 0] #[b, c, h, w]
+        x_maxmax = x[:, 1]
+        x_minmin1 = x[:, 2]
+        x_maxmax1 = x[:, 3]
+        tokens = x[:, 4]
+
+        x_minmin = self.hpf(x_minmin)
+        x_maxmax = self.hpf(x_maxmax)
+        x_minmin1 = self.hpf(x_minmin1)
+        x_maxmax1 = self.hpf(x_maxmax1)
+
+        with torch.no_grad():
+            
+            clip_mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073])
+            clip_mean = clip_mean.to(tokens, non_blocking=True).view(3, 1, 1)
+            clip_std = torch.Tensor([0.26862954, 0.26130258, 0.27577711])
+            clip_std = clip_std.to(tokens, non_blocking=True).view(3, 1, 1)
+            dinov2_mean = torch.Tensor([0.485, 0.456, 0.406]).to(tokens, non_blocking=True).view(3, 1, 1)
+            dinov2_std = torch.Tensor([0.229, 0.224, 0.225]).to(tokens, non_blocking=True).view(3, 1, 1)
+
+            local_convnext_image_feats = self.openclip_convnext_xxl(
+                tokens * (dinov2_std / clip_std) + (dinov2_mean - clip_mean) / clip_std
+            ) #[b, 3072, 8, 8]
+            assert local_convnext_image_feats.size()[1:] == (3072, 8, 8)
+            local_convnext_image_feats = self.avgpool(local_convnext_image_feats).view(tokens.size(0), -1)
+            x_0 = self.convnext_proj(local_convnext_image_feats)
+
+        x_min = self.model_min(x_minmin)
+        x_max = self.model_max(x_maxmax)
+        x_min1 = self.model_min(x_minmin1)
+        x_max1 = self.model_max(x_maxmax1)
+
+        x_1 = (x_min + x_max + x_min1 + x_max1) / 4
+
+        x = torch.cat([x_0, x_1], dim=1)
+
+        x = self.fc(x)
+
+        return x
+
+def AIDE(resnet_path, convnext_path):
+    model = AIDE_Model(resnet_path, convnext_path)
+    return model
+

clean/image/aide/models/srm_filter_kernel.py

@@ -0,0 +1,220 @@
+
+import numpy as np
+
+filter_class_1 = [
+  np.array([
+    [1, 0, 0],
+    [0, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 1, 0],
+    [0, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 1],
+    [0, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [1, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 1],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 0],
+    [1, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 0],
+    [0, 1, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 0],
+    [0, 0, 1]
+  ], dtype=np.float32)
+]
+
+
+filter_class_2 = [
+  np.array([
+    [1, 0, 0],
+    [0, -2, 0],
+    [0, 0, 1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 1, 0],
+    [0, -2, 0],
+    [0, 1, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 1],
+    [0, -2, 0],
+    [1, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [1, -2, 1],
+    [0, 0, 0]
+  ], dtype=np.float32),
+]
+
+
+filter_class_3 = [
+  np.array([
+    [-1, 0, 0, 0, 0],
+    [0, 3, 0, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 0, 1, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, -1, 0, 0],
+    [0, 0, 3, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 1, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, -1],
+    [0, 0, 0, 3, 0],
+    [0, 0, -3, 0, 0],
+    [0, 1, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0],
+    [0, 1, -3, 3, -1],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 1, 0, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 0, 3, 0],
+    [0, 0, 0, 0, -1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 1, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 3, 0, 0],
+    [0, 0, -1, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 1, 0],
+    [0, 0, -3, 0, 0],
+    [0, 3, 0, 0, 0],
+    [-1, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0],
+    [-1, 3, -3, 1, 0],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32)
+]
+
+
+filter_edge_3x3 = [
+  np.array([
+    [-1, 2, -1],
+    [2, -4, 2],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 2, -1],
+    [0, -4, 2],
+    [0, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [2, -4, 2],
+    [-1, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [-1, 2, 0],
+    [2, -4, 0],
+    [-1, 2, 0]
+  ], dtype=np.float32),
+]
+
+filter_edge_5x5 = [
+  np.array([
+    [-1, 2, -2, 2, -1],
+    [2, -6, 8, -6, 2],
+    [-2, 8, -12, 8, -2],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, -2, 2, -1],
+    [0, 0, 8, -6, 2],
+    [0, 0, -12, 8, -2],
+    [0, 0, 8, -6, 2],
+    [0, 0, -2, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0],
+    [-2, 8, -12, 8, -2],
+    [2, -6, 8, -6, 2],
+    [-1, 2, -2, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [-1, 2, -2, 0, 0],
+    [2, -6, 8, 0, 0],
+    [-2, 8, -12, 0, 0],
+    [2, -6, 8, 0, 0],
+    [-1, 2, -2, 0, 0]
+  ], dtype=np.float32),
+]
+
+square_3x3 = np.array([
+  [-1, 2, -1],
+  [2, -4, 2],
+  [-1, 2, -1]
+], dtype=np.float32)
+
+square_5x5 = np.array([
+  [-1, 2, -2, 2, -1],
+  [2, -6, 8, -6, 2],
+  [-2, 8, -12, 8, -2],
+  [2, -6, 8, -6, 2],
+  [-1, 2, -2, 2, -1]
+], dtype=np.float32)
+
+
+all_hpf_list = filter_class_1 + filter_class_2 + filter_class_3 + filter_edge_3x3 + filter_edge_5x5 + [square_3x3, square_5x5]
+
+hpf_3x3_list = filter_class_1 + filter_class_2 + filter_edge_3x3 + [square_3x3]
+hpf_5x5_list = filter_class_3 + filter_edge_5x5 + [square_5x5]
+
+normalized_filter_class_2 = [hpf / 2 for hpf in filter_class_2]
+normalized_filter_class_3 = [hpf / 3 for hpf in filter_class_3]
+normalized_filter_edge_3x3 = [hpf / 4 for hpf in filter_edge_3x3]
+normalized_square_3x3 = square_3x3 / 4
+normalized_filter_edge_5x5 = [hpf / 12 for hpf in filter_edge_5x5]
+normalized_square_5x5 = square_5x5 / 12
+
+all_normalized_hpf_list = filter_class_1 + normalized_filter_class_2 + normalized_filter_class_3 + \
+ normalized_filter_edge_3x3 + normalized_filter_edge_5x5 + [normalized_square_3x3, normalized_square_5x5]
+
+normalized_hpf_3x3_list = filter_class_1 + normalized_filter_class_2 + normalized_filter_edge_3x3 + [normalized_square_3x3]
+normalized_hpf_5x5_list = normalized_filter_class_3 + normalized_filter_edge_5x5 + [normalized_square_5x5]
+
+normalized_3x3_list = normalized_filter_edge_3x3 + [normalized_square_3x3]
+normalized_5x5_list = normalized_filter_edge_5x5 + [normalized_square_5x5]

clean/image/aide/models/utils.py

@@ -0,0 +1,116 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import numpy.random as random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from MinkowskiEngine import SparseTensor
+
+# class MinkowskiGRN(nn.Module):
+#     """ GRN layer for sparse tensors.
+#     """
+#     def __init__(self, dim):
+#         super().__init__()
+#         self.gamma = nn.Parameter(torch.zeros(1, dim))
+#         self.beta = nn.Parameter(torch.zeros(1, dim))
+
+#     def forward(self, x):
+#         cm = x.coordinate_manager
+#         in_key = x.coordinate_map_key
+
+#         Gx = torch.norm(x.F, p=2, dim=0, keepdim=True)
+#         Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+#         return SparseTensor(
+#                 self.gamma * (x.F * Nx) + self.beta + x.F,
+#                 coordinate_map_key=in_key,
+#                 coordinate_manager=cm)
+
+# class MinkowskiDropPath(nn.Module):
+#     """ Drop Path for sparse tensors.
+#     """
+
+#     def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+#         super(MinkowskiDropPath, self).__init__()
+#         self.drop_prob = drop_prob
+#         self.scale_by_keep = scale_by_keep
+    
+#     def forward(self, x):
+#         if self.drop_prob == 0. or not self.training:
+#             return x
+#         cm = x.coordinate_manager
+#         in_key = x.coordinate_map_key
+#         keep_prob = 1 - self.drop_prob
+#         mask = torch.cat([
+#             torch.ones(len(_)) if random.uniform(0, 1) > self.drop_prob
+#             else torch.zeros(len(_)) for _ in x.decomposed_coordinates
+#         ]).view(-1, 1).to(x.device)
+#         if keep_prob > 0.0 and self.scale_by_keep:
+#             mask.div_(keep_prob)
+#         return SparseTensor(
+#                 x.F * mask,
+#                 coordinate_map_key=in_key,
+#                 coordinate_manager=cm)
+
+# class MinkowskiLayerNorm(nn.Module):
+#     """ Channel-wise layer normalization for sparse tensors.
+#     """
+
+#     def __init__(
+#         self,
+#         normalized_shape,
+#         eps=1e-6,
+#     ):
+#         super(MinkowskiLayerNorm, self).__init__()
+#         self.ln = nn.LayerNorm(normalized_shape, eps=eps)
+#     def forward(self, input):
+#         output = self.ln(input.F)
+#         return SparseTensor(
+#             output,
+#             coordinate_map_key=input.coordinate_map_key,
+#             coordinate_manager=input.coordinate_manager)
+            
+class LayerNorm(nn.Module):
+    """ LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+    with shape (batch_size, channels, height, width).
+    """
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError 
+        self.normalized_shape = (normalized_shape, )
+    
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+class GRN(nn.Module):
+    """ GRN (Global Response Normalization) layer
+    """
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1,2), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x

clean/image/aide/optim_factory.py

@@ -0,0 +1,222 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+from torch import optim as optim
+
+from timm.optim.adafactor import Adafactor
+from timm.optim.adahessian import Adahessian
+from timm.optim.adamp import AdamP
+from timm.optim.lookahead import Lookahead
+from timm.optim.nadam import Nadam
+# from timm.optim.novograd import NovoGrad
+# from timm.optim.nvnovograd import NvNovoGrad
+from timm.optim.radam import RAdam
+from timm.optim.rmsprop_tf import RMSpropTF
+from timm.optim.sgdp import SGDP
+
+import json
+
+try:
+    from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+
+def get_num_layer_for_convnext_single(var_name, depths):
+    """
+    Each layer is assigned distinctive layer ids
+    """
+    if var_name.startswith("downsample_layers"):
+        stage_id = int(var_name.split('.')[1])
+        layer_id = sum(depths[:stage_id]) + 1
+        return layer_id
+    
+    elif var_name.startswith("stages"):
+        stage_id = int(var_name.split('.')[1])
+        block_id = int(var_name.split('.')[2])
+        layer_id = sum(depths[:stage_id]) + block_id + 1
+        return layer_id
+    
+    else:
+        return sum(depths) + 1
+
+
+def get_num_layer_for_convnext(var_name):
+    """
+    Divide [3, 3, 27, 3] layers into 12 groups; each group is three 
+    consecutive blocks, including possible neighboring downsample layers;
+    adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py
+    """
+    num_max_layer = 12
+    if var_name.startswith("downsample_layers"):
+        stage_id = int(var_name.split('.')[1])
+        if stage_id == 0:
+            layer_id = 0
+        elif stage_id == 1 or stage_id == 2:
+            layer_id = stage_id + 1
+        elif stage_id == 3:
+            layer_id = 12
+        return layer_id
+
+    elif var_name.startswith("stages"):
+        stage_id = int(var_name.split('.')[1])
+        block_id = int(var_name.split('.')[2])
+        if stage_id == 0 or stage_id == 1:
+            layer_id = stage_id + 1
+        elif stage_id == 2:
+            layer_id = 3 + block_id // 3 
+        elif stage_id == 3:
+            layer_id = 12
+        return layer_id
+    else:
+        return num_max_layer + 1
+
+class LayerDecayValueAssigner(object):
+    def __init__(self, values, depths=[3,3,27,3], layer_decay_type='single'):
+        self.values = values
+        self.depths = depths
+        self.layer_decay_type = layer_decay_type
+
+    def get_scale(self, layer_id):
+        return self.values[layer_id]
+
+    def get_layer_id(self, var_name):
+        if self.layer_decay_type == 'single':
+            return get_num_layer_for_convnext_single(var_name, self.depths)
+        else:
+            return get_num_layer_for_convnext(var_name)
+
+
+def get_parameter_groups(model, weight_decay=1e-5, skip_list=(), get_num_layer=None, get_layer_scale=None):
+    parameter_group_names = {}
+    parameter_group_vars = {}
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list or \
+            name.endswith(".gamma") or name.endswith(".beta"):
+            group_name = "no_decay"
+            this_weight_decay = 0.
+        else:
+            group_name = "decay"
+            this_weight_decay = weight_decay
+        if get_num_layer is not None:
+            layer_id = get_num_layer(name)
+            group_name = "layer_%d_%s" % (layer_id, group_name)
+        else:
+            layer_id = None
+
+        if group_name not in parameter_group_names:
+            if get_layer_scale is not None:
+                scale = get_layer_scale(layer_id)
+            else:
+                scale = 1.
+
+            parameter_group_names[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale
+            }
+            parameter_group_vars[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale
+            }
+
+        parameter_group_vars[group_name]["params"].append(param)
+        parameter_group_names[group_name]["params"].append(name)
+    print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))
+    return list(parameter_group_vars.values())
+
+
+def create_optimizer(args, model, get_num_layer=None, get_layer_scale=None, filter_bias_and_bn=True, skip_list=None):
+    opt_lower = args.opt.lower()
+    weight_decay = args.weight_decay
+    # if weight_decay and filter_bias_and_bn:
+    if filter_bias_and_bn:
+        skip = {}
+        if skip_list is not None:
+            skip = skip_list
+        elif hasattr(model, 'no_weight_decay'):
+            skip = model.no_weight_decay()
+        parameters = get_parameter_groups(model, weight_decay, skip, get_num_layer, get_layer_scale)
+        weight_decay = 0.
+    else:
+        parameters = model.parameters()
+
+    if 'fused' in opt_lower:
+        assert has_apex and torch.cuda.is_available(), 'APEX and CUDA required for fused optimizers'
+
+    opt_args = dict(lr=args.lr, weight_decay=weight_decay)
+    if hasattr(args, 'opt_eps') and args.opt_eps is not None:
+        opt_args['eps'] = args.opt_eps
+    if hasattr(args, 'opt_betas') and args.opt_betas is not None:
+        opt_args['betas'] = args.opt_betas
+
+    opt_split = opt_lower.split('_')
+    opt_lower = opt_split[-1]
+    if opt_lower == 'sgd' or opt_lower == 'nesterov':
+        opt_args.pop('eps', None)
+        optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'momentum':
+        opt_args.pop('eps', None)
+        optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
+    elif opt_lower == 'adam':
+        optimizer = optim.Adam(parameters, **opt_args)
+    elif opt_lower == 'adamw':
+        optimizer = optim.AdamW(parameters, **opt_args)
+    elif opt_lower == 'nadam':
+        optimizer = Nadam(parameters, **opt_args)
+    elif opt_lower == 'radam':
+        optimizer = RAdam(parameters, **opt_args)
+    elif opt_lower == 'adamp':
+        optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args)
+    elif opt_lower == 'sgdp':
+        optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'adadelta':
+        optimizer = optim.Adadelta(parameters, **opt_args)
+    elif opt_lower == 'adafactor':
+        if not args.lr:
+            opt_args['lr'] = None
+        optimizer = Adafactor(parameters, **opt_args)
+    elif opt_lower == 'adahessian':
+        optimizer = Adahessian(parameters, **opt_args)
+    elif opt_lower == 'rmsprop':
+        optimizer = optim.RMSprop(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+    elif opt_lower == 'rmsproptf':
+        optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+    elif opt_lower == 'novograd':
+        optimizer = NovoGrad(parameters, **opt_args)
+    elif opt_lower == 'nvnovograd':
+        optimizer = NvNovoGrad(parameters, **opt_args)
+    elif opt_lower == 'fusedsgd':
+        opt_args.pop('eps', None)
+        optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'fusedmomentum':
+        opt_args.pop('eps', None)
+        optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
+    elif opt_lower == 'fusedadam':
+        optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args)
+    elif opt_lower == 'fusedadamw':
+        optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args)
+    elif opt_lower == 'fusedlamb':
+        optimizer = FusedLAMB(parameters, **opt_args)
+    elif opt_lower == 'fusednovograd':
+        opt_args.setdefault('betas', (0.95, 0.98))
+        optimizer = FusedNovoGrad(parameters, **opt_args)
+    else:
+        assert False and "Invalid optimizer"
+
+    if len(opt_split) > 1:
+        if opt_split[0] == 'lookahead':
+            optimizer = Lookahead(optimizer)
+
+    return optimizer

clean/image/aide/requirements.txt

@@ -0,0 +1,37 @@
+einops==0.6.1
+fairscale==0.4.13
+filelock==3.13.1
+ftfy==6.1.3
+h5py==3.10.0
+imgaug==0.2.6
+keras==2.11.0
+kornia==0.7.2
+kornia_rs==0.1.2
+lmdb==1.4.1
+matplotlib==3.7.4
+matplotlib-inline==0.1.6
+numpy==1.24.3
+omegaconf==2.3.0
+open-clip-torch==2.24.0
+openai-clip==1.0.1
+openpyxl==3.1.2
+pandas==2.0.3
+Pillow==9.5.0
+safetensors==0.4.1
+scikit-image==0.20.0
+scikit-learn==1.3.2
+scipy==1.9.1
+sentencepiece==0.2.0
+streamlit==1.30.0
+tenacity==8.2.3
+tensorboard==2.11.2
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+tensorboardX==2.6.2.2
+timm==0.9.6
+torch==1.11.0
+torch-fidelity==0.3.0
+torchmetrics==0.6.0
+torchsummary==1.5.1
+torchvision==0.12.0
+tqdm==4.66.1