### WaterMamba YAML Configuration Example

Source: https://context7.com/guan-ms/watermamba/llms.txt

A sample YAML configuration file for training the WaterMamba model, detailing network architecture, datasets, optimizer, scheduler, and validation settings.

```yaml
# Mamba/Options/WaterMamba.yml
name: WaterMamba
model_type: ImageCleanModel
scale: 1
num_gpu: 8
manual_seed: 100

# Model architecture
network_g:
  type: WaterMamba
  inp_channels: 3
  out_channels: 3
  dim: 24
  num_blocks: [2, 2, 2, 2]
  mam_blocks: [2, 2, 2, 2]
  num_refinement_blocks: 4
  d_state: 16
  bias: False
  dual_pixel_task: False

# Training dataset (paired LQ/GT images)
datasets:
  train:
    name: TrainSet
    type: Dataset_PairedImage
    dataroot_gt: /path/to/train/GT
    dataroot_lq: /path/to/train/LQ
    geometric_augs: true
    filename_tmpl: '{}'
    io_backend:
      type: disk
    # Progressive patch-size training schedule
    train_size: [128, 160, 192, 256]
    batch_size_per_gpu: 8
    num_worker_per_gpu: 4
    use_shuffle: true

  val:
    name: ValSet
    type: Dataset_PairedImage
    dataroot_gt: /path/to/val/GT
    dataroot_lq: /path/to/val/LQ
    io_backend:
      type: disk

# Path configuration
path:
  pretrain_network_g: ~
  resume_state: ~

# Training schedule
train:
  total_iter: 300000
  warmup_iter: -1
  use_grad_clip: true
  scheduler:
    type: CosineAnnealingRestartCyclicLR
    periods: [92000, 208000]
    restart_weights: [1, 1]
    eta_mins: [0.0003, 0.000001]
  optim_g:
    type: AdamW
    lr: !!float 3e-4
    weight_decay: !!float 1e-4
    betas: [0.9, 0.999]
  pixel_opt:
    type: L1Loss
    loss_weight: 1
    reduction: mean

# Validation settings
val:
  val_freq: !!float 4e3
  save_img: false
  metrics:
    psnr:
      type: calculate_psnr
      crop_border: 0
      test_y_channel: false

# Logging
logger:
  print_freq: 1000
  save_checkpoint_freq: !!float 4e3
  use_tb_logger: true
```

--------------------------------

### BasicSR YAML Configuration for Testing

Source: https://context7.com/guan-ms/watermamba/llms.txt

Example YAML configuration for BasicSR, defining test datasets and model paths. This setup is used by `basicsr/test.py` to specify datasets for evaluation and the pre-trained model to load.

```yaml
# Add test dataset entries to WaterMamba.yml:
datasets:
  test_1:
    name: UIEB
    type: Dataset_PairedImage
    dataroot_gt: datasets/UIEB/GT
    dataroot_lq: datasets/UIEB/LQ
    io_backend:
      type: disk

  test_2:
    name: SQUID
    type: Dataset_PairedImage
    dataroot_gt: datasets/SQUID/GT
    dataroot_lq: datasets/SQUID/LQ
    io_backend:
      type: disk

path:
  pretrain_network_g: pretrained/WaterMamba.pth
```

--------------------------------

### Install Dependencies for WaterMamba

Source: https://github.com/guan-ms/watermamba/blob/main/README.md

Install the necessary libraries for selective scan with efficient hardware design. Ensure you have the correct versions of causal_conv1d and mamba_ssm.

```bash
pip install causal_conv1d==1.0.0
pip install mamba_ssm==1.1.1
```

--------------------------------

### Implement and Use Loss Functions in PyTorch

Source: https://context7.com/guan-ms/watermamba/llms.txt

Shows the implementation and usage of L1Loss, CharbonnierLoss, and PSNRLoss from basicsr.models.losses.losses. Includes an example of combining L1 and Charbonnier losses for mixed training.

