wavedl 1.5.7__tar.gz → 1.6.0__tar.gz

{wavedl-1.5.7/src/wavedl.egg-info → wavedl-1.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.5.7
+Version: 1.6.0
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -23,6 +23,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
 Requires-Dist: torchvision>=0.15.0
+Requires-Dist: timm>=0.9.0
 Requires-Dist: accelerate>=0.20.0
 Requires-Dist: numpy>=1.24.0
 Requires-Dist: scipy>=1.10.0
@@ -117,7 +118,7 @@ Train on datasets larger than RAM:
 
 **🧠 Models? We've Got Options**
 
-38 architectures, ready to go:
+57 architectures, ready to go:
 - CNNs, ResNets, ViTs, EfficientNets...
 - All adapted for regression
 - [Add your own](#adding-custom-models) in one line
@@ -202,7 +203,7 @@ Deploy models anywhere:
 #### From PyPI (recommended for all users)
 
 ```bash
-pip install wavedl
+pip install --upgrade wavedl
 ```
 
 This installs everything you need: training, inference, HPO, ONNX export.
@@ -358,22 +359,10 @@ WaveDL/
 │       ├── hpo.py                # Hyperparameter optimization
 │       ├── hpc.py                # HPC distributed training launcher
 │       │
-│       ├── models/               # Model architectures (38 variants)
+│       ├── models/               # Model Zoo (57 architectures)
 │       │   ├── registry.py       # Model factory (@register_model)
 │       │   ├── base.py           # Abstract base class
-│       │   ├── cnn.py            # Baseline CNN (1D/2D/3D)
-│       │   ├── resnet.py         # ResNet-18/34/50 (1D/2D/3D)
-│       │   ├── resnet3d.py       # ResNet3D-18, MC3-18 (3D only)
-│       │   ├── tcn.py            # TCN (1D only)
-│       │   ├── efficientnet.py   # EfficientNet-B0/B1/B2 (2D)
-│       │   ├── efficientnetv2.py # EfficientNetV2-S/M/L (2D)
-│       │   ├── mobilenetv3.py    # MobileNetV3-Small/Large (2D)
-│       │   ├── regnet.py         # RegNetY variants (2D)
-│       │   ├── swin.py           # Swin Transformer (2D)
-│       │   ├── vit.py            # Vision Transformer (1D/2D)
-│       │   ├── convnext.py       # ConvNeXt (1D/2D/3D)
-│       │   ├── densenet.py       # DenseNet-121/169 (1D/2D/3D)
-│       │   └── unet.py           # U-Net Regression
+│       │   └── ...               # See "Available Models" section
 │       │
 │       └── utils/                # Utilities
 │           ├── data.py           # Memory-mapped data pipeline
@@ -388,7 +377,7 @@ WaveDL/
 ├── configs/                      # YAML config templates
 ├── examples/                     # Ready-to-run examples
 ├── notebooks/                    # Jupyter notebooks
-├── unit_tests/                   # Pytest test suite (903 tests)
+├── unit_tests/                   # Pytest test suite
 │
 ├── pyproject.toml                # Package config, dependencies
 ├── CHANGELOG.md                  # Version history
@@ -411,71 +400,96 @@ WaveDL/
 > ```
 
 <details>
-<summary><b>Available Models</b> — 38 architectures</summary>
+<summary><b>Available Models</b> — 57 architectures</summary>
 
-| Model | Params | Dim |
-|-------|--------|-----|
+| Model | Backbone Params | Dim |
+|-------|-----------------|-----|
 | **CNN** — Convolutional Neural Network |||
-| `cnn` | 1.7M | 1D/2D/3D |
+| `cnn` | 1.6M | 1D/2D/3D |
 | **ResNet** — Residual Network |||
-| `resnet18` | 11.4M | 1D/2D/3D |
-| `resnet34` | 21.5M | 1D/2D/3D |
-| `resnet50` | 24.6M | 1D/2D/3D |
-| `resnet18_pretrained` ⭐ | 11.4M | 2D |
-| `resnet50_pretrained` ⭐ | 24.6M | 2D |
+| `resnet18` | 11.2M | 1D/2D/3D |
+| `resnet34` | 21.3M | 1D/2D/3D |
+| `resnet50` | 23.5M | 1D/2D/3D |
+| `resnet18_pretrained` ⭐ | 11.2M | 2D |
+| `resnet50_pretrained` ⭐ | 23.5M | 2D |
 | **ResNet3D** — 3D Residual Network |||
-| `resnet3d_18` | 33.6M | 3D |
-| `mc3_18` — Mixed Convolution 3D | 11.9M | 3D |
+| `resnet3d_18` | 33.2M | 3D |
+| `mc3_18` — Mixed Convolution 3D | 11.5M | 3D |
 | **TCN** — Temporal Convolutional Network |||
-| `tcn_small` | 1.0M | 1D |
-| `tcn` | 7.0M | 1D |
-| `tcn_large` | 10.2M | 1D |
+| `tcn_small` | 0.9M | 1D |
+| `tcn` | 6.9M | 1D |
+| `tcn_large` | 10.0M | 1D |
 | **EfficientNet** — Efficient Neural Network |||
-| `efficientnet_b0` ⭐ | 4.7M | 2D |
-| `efficientnet_b1` ⭐ | 7.2M | 2D |
-| `efficientnet_b2` ⭐ | 8.4M | 2D |
+| `efficientnet_b0` ⭐ | 4.0M | 2D |
+| `efficientnet_b1` ⭐ | 6.5M | 2D |
+| `efficientnet_b2` ⭐ | 7.7M | 2D |
 | **EfficientNetV2** — Efficient Neural Network V2 |||
-| `efficientnet_v2_s` ⭐ | 21.0M | 2D |
-| `efficientnet_v2_m` ⭐ | 53.6M | 2D |
-| `efficientnet_v2_l` ⭐ | 118.0M | 2D |
+| `efficientnet_v2_s` ⭐ | 20.2M | 2D |
+| `efficientnet_v2_m` ⭐ | 52.9M | 2D |
+| `efficientnet_v2_l` ⭐ | 117.2M | 2D |
 | **MobileNetV3** — Mobile Neural Network V3 |||
-| `mobilenet_v3_small` ⭐ | 1.1M | 2D |
-| `mobilenet_v3_large` ⭐ | 3.2M | 2D |
+| `mobilenet_v3_small` ⭐ | 0.9M | 2D |
+| `mobilenet_v3_large` ⭐ | 3.0M | 2D |
 | **RegNet** — Regularized Network |||
-| `regnet_y_400mf` ⭐ | 4.0M | 2D |
-| `regnet_y_800mf` ⭐ | 5.8M | 2D |
-| `regnet_y_1_6gf` ⭐ | 10.5M | 2D |
-| `regnet_y_3_2gf` ⭐ | 18.3M | 2D |
-| `regnet_y_8gf` ⭐ | 37.9M | 2D |
+| `regnet_y_400mf` ⭐ | 3.9M | 2D |
+| `regnet_y_800mf` ⭐ | 5.7M | 2D |
+| `regnet_y_1_6gf` ⭐ | 10.3M | 2D |
+| `regnet_y_3_2gf` ⭐ | 17.9M | 2D |
+| `regnet_y_8gf` ⭐ | 37.4M | 2D |
 | **Swin** — Shifted Window Transformer |||
-| `swin_t` ⭐ | 28.0M | 2D |
-| `swin_s` ⭐ | 49.4M | 2D |
-| `swin_b` ⭐ | 87.4M | 2D |
+| `swin_t` ⭐ | 27.5M | 2D |
+| `swin_s` ⭐ | 48.8M | 2D |
+| `swin_b` ⭐ | 86.7M | 2D |
 | **ConvNeXt** — Convolutional Next |||
-| `convnext_tiny` | 28.2M | 1D/2D/3D |
-| `convnext_small` | 49.8M | 1D/2D/3D |
-| `convnext_base` | 88.1M | 1D/2D/3D |
-| `convnext_tiny_pretrained` ⭐ | 28.2M | 2D |
+| `convnext_tiny` | 27.8M | 1D/2D/3D |
+| `convnext_small` | 49.5M | 1D/2D/3D |
+| `convnext_base` | 87.6M | 1D/2D/3D |
+| `convnext_tiny_pretrained` ⭐ | 27.8M | 2D |
 | **DenseNet** — Densely Connected Network |||
-| `densenet121` | 7.5M | 1D/2D/3D |
-| `densenet169` | 13.3M | 1D/2D/3D |
-| `densenet121_pretrained` ⭐ | 7.5M | 2D |
+| `densenet121` | 7.0M | 1D/2D/3D |
+| `densenet169` | 12.5M | 1D/2D/3D |
+| `densenet121_pretrained` ⭐ | 7.0M | 2D |
 | **ViT** — Vision Transformer |||
-| `vit_tiny` | 5.5M | 1D/2D |
-| `vit_small` | 21.6M | 1D/2D |
-| `vit_base` | 85.6M | 1D/2D |
+| `vit_tiny` | 5.4M | 1D/2D |
+| `vit_small` | 21.4M | 1D/2D |
+| `vit_base` | 85.3M | 1D/2D |
+| **ConvNeXt V2** — ConvNeXt with GRN |||
+| `convnext_v2_tiny` | 27.9M | 1D/2D/3D |
+| `convnext_v2_small` | 49.6M | 1D/2D/3D |
+| `convnext_v2_base` | 87.7M | 1D/2D/3D |
+| `convnext_v2_tiny_pretrained` ⭐ | 27.9M | 2D |
+| **Mamba** — State Space Model |||
+| `mamba_1d` | 3.4M | 1D |
+| **Vision Mamba (ViM)** — 2D Mamba |||
+| `vim_tiny` | 6.6M | 2D |
+| `vim_small` | 51.1M | 2D |
+| `vim_base` | 201.4M | 2D |
+| **MaxViT** — Multi-Axis ViT |||
+| `maxvit_tiny` ⭐ | 30.1M | 2D |
+| `maxvit_small` ⭐ | 67.6M | 2D |
+| `maxvit_base` ⭐ | 119.1M | 2D |
+| **FastViT** — Fast Hybrid CNN-ViT |||
+| `fastvit_t8` ⭐ | 4.0M | 2D |
+| `fastvit_t12` ⭐ | 6.8M | 2D |
+| `fastvit_s12` ⭐ | 8.8M | 2D |
+| `fastvit_sa12` ⭐ | 10.9M | 2D |
+| **CAFormer** — MetaFormer with Attention |||
+| `caformer_s18` ⭐ | 26.3M | 2D |
+| `caformer_s36` ⭐ | 39.2M | 2D |
+| `caformer_m36` ⭐ | 56.9M | 2D |
+| `poolformer_s12` ⭐ | 11.9M | 2D |
 | **U-Net** — U-shaped Network |||
-| `unet_regression` | 31.1M | 1D/2D/3D |
+| `unet_regression` | 31.0M | 1D/2D/3D |
+
 
 ⭐ = **Pretrained on ImageNet** (recommended for smaller datasets). Weights are downloaded automatically on first use.
 - **Cache location**: `~/.cache/torch/hub/checkpoints/` (or `./.torch_cache/` on HPC if home is not writable)
-- **Size**: ~20–350 MB per model depending on architecture
 - **Train from scratch**: Use `--no_pretrained` to disable pretrained weights
 
 **💡 HPC Users**: If compute nodes block internet, pre-download weights on the login node:
 
 ```bash
-# Run once on login node (with internet) — downloads ALL pretrained weights (~1.5 GB total)
+# Run once on login node (with internet) — downloads ALL pretrained weights
 python -c "
 import os
 os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
@@ -483,7 +497,7 @@ os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
 from torchvision import models as m
 from torchvision.models import video as v
 
-# Model name -> Weights class mapping
+# === TorchVision Models ===
 weights = {
     'resnet18': m.ResNet18_Weights, 'resnet50': m.ResNet50_Weights,
     'efficientnet_b0': m.EfficientNet_B0_Weights, 'efficientnet_b1': m.EfficientNet_B1_Weights,
@@ -501,6 +515,20 @@ for name, w in weights.items():
 # 3D video models
 v.r3d_18(weights=v.R3D_18_Weights.DEFAULT); print('✓ r3d_18')
 v.mc3_18(weights=v.MC3_18_Weights.DEFAULT); print('✓ mc3_18')
+
+# === Timm Models (MaxViT, FastViT, CAFormer, ConvNeXt V2) ===
+import timm
+
+timm_models = [
+    'maxvit_tiny_tf_224.in1k', 'maxvit_small_tf_224.in1k', 'maxvit_base_tf_224.in1k',
+    'fastvit_t8.apple_in1k', 'fastvit_t12.apple_in1k', 'fastvit_s12.apple_in1k', 'fastvit_sa12.apple_in1k',
+    'caformer_s18.sail_in1k', 'caformer_s36.sail_in22k_ft_in1k', 'caformer_m36.sail_in22k_ft_in1k',
+    'poolformer_s12.sail_in1k',
+    'convnextv2_tiny.fcmae_ft_in1k',
+]
+for name in timm_models:
+    timm.create_model(name, pretrained=True); print(f'✓ {name}')
+
 print('\\n✓ All pretrained weights cached!')
 "
 ```

{wavedl-1.5.7 → wavedl-1.6.0}/README.md

@@ -71,7 +71,7 @@ Train on datasets larger than RAM:
 
 **🧠 Models? We've Got Options**
 
-38 architectures, ready to go:
+57 architectures, ready to go:
 - CNNs, ResNets, ViTs, EfficientNets...
 - All adapted for regression
 - [Add your own](#adding-custom-models) in one line
@@ -156,7 +156,7 @@ Deploy models anywhere:
 #### From PyPI (recommended for all users)
 
 ```bash
-pip install wavedl
+pip install --upgrade wavedl
 ```
 
 This installs everything you need: training, inference, HPO, ONNX export.
@@ -312,22 +312,10 @@ WaveDL/
 │       ├── hpo.py                # Hyperparameter optimization
 │       ├── hpc.py                # HPC distributed training launcher
 │       │
-│       ├── models/               # Model architectures (38 variants)
+│       ├── models/               # Model Zoo (57 architectures)
 │       │   ├── registry.py       # Model factory (@register_model)
 │       │   ├── base.py           # Abstract base class
-│       │   ├── cnn.py            # Baseline CNN (1D/2D/3D)
-│       │   ├── resnet.py         # ResNet-18/34/50 (1D/2D/3D)
-│       │   ├── resnet3d.py       # ResNet3D-18, MC3-18 (3D only)
-│       │   ├── tcn.py            # TCN (1D only)
-│       │   ├── efficientnet.py   # EfficientNet-B0/B1/B2 (2D)
-│       │   ├── efficientnetv2.py # EfficientNetV2-S/M/L (2D)
-│       │   ├── mobilenetv3.py    # MobileNetV3-Small/Large (2D)
-│       │   ├── regnet.py         # RegNetY variants (2D)
-│       │   ├── swin.py           # Swin Transformer (2D)
-│       │   ├── vit.py            # Vision Transformer (1D/2D)
-│       │   ├── convnext.py       # ConvNeXt (1D/2D/3D)
-│       │   ├── densenet.py       # DenseNet-121/169 (1D/2D/3D)
-│       │   └── unet.py           # U-Net Regression
+│       │   └── ...               # See "Available Models" section
 │       │
 │       └── utils/                # Utilities
 │           ├── data.py           # Memory-mapped data pipeline
@@ -342,7 +330,7 @@ WaveDL/
 ├── configs/                      # YAML config templates
 ├── examples/                     # Ready-to-run examples
 ├── notebooks/                    # Jupyter notebooks
-├── unit_tests/                   # Pytest test suite (903 tests)
+├── unit_tests/                   # Pytest test suite
 │
 ├── pyproject.toml                # Package config, dependencies
 ├── CHANGELOG.md                  # Version history
@@ -365,71 +353,96 @@ WaveDL/
 > ```
 
 <details>
-<summary><b>Available Models</b> — 38 architectures</summary>
+<summary><b>Available Models</b> — 57 architectures</summary>
 
-| Model | Params | Dim |
-|-------|--------|-----|
+| Model | Backbone Params | Dim |
+|-------|-----------------|-----|
 | **CNN** — Convolutional Neural Network |||
-| `cnn` | 1.7M | 1D/2D/3D |
+| `cnn` | 1.6M | 1D/2D/3D |
 | **ResNet** — Residual Network |||
-| `resnet18` | 11.4M | 1D/2D/3D |
-| `resnet34` | 21.5M | 1D/2D/3D |
-| `resnet50` | 24.6M | 1D/2D/3D |
-| `resnet18_pretrained` ⭐ | 11.4M | 2D |
-| `resnet50_pretrained` ⭐ | 24.6M | 2D |
+| `resnet18` | 11.2M | 1D/2D/3D |
+| `resnet34` | 21.3M | 1D/2D/3D |
+| `resnet50` | 23.5M | 1D/2D/3D |
+| `resnet18_pretrained` ⭐ | 11.2M | 2D |
+| `resnet50_pretrained` ⭐ | 23.5M | 2D |
 | **ResNet3D** — 3D Residual Network |||
-| `resnet3d_18` | 33.6M | 3D |
-| `mc3_18` — Mixed Convolution 3D | 11.9M | 3D |
+| `resnet3d_18` | 33.2M | 3D |
+| `mc3_18` — Mixed Convolution 3D | 11.5M | 3D |
 | **TCN** — Temporal Convolutional Network |||
-| `tcn_small` | 1.0M | 1D |
-| `tcn` | 7.0M | 1D |
-| `tcn_large` | 10.2M | 1D |
+| `tcn_small` | 0.9M | 1D |
+| `tcn` | 6.9M | 1D |
+| `tcn_large` | 10.0M | 1D |
 | **EfficientNet** — Efficient Neural Network |||
-| `efficientnet_b0` ⭐ | 4.7M | 2D |
-| `efficientnet_b1` ⭐ | 7.2M | 2D |
-| `efficientnet_b2` ⭐ | 8.4M | 2D |
+| `efficientnet_b0` ⭐ | 4.0M | 2D |
+| `efficientnet_b1` ⭐ | 6.5M | 2D |
+| `efficientnet_b2` ⭐ | 7.7M | 2D |
 | **EfficientNetV2** — Efficient Neural Network V2 |||
-| `efficientnet_v2_s` ⭐ | 21.0M | 2D |
-| `efficientnet_v2_m` ⭐ | 53.6M | 2D |
-| `efficientnet_v2_l` ⭐ | 118.0M | 2D |
+| `efficientnet_v2_s` ⭐ | 20.2M | 2D |
+| `efficientnet_v2_m` ⭐ | 52.9M | 2D |
+| `efficientnet_v2_l` ⭐ | 117.2M | 2D |
 | **MobileNetV3** — Mobile Neural Network V3 |||
-| `mobilenet_v3_small` ⭐ | 1.1M | 2D |
-| `mobilenet_v3_large` ⭐ | 3.2M | 2D |
+| `mobilenet_v3_small` ⭐ | 0.9M | 2D |
+| `mobilenet_v3_large` ⭐ | 3.0M | 2D |
 | **RegNet** — Regularized Network |||
-| `regnet_y_400mf` ⭐ | 4.0M | 2D |
-| `regnet_y_800mf` ⭐ | 5.8M | 2D |
-| `regnet_y_1_6gf` ⭐ | 10.5M | 2D |
-| `regnet_y_3_2gf` ⭐ | 18.3M | 2D |
-| `regnet_y_8gf` ⭐ | 37.9M | 2D |
+| `regnet_y_400mf` ⭐ | 3.9M | 2D |
+| `regnet_y_800mf` ⭐ | 5.7M | 2D |
+| `regnet_y_1_6gf` ⭐ | 10.3M | 2D |
+| `regnet_y_3_2gf` ⭐ | 17.9M | 2D |
+| `regnet_y_8gf` ⭐ | 37.4M | 2D |
 | **Swin** — Shifted Window Transformer |||
-| `swin_t` ⭐ | 28.0M | 2D |
-| `swin_s` ⭐ | 49.4M | 2D |
-| `swin_b` ⭐ | 87.4M | 2D |
+| `swin_t` ⭐ | 27.5M | 2D |
+| `swin_s` ⭐ | 48.8M | 2D |
+| `swin_b` ⭐ | 86.7M | 2D |
 | **ConvNeXt** — Convolutional Next |||
-| `convnext_tiny` | 28.2M | 1D/2D/3D |
-| `convnext_small` | 49.8M | 1D/2D/3D |
-| `convnext_base` | 88.1M | 1D/2D/3D |
-| `convnext_tiny_pretrained` ⭐ | 28.2M | 2D |
+| `convnext_tiny` | 27.8M | 1D/2D/3D |
+| `convnext_small` | 49.5M | 1D/2D/3D |
+| `convnext_base` | 87.6M | 1D/2D/3D |
+| `convnext_tiny_pretrained` ⭐ | 27.8M | 2D |
 | **DenseNet** — Densely Connected Network |||
-| `densenet121` | 7.5M | 1D/2D/3D |
-| `densenet169` | 13.3M | 1D/2D/3D |
-| `densenet121_pretrained` ⭐ | 7.5M | 2D |
+| `densenet121` | 7.0M | 1D/2D/3D |
+| `densenet169` | 12.5M | 1D/2D/3D |
+| `densenet121_pretrained` ⭐ | 7.0M | 2D |
 | **ViT** — Vision Transformer |||
-| `vit_tiny` | 5.5M | 1D/2D |
-| `vit_small` | 21.6M | 1D/2D |
-| `vit_base` | 85.6M | 1D/2D |
+| `vit_tiny` | 5.4M | 1D/2D |
+| `vit_small` | 21.4M | 1D/2D |
+| `vit_base` | 85.3M | 1D/2D |
+| **ConvNeXt V2** — ConvNeXt with GRN |||
+| `convnext_v2_tiny` | 27.9M | 1D/2D/3D |
+| `convnext_v2_small` | 49.6M | 1D/2D/3D |
+| `convnext_v2_base` | 87.7M | 1D/2D/3D |
+| `convnext_v2_tiny_pretrained` ⭐ | 27.9M | 2D |
+| **Mamba** — State Space Model |||
+| `mamba_1d` | 3.4M | 1D |
+| **Vision Mamba (ViM)** — 2D Mamba |||
+| `vim_tiny` | 6.6M | 2D |
+| `vim_small` | 51.1M | 2D |
+| `vim_base` | 201.4M | 2D |
+| **MaxViT** — Multi-Axis ViT |||
+| `maxvit_tiny` ⭐ | 30.1M | 2D |
+| `maxvit_small` ⭐ | 67.6M | 2D |
+| `maxvit_base` ⭐ | 119.1M | 2D |
+| **FastViT** — Fast Hybrid CNN-ViT |||
+| `fastvit_t8` ⭐ | 4.0M | 2D |
+| `fastvit_t12` ⭐ | 6.8M | 2D |
+| `fastvit_s12` ⭐ | 8.8M | 2D |
+| `fastvit_sa12` ⭐ | 10.9M | 2D |
+| **CAFormer** — MetaFormer with Attention |||
+| `caformer_s18` ⭐ | 26.3M | 2D |
+| `caformer_s36` ⭐ | 39.2M | 2D |
+| `caformer_m36` ⭐ | 56.9M | 2D |
+| `poolformer_s12` ⭐ | 11.9M | 2D |
 | **U-Net** — U-shaped Network |||
-| `unet_regression` | 31.1M | 1D/2D/3D |
+| `unet_regression` | 31.0M | 1D/2D/3D |
+
 
 ⭐ = **Pretrained on ImageNet** (recommended for smaller datasets). Weights are downloaded automatically on first use.
 - **Cache location**: `~/.cache/torch/hub/checkpoints/` (or `./.torch_cache/` on HPC if home is not writable)
-- **Size**: ~20–350 MB per model depending on architecture
 - **Train from scratch**: Use `--no_pretrained` to disable pretrained weights
 
 **💡 HPC Users**: If compute nodes block internet, pre-download weights on the login node:
 
 ```bash
-# Run once on login node (with internet) — downloads ALL pretrained weights (~1.5 GB total)
+# Run once on login node (with internet) — downloads ALL pretrained weights
 python -c "
 import os
 os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
@@ -437,7 +450,7 @@ os.environ['TORCH_HOME'] = '.torch_cache'  # Match WaveDL's HPC cache location
 from torchvision import models as m
 from torchvision.models import video as v
 
-# Model name -> Weights class mapping
+# === TorchVision Models ===
 weights = {
     'resnet18': m.ResNet18_Weights, 'resnet50': m.ResNet50_Weights,
     'efficientnet_b0': m.EfficientNet_B0_Weights, 'efficientnet_b1': m.EfficientNet_B1_Weights,
@@ -455,6 +468,20 @@ for name, w in weights.items():
 # 3D video models
 v.r3d_18(weights=v.R3D_18_Weights.DEFAULT); print('✓ r3d_18')
 v.mc3_18(weights=v.MC3_18_Weights.DEFAULT); print('✓ mc3_18')
+
+# === Timm Models (MaxViT, FastViT, CAFormer, ConvNeXt V2) ===
+import timm
+
+timm_models = [
+    'maxvit_tiny_tf_224.in1k', 'maxvit_small_tf_224.in1k', 'maxvit_base_tf_224.in1k',
+    'fastvit_t8.apple_in1k', 'fastvit_t12.apple_in1k', 'fastvit_s12.apple_in1k', 'fastvit_sa12.apple_in1k',
+    'caformer_s18.sail_in1k', 'caformer_s36.sail_in22k_ft_in1k', 'caformer_m36.sail_in22k_ft_in1k',
+    'poolformer_s12.sail_in1k',
+    'convnextv2_tiny.fcmae_ft_in1k',
+]
+for name in timm_models:
+    timm.create_model(name, pretrained=True); print(f'✓ {name}')
+
 print('\\n✓ All pretrained weights cached!')
 "
 ```

{wavedl-1.5.7 → wavedl-1.6.0}/pyproject.toml

@@ -52,6 +52,7 @@ dependencies = [
     # Core ML stack
     "torch>=2.0.0",
     "torchvision>=0.15.0",
+    "timm>=0.9.0",  # Pretrained models (MaxViT, FastViT, CAFormer)
     "accelerate>=0.20.0",
     "numpy>=1.24.0",
     "scipy>=1.10.0",

{wavedl-1.5.7 → wavedl-1.6.0}/src/wavedl/__init__.py

@@ -18,7 +18,7 @@ For inference:
     # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
 """
 
-__version__ = "1.5.7"
+__version__ = "1.6.0"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

{wavedl-1.5.7 → wavedl-1.6.0}/src/wavedl/models/__init__.py

@@ -6,10 +6,11 @@ This module provides a centralized registry for neural network architectures,
 enabling dynamic model selection via command-line arguments.
 
 **Dimensionality Coverage**:
-    - 1D (waveforms): TCN, CNN, ResNet, ConvNeXt, DenseNet, ViT
-    - 2D (images): CNN, ResNet, ConvNeXt, DenseNet, ViT, UNet,
-                   EfficientNet, MobileNetV3, RegNet, Swin
-    - 3D (volumes): ResNet3D, CNN, ResNet, ConvNeXt, DenseNet
+    - 1D (waveforms): TCN, CNN, ResNet, ConvNeXt, ConvNeXt V2, DenseNet, ViT, Mamba
+    - 2D (images): CNN, ResNet, ConvNeXt, ConvNeXt V2, DenseNet, ViT, UNet,
+                   EfficientNet, MobileNetV3, RegNet, Swin, MaxViT, FastViT,
+                   CAFormer, PoolFormer, Vision Mamba
+    - 3D (volumes): ResNet3D, CNN, ResNet, ConvNeXt, ConvNeXt V2, DenseNet
 
 Usage:
     from wavedl.models import get_model, list_models, MODEL_REGISTRY
@@ -46,9 +47,19 @@ from .base import BaseModel
 # Import model implementations (triggers registration via decorators)
 from .cnn import CNN
 from .convnext import ConvNeXtBase_, ConvNeXtSmall, ConvNeXtTiny
+
+# New models (v1.6+)
+from .convnext_v2 import (
+    ConvNeXtV2Base,
+    ConvNeXtV2BaseLarge,
+    ConvNeXtV2Small,
+    ConvNeXtV2Tiny,
+    ConvNeXtV2TinyPretrained,
+)
 from .densenet import DenseNet121, DenseNet169
 from .efficientnet import EfficientNetB0, EfficientNetB1, EfficientNetB2
 from .efficientnetv2 import EfficientNetV2L, EfficientNetV2M, EfficientNetV2S
+from .mamba import Mamba1D, VimBase, VimSmall, VimTiny
 from .mobilenetv3 import MobileNetV3Large, MobileNetV3Small
 from .registry import (
     MODEL_REGISTRY,
@@ -66,6 +77,17 @@ from .unet import UNetRegression
 from .vit import ViTBase_, ViTSmall, ViTTiny
 
 
+# Optional timm-based models (imported conditionally)
+try:
+    from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
+    from .fastvit import FastViTS12, FastViTSA12, FastViTT8, FastViTT12
+    from .maxvit import MaxViTBaseLarge, MaxViTSmall, MaxViTTiny
+
+    _HAS_TIMM_MODELS = True
+except ImportError:
+    _HAS_TIMM_MODELS = False
+
+
 # Export public API (sorted alphabetically per RUF022)
 # See module docstring for dimensionality support details
 __all__ = [
@@ -77,6 +99,11 @@ __all__ = [
     "ConvNeXtBase_",
     "ConvNeXtSmall",
     "ConvNeXtTiny",
+    "ConvNeXtV2Base",
+    "ConvNeXtV2BaseLarge",
+    "ConvNeXtV2Small",
+    "ConvNeXtV2Tiny",
+    "ConvNeXtV2TinyPretrained",
     "DenseNet121",
     "DenseNet169",
     "EfficientNetB0",
@@ -85,6 +112,7 @@ __all__ = [
     "EfficientNetV2L",
     "EfficientNetV2M",
     "EfficientNetV2S",
+    "Mamba1D",
     "MobileNetV3Large",
     "MobileNetV3Small",
     "RegNetY1_6GF",
@@ -105,8 +133,28 @@ __all__ = [
     "ViTBase_",
     "ViTSmall",
     "ViTTiny",
+    "VimBase",
+    "VimSmall",
+    "VimTiny",
     "build_model",
     "get_model",
     "list_models",
     "register_model",
 ]
+
+# Add timm-based models to __all__ if available
+if _HAS_TIMM_MODELS:
+    __all__.extend(
+        [
+            "CaFormerS18",
+            "CaFormerS36",
+            "FastViTS12",
+            "FastViTSA12",
+            "FastViTT8",
+            "FastViTT12",
+            "MaxViTBaseLarge",
+            "MaxViTSmall",
+            "MaxViTTiny",
+            "PoolFormerS12",
+        ]
+    )