```python
import torch
from basicsr.models.losses.losses import L1Loss, CharbonnierLoss, PSNRLoss, MSELoss

pred = torch.rand(2, 3, 256, 256, requires_grad=True).cuda()
gt   = torch.rand(2, 3, 256, 256).cuda()

# L1 loss (used by default in WaterMamba)
l1 = L1Loss(loss_weight=1.0, reduction='mean').cuda()
loss = l1(pred, gt)
print(f"L1 loss: {loss.item():.4f}")

# Charbonnier loss (smooth L1, robust to outliers)
charb = CharbonnierLoss(loss_weight=1.0, reduction='mean', eps=1e-3).cuda()
loss = charb(pred, gt)
print(f"Charbonnier loss: {loss.item():.4f}")

# PSNR loss (negative PSNR, suitable for high-quality fine-tuning)
psnr_loss = PSNRLoss(loss_weight=1.0, reduction='mean').cuda()
loss = psnr_loss(pred, gt)
print(f"PSNR loss: {loss.item():.4f}")   # Negative value (higher PSNR = lower loss)

# Combined loss (e.g., L1 + Charbonnier for mixed training)
total_loss = l1(pred, gt) + 0.05 * charb(pred, gt)
total_loss.backward()
```

--------------------------------

### Calculate PSNR and SSIM with NumPy and Tensors

Source: https://context7.com/guan-ms/watermamba/llms.txt

Demonstrates calculating Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index Measure (SSIM) using both NumPy arrays (uint8 range [0, 255]) and PyTorch Tensors (float range [0, 1]). Includes an example for Y-channel only calculation.

```python
pred_np = np.random.randint(0, 256, (256, 256, 3), dtype=np.uint8)
gt_np   = np.random.randint(0, 256, (256, 256, 3), dtype=np.uint8)

psnr = calculate_psnr(pred_np, gt_np, crop_border=0, test_y_channel=False)
ssim = calculate_ssim(pred_np, gt_np, crop_border=0, test_y_channel=False)
print(f"PSNR: {psnr:.2f} dB")    # e.g., PSNR: 8.23 dB (random → low)
print(f"SSIM: {ssim:.4f}")        # e.g., SSIM: 0.0012

# --- Tensor usage (values in [0, 1] float range) ---
pred_t = torch.rand(3, 256, 256)
gt_t   = torch.rand(3, 256, 256)

psnr_t = calculate_psnr(pred_t, gt_t, crop_border=0)
ssim_t = calculate_ssim(pred_t, gt_t, crop_border=0)
print(f"PSNR (tensor): {psnr_t:.2f} dB")
print(f"SSIM (tensor): {ssim_t:.4f}")

# --- Y-channel only (for legacy benchmarks) ---
psnr_y = calculate_psnr(pred_np, gt_np, crop_border=4, test_y_channel=True)
```

--------------------------------

### BasicSR Training Script Commands

Source: https://context7.com/guan-ms/watermamba/llms.txt

Demonstrates various command-line options for training models using the BasicSR `train.py` script, including single-GPU, multi-GPU distributed training, and resuming from checkpoints.

```bash
# Single-GPU training
python basicsr/train.py -opt Mamba/Options/WaterMamba.yml

# Multi-GPU distributed training (8 GPUs) via train.sh
bash train.sh Mamba/Options/WaterMamba.yml

# Resume from checkpoint (set resume_state path in YAML, or use --auto_resume)
python basicsr/train.py -opt Mamba/Options/WaterMamba.yml --auto_resume

# Launch manually with torch.distributed
python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --master_port=4321 \
    basicsr/train.py \
    -opt Mamba/Options/WaterMamba.yml \
    --launcher pytorch
```

--------------------------------

### Load Model Checkpoints with BasicSR

Source: https://context7.com/guan-ms/watermamba/llms.txt

Demonstrates loading model weights from a BasicSR-formatted checkpoint file, including handling EMA weights and optimizer states. This is crucial for inference with pre-trained models.

```python
# --- In training: checkpoints are saved automatically by the train loop ---
# Checkpoint structure saved at: experiments/<name>/models/net_g_<iter>.pth
# Training state saved at:       experiments/<name>/training_states/<iter>.state

# --- Manual checkpoint loading for inference ---
import torch
from basicsr.models.archs.WaterMamba_arch import WaterMamba

model = WaterMamba(
    inp_channels=3, out_channels=3, dim=24,
    num_blocks=[2,2,2,2], mam_blocks=[2,2,2,2],
    num_refinement_blocks=4, d_state=16, bias=False
).cuda()

# Load a BasicSR-format checkpoint
ckpt = torch.load('experiments/WaterMamba/models/net_g_300000.pth', map_location='cuda')
# BasicSR checkpoints store weights under 'params' or 'params_ema'
state_dict = ckpt.get('params_ema', ckpt.get('params', ckpt))
model.load_state_dict(state_dict, strict=True)
model.eval()
print("Checkpoint loaded successfully")

# The Google Drive pretrained model uses the same format:
ckpt = torch.load('WaterMamba_pretrained.pth', map_location='cuda')
model.load_state_dict(ckpt['params'])
```

--------------------------------

### Parse YAML Configuration Files with BasicSR

Source: https://context7.com/guan-ms/watermamba/llms.txt

Shows how to use the `parse` function from `basicsr.utils.options` to load and process YAML configuration files for training. This utility resolves dataset paths and sets up experiment directories.

```python
from basicsr.utils.options import parse, dict2str

# Parse a training config
opt = parse('Mamba/Options/WaterMamba.yml', is_train=True)
```

--------------------------------

### Load Paired Image Datasets with BasicSR

Source: https://context7.com/guan-ms/watermamba/llms.txt

Demonstrates initializing `Dataset_PairedImage` for training and validation. The training dataset includes augmentations like cropping and flipping, while the validation dataset uses full images without augmentation.

```python
from torch.utils.data import DataLoader
from basicsr.data.paired_image_dataset import Dataset_PairedImage

# Training dataset with augmentation
train_dataset = Dataset_PairedImage({
    'dataroot_gt': 'datasets/UIEB/train/GT',
    'dataroot_lq': 'datasets/UIEB/train/LQ',
    'gt_size': 256,             # Random crop size for training
    'use_hflip': True,
    'use_rot': True,
    'geometric_augs': True,
    'scale': 1,
    'phase': 'train',
    'io_backend': {'type': 'disk'},
})

# Inference / validation dataset (no augmentation, full image)
val_dataset = Dataset_PairedImage({
    'dataroot_gt': 'datasets/UIEB/val/GT',
    'dataroot_lq': 'datasets/UIEB/val/LQ',
    'scale': 1,
    'phase': 'val',
    'io_backend': {'type': 'disk'},
    # gt_size omitted → full resolution images returned
})

loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
batch = next(iter(loader))
print(batch['lq'].shape)   # torch.Size([4, 3, 256, 256]) — degraded input
print(batch['gt'].shape)   # torch.Size([4, 3, 256, 256]) — clean target
print(batch['lq_path'])    # List of source file paths
```

--------------------------------

### Configure and Run WaterMamba Inference

Source: https://context7.com/guan-ms/watermamba/llms.txt

Sets up directories, loads the WaterMamba model from a checkpoint, and performs inference on a directory of images. Optionally calculates PSNR and SSIM against ground truth.

```python
inp_dir   = 'demo/input'        # Directory with degraded underwater images
out_dir   = 'demo/results'
gt_dir    = 'demo/target'       # Optional: ground-truth for PSNR/SSIM evaluation
weights   = 'pretrained/WaterMamba.pth'
os.makedirs(out_dir, exist_ok=True)

# --- Load model ---
model = WaterMamba(
    inp_channels=3, out_channels=3, dim=24,
    num_blocks=[2,2,2,2], mam_blocks=[2,2,2,2],
    num_refinement_blocks=4, d_state=16, bias=False
)
checkpoint = torch.load(weights)
model.load_state_dict(checkpoint['params'] if 'params' in checkpoint else checkpoint)
model.eval().cuda()

# --- Inference loop ---
psnr_list, ssim_list = [], []
for fname in sorted(os.listdir(inp_dir)):
    img_path = os.path.join(inp_dir, fname)
    img = load_img(img_path)                        # Returns float32 numpy (H,W,3)
    inp = torch.from_numpy(img).permute(2,0,1).unsqueeze(0).cuda()  # (1,3,H,W)

    # Pad to multiple of 8
    h, w = inp.shape[2], inp.shape[3]
    H = ((h + 7) // 8) * 8
    W = ((w + 7) // 8) * 8
    inp_pad = F.pad(inp, [0, W-w, 0, H-h], mode='reflect')

    with torch.no_grad():
        out_pad = model(inp_pad)

    out = out_pad[:, :, :h, :w].clamp(0, 1)
    out_np = out.squeeze(0).permute(1,2,0).cpu().numpy()

    save_img(os.path.join(out_dir, fname), img_as_ubyte(out_np))

    # Optional PSNR/SSIM against ground truth
    gt_path = os.path.join(gt_dir, fname)
    if os.path.exists(gt_path):
        gt = load_img(gt_path)
        psnr_list.append(calculate_psnr(out_np, gt))
        ssim_list.append(calculate_ssim(out_np, gt))

if psnr_list:
    print(f"Mean PSNR: {np.mean(psnr_list):.2f} dB")
    print(f"Mean SSIM: {np.mean(ssim_list):.4f}")
# Expected output (UIEB dataset): Mean PSNR: 24.70 dB, Mean SSIM: 0.9300
```

--------------------------------

### Initialize and Use ChannelMamba Model

Source: https://context7.com/guan-ms/watermamba/llms.txt

Instantiates the ChannelMamba model for bidirectional channel-wise state space modeling. Requires CUDA-enabled device.

```python
import torch
from basicsr.models.archs.WaterMamba_arch import ChannelMamba

cm = ChannelMamba(
    dim=24,
    d_state=16,
    bias=False
).cuda()

# Input: (B, H, W, C)
x = torch.randn(2, 64, 64, 24).cuda()
out = cm(x)
print(out.shape)   # torch.Size([2, 64, 64, 24])
```

--------------------------------

### Accessing Experiment Directories

Source: https://context7.com/guan-ms/watermamba/llms.txt

Prints the paths for experiment-related directories. These are automatically set up by the project.

```python
print(opt['path']['experiments_root'])  # .../experiments/WaterMamba
print(opt['path']['models'])            # .../experiments/WaterMamba/models
print(opt['path']['log'])               # .../experiments/WaterMamba
print(opt['path']['visualization'])     # .../experiments/WaterMamba/visualization
```

--------------------------------

### Run BasicSR Evaluation with YAML Config

Source: https://context7.com/guan-ms/watermamba/llms.txt

Executes the BasicSR testing script using a YAML configuration file. This process loads a specified model and weights, runs inference on test datasets, saves enhanced images, and logs performance metrics.

```bash
# Run evaluation using YAML config (test mode: is_train=False)
python basicsr/test.py -opt Mamba/Options/WaterMamba.yml

# The script will:
# 1. Load the model specified in network_g
# 2. Load weights from path.pretrain_network_g
# 3. Run inference on all datasets.test_* entries
# 4. Save enhanced images to results/<name>/visualization/
# 5. Log PSNR/SSIM metrics to results/<name>/<name>.log
```

--------------------------------

### Accessing Training Dataset Configuration

Source: https://context7.com/guan-ms/watermamba/llms.txt

Retrieves and prints specific configuration details for the training dataset, including the ground truth data root, phase, and size.

```python
train_cfg = opt['datasets']['train']
print(train_cfg['dataroot_gt'])   # Expanded absolute path
print(train_cfg['phase'])         # 'train'
print(train_cfg['gt_size'])       # 256
```

--------------------------------

### Initialize and Use MSFFN Model

Source: https://context7.com/guan-ms/watermamba/llms.txt

Instantiates the Multi-Scale Feed-Forward Network (MSFFN) for capturing multi-scale local texture. Assumes channel-first input layout and requires CUDA.

```python
import torch
from basicsr.models.archs.WaterMamba_arch import MSFFN

ffn = MSFFN(
    dim=24,
    ffn_expansion_factor=2.66,
    bias=False
).cuda()

# Input: (B, C, H, W) — channel-first layout
x = torch.randn(2, 24, 64, 64).cuda()
out = ffn(x)
print(out.shape)   # torch.Size([2, 24, 64, 64])
```

--------------------------------

### Instantiate and Use WaterMamba Model

Source: https://context7.com/guan-ms/watermamba/llms.txt

Instantiate the WaterMamba U-Net model with specified dimensions and block configurations. The model expects input in the range [0,1] and outputs enhanced images.

```python
import torch
from basicsr.models.archs.WaterMamba_arch import WaterMamba

# Instantiate WaterMamba with the default paper configuration
model = WaterMamba(
    inp_channels=3,
    out_channels=3,
    dim=24,
    num_blocks=[2, 2, 2, 2],       # Transformer-style blocks per level
    mam_blocks=[2, 2, 2, 2],       # SCOSS (Mamba) blocks per level
    num_refinement_blocks=4,
    d_state=16,                     # SSM state dimension
    bias=False,
    dual_pixel_task=False
).cuda()

# Forward pass on a batch of degraded underwater images
degraded = torch.randn(1, 3, 256, 256).cuda()   # (B, C, H, W), values in [0,1]
with torch.no_grad():
    enhanced = model(degraded)                   # Output: (1, 3, 256, 256)

print(enhanced.shape)   # torch.Size([1, 3, 256, 256])
print(enhanced.min().item(), enhanced.max().item())
```

--------------------------------

### Pretty-Printing Options

Source: https://context7.com/guan-ms/watermamba/llms.txt

Prints all options in a human-readable dictionary format for debugging purposes.

```python
print(dict2str(opt))
```

--------------------------------

### Use SCOSSBlock in WaterMamba

Source: https://context7.com/guan-ms/watermamba/llms.txt

Demonstrates the standalone usage of the SCOSSBlock, which combines spatial and channel scanning. The block expects input in HWC layout.

```python
import torch
from basicsr.models.archs.WaterMamba_arch import SCOSSBlock

# Standalone SCOSSBlock usage
block = SCOSSBlock(
    dim=24,        # channel dimension
    d_state=16,    # SSM hidden state size
    bias=False
).cuda()

# Input: (Batch, Height, Width, Channels) — note HWC layout used internally
x = torch.randn(2, 64, 64, 24).cuda()
out = block(x)
print(out.shape)   # torch.Size([2, 64, 64, 24])
```

--------------------------------

### Parsing Test Configuration

Source: https://context7.com/guan-ms/watermamba/llms.txt

Parses a specific YAML configuration file for testing and prints the results root path and training status.

```python
opt_test = parse('Mamba/Options/WaterMamba.yml', is_train=False)
print(opt_test['path']['results_root'])  # .../results/WaterMamba
print(opt_test['is_train'])              # False
```

--------------------------------

### Dataset Structure for WaterMamba

Source: https://github.com/guan-ms/watermamba/blob/main/README.md

Organize your dataset in the specified directory structure for WaterMamba to process. This includes 'train' and 'test' sets, each with 'input' and 'target' subfolders.

```text
WaterMamba
├─ other files and folders
├─ dataset
│	├─ demo
│	│	├─ train
│	│	│	├─ input
│	│	│	│	├─ fig1.png
│	│	│	│	├─ ...
│	│	│	├─ target
│	│	│	│	├─ fig1.png
│	│	│	│	├─ ...
│	│	├─ test
│	│	│	│	├─ ...

```

--------------------------------

### Load, Save, and Evaluate Images with Mamba Utils

Source: https://context7.com/guan-ms/watermamba/llms.txt

Utilizes Mamba/utils.py for image operations. Loads images as float32 NumPy arrays in the [0, 1] range, saves them after conversion to uint8, and calculates PSNR/SSIM between enhanced and ground truth images.

```python
from Mamba.utils import load_img, save_img, calculate_psnr, calculate_ssim
import numpy as np

# Load image as float32 numpy array (H, W, 3), values in [0, 1]
img = load_img('demo/input/663.jpg')
print(img.shape, img.dtype)   # (H, W, 3) float32

# Save float32 array back to disk (converts to uint8 internally via img_as_ubyte)
from skimage import img_as_ubyte
save_img('demo/results/663_enhanced.jpg', img_as_ubyte(img))

# PSNR/SSIM on float32 arrays in [0, 1]
enhanced = load_img('demo/results/663_enhanced.jpg')
gt        = load_img('demo/target/663.jpg')

psnr = calculate_psnr(enhanced, gt)
ssim = calculate_ssim(enhanced, gt)
print(f"PSNR: {psnr:.2f} dB")
print(f"SSIM: {ssim:.4f}")
```

--------------------------------

### SS2D (2D Selective State Space Module)

Source: https://context7.com/guan-ms/watermamba/llms.txt

Implements the 4-directional spatial scanning core of WaterMamba, capturing long-range spatial dependencies with linear complexity.

```APIDOC
## SS2D (2D Selective State Space Module)

### Description
`SS2D` implements the 4-directional spatial scanning core of WaterMamba. Input features are unfolded into four directional sequences, passed through a shared Mamba-style SSM, then folded back and summed. This captures long-range spatial dependencies at O(N) complexity.

### Usage
```python
import torch
from basicsr.models.archs.WaterMamba_arch import SS2D

ss2d = SS2D(
    d_model=24,     # input/output channel dimension
    d_state=16,     # SSM state dimension
    d_conv=3,       # local convolution kernel size
    expand=2,       # inner expansion factor
    dropout=0.0,
    bias=False
).cuda()

# Input: (B, H, W, C) channel-last format
x = torch.randn(2, 64, 64, 24).cuda()
out = ss2d(x)
print(out.shape)   # torch.Size([2, 64, 64, 24])
```

### Parameters
- **d_model** (int) - Input/output channel dimension.
- **d_state** (int) - SSM state dimension.
- **d_conv** (int) - Local convolution kernel size.
- **expand** (int) - Inner expansion factor.
- **dropout** (float) - Dropout rate.
- **bias** (bool) - Whether to use bias.
```

--------------------------------

### Standalone Inference Script Imports

Source: https://context7.com/guan-ms/watermamba/llms.txt

Imports necessary libraries for the standalone inference script, including PyTorch, image processing utilities, and model components.

```python
# Mamba/test.py — run from repository root
# Usage: python Mamba/test.py
# Edit the configuration section at the top of the file before running:

import torch
import torch.nn.functional as F
from skimage import img_as_ubyte
from basicsr.models.archs.WaterMamba_arch import WaterMamba
from Mamba.utils import load_img, save_img, calculate_psnr, calculate_ssim
import os
import numpy as np
```

--------------------------------

### Calculate PSNR and SSIM Metrics

Source: https://context7.com/guan-ms/watermamba/llms.txt

Imports and utilizes `calculate_psnr` and `calculate_ssim` functions from `basicsr.metrics.psnr_ssim`. These functions can process both NumPy arrays and PyTorch tensors, supporting evaluation on RGB or Y channel.

```python
import numpy as np
import torch
from basicsr.metrics.psnr_ssim import calculate_psnr, calculate_ssim
```

--------------------------------

### WaterMamba (Top-Level Model)

Source: https://context7.com/guan-ms/watermamba/llms.txt

The main U-Net encoder-decoder model for underwater image enhancement. It stacks SCOSSBlock layers and can be configured via several arguments to control its architecture size.

```APIDOC
## WaterMamba (Top-Level Model)

### Description
The `WaterMamba` class is the main U-Net encoder-decoder model. It stacks `SCOSSBlock` layers at each of four encoder and decoder levels, connected by patch-embedding, downsampling, and upsampling stages. The `inp_channels`, `out_channels`, `dim`, `num_blocks`, `mam_blocks`, `num_refinement_blocks`, `d_state`, and `bias` arguments fully control the architecture size.

### Usage
```python
import torch
from basicsr.models.archs.WaterMamba_arch import WaterMamba

# Instantiate WaterMamba with the default paper configuration
model = WaterMamba(
    inp_channels=3,
    out_channels=3,
    dim=24,
    num_blocks=[2, 2, 2, 2],       # Transformer-style blocks per level
    mam_blocks=[2, 2, 2, 2],       # SCOSS (Mamba) blocks per level
    num_refinement_blocks=4,
    d_state=16,                     # SSM state dimension
    bias=False,
    dual_pixel_task=False
).cuda()

# Forward pass on a batch of degraded underwater images
de-graded = torch.randn(1, 3, 256, 256).cuda()   # (B, C, H, W), values in [0,1]
with torch.no_grad():
    enhanced = model(de-graded)                   # Output: (1, 3, 256, 256)

print(enhanced.shape)   # torch.Size([1, 3, 256, 256])
print(enhanced.min().item(), enhanced.max().item())
```

### Parameters
- **inp_channels** (int) - Input channels.
- **out_channels** (int) - Output channels.
- **dim** (int) - Dimension of the model.
- **num_blocks** (list[int]) - Number of Transformer-style blocks per level.
- **mam_blocks** (list[int]) - Number of SCOSS (Mamba) blocks per level.
- **num_refinement_blocks** (int) - Number of refinement blocks.
- **d_state** (int) - SSM state dimension.
- **bias** (bool) - Whether to use bias.
- **dual_pixel_task** (bool) - Whether to use dual pixel task.
```

--------------------------------

### Use SS2D Module for Spatial Scanning

Source: https://context7.com/guan-ms/watermamba/llms.txt

Utilizes the SS2D module, which implements the 4-directional spatial scanning for capturing long-range dependencies. Input should be in channel-last format (B, H, W, C).

```python
import torch
from basicsr.models.archs.WaterMamba_arch import SS2D

ss2d = SS2D(
    d_model=24,     # input/output channel dimension
    d_state=16,     # SSM state dimension
    d_conv=3,       # local convolution kernel size
    expand=2,       # inner expansion factor
    dropout=0.0,
    bias=False
).cuda()

# Input: (B, H, W, C) channel-last format
x = torch.randn(2, 64, 64, 24).cuda()
out = ss2d(x)
print(out.shape)   # torch.Size([2, 64, 64, 24])
```

--------------------------------

### SCOSSBlock (Spatial-Channel Omnidirectional Selective Scan Block)

Source: https://context7.com/guan-ms/watermamba/llms.txt

The fundamental building block of WaterMamba, combining spatial and channel scanning with a multi-scale feed-forward network.

```APIDOC
## SCOSSBlock (Spatial-Channel Omnidirectional Selective Scan Block)

### Description
`SCOSSBlock` is the fundamental building block of WaterMamba. It wraps a `SS2D` spatial scanner and a `ChannelMamba` channel scanner inside a residual block with layer normalization. The spatial `SS2D` module scans feature maps in four directions, while `ChannelMamba` performs bidirectional scanning across the channel dimension. Both branches are fused and passed through the `MSFFN` multi-scale feed-forward network.

### Usage
```python
import torch
from basicsr.models.archs.WaterMamba_arch import SCOSSBlock

# Standalone SCOSSBlock usage
block = SCOSSBlock(
    dim=24,        # channel dimension
    d_state=16,    # SSM hidden state size
    bias=False
).cuda()

# Input: (Batch, Height, Width, Channels) — note HWC layout used internally
x = torch.randn(2, 64, 64, 24).cuda()
out = block(x)
print(out.shape)   # torch.Size([2, 64, 64, 24])
```

### Parameters
- **dim** (int) - Channel dimension.
- **d_state** (int) - SSM hidden state size.
- **bias** (bool) - Whether to use bias.
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.