MambaGuard: A CLIP-Mamba Approach for OOD Generated Image Detection

This repository contains the official implementation of the paper MambaGuard: A CLIP-Mamba Approach for OOD Generated Image Detection.

MCDF, short for MambaClip-DeFuseNet, is a generated image detector designed for out-of-distribution (OOD) generalization. The project combines CLIP-Mamba style visual backbones with spectral feature refinement and contrastive supervision for real/fake image classification.

Overview

The official paper implementation in this repository is centered on:

• main.py as the training and evaluation entry point
• configs/config_mcdf.py as the default experiment configuration
• models/MCDF.py and models/Dual_MCDF.py as the main detector variants

In the paper setting, MCDF uses:

• a CLIP-Mamba backbone for robust visual representation learning
• an Einstein FFT Disentangled Block for spectral feature refinement
• a hybrid spectro-contrastive objective for improved OOD generalization

Highlights

• Official implementation of MambaGuard: A CLIP-Mamba Approach for OOD Generated Image Detection
• Support for multiple backbones, including CLIP-VMamba, CLIP-Simba, and CLIP baselines
• MCDF and DualBranchMCDF detector variants
• Spectral feature modeling with EinFFT
• Supervised contrastive training support
• Generator-wise evaluation on OOD test sets

Repository Structure

MCDF/
|-- configs/
|   |-- config_mcdf.py
|   `-- config_setting.py
|-- data/
|-- datasets/
|   `-- dataset.py
|-- kernels/
|   `-- selective_scan/
|-- mamba/
|-- models/
|   |-- MCDF.py
|   |-- Dual_MCDF.py
|   |-- Dual_MCDF_2d.py
|   |-- PIMNet.py
|   |-- gilm_det.py
|   |-- glim_clip.py
|   |-- vmamba/
|   `-- clip/
|-- pre_trained_weights/
|-- results/
|-- main.py
|-- engine.py
|-- earlystop.py
`-- utils.py

Supported Models

Detector Variants

• mcdf
• dual_mcdf
• mamba_clip
• vmunet

Supported Visual Backbones

• CLIP_VMamba_S
• CLIP_VMamba_B
• CLIP_VMamba_T220
• CLIP_Simba_S
• CLIP_Simba_B
• CLIP_Simba_L
• RN50
• ViT-B/32
• ViT-B/16
• ViT-L/14

Installation

Basic Dependencies

Install the dependencies required by the current codebase:

pip install torch torchvision
pip install timm tensorboardX scikit-learn opencv-python pillow scipy tqdm einops transformers
pip install matplotlib thop SimpleITK medpy h5py

CUDA and Custom Kernels

The repository includes selective scan / Mamba-related CUDA components. Depending on your environment, you may also need:

• CUDA Toolkit
• ninja
• a compatible C/C++ build toolchain

Data Preparation

The default official configuration in configs/config_mcdf.py expects:

• CNNSpot-style training data
• FDMAS-style OOD test data

Default Paths

wang2020_data_path = "../../datasets/CNNSpot"
test_root = "../../datasets/fdmas"

Training Data Layout

The default wang2020 mode recursively searches for samples under:

../../datasets/CNNSpot/
|-- train/
|   `-- progan/
|       `-- .../
|           |-- 0_real/
|           `-- 1_fake/
`-- test/
    `-- progan/
        `-- .../
            |-- 0_real/
            `-- 1_fake/

OOD Test Data Layout

For custom generator-wise testing, the code expects one subdirectory per generator under test_root:

../../datasets/fdmas/
|-- generator_a/
|   |-- 0_real/
|   `-- 1_fake/
|-- generator_b/
|   |-- 0_real/
|   `-- 1_fake/
`-- ...

Nested generator layouts are also supported as long as the final folders are named 0_real and 1_fake.

Pretrained Weights

Place the required pretrained weights in the paths referenced by the config files. The default MCDF setup uses:

models/clip/VMamba_B_clip.pt
pre_trained_weights/vmamba_small_e238_ema.pth
pre_trained_weights/vmamba_tiny_e292.pth

You can switch to other backbones by editing configs/config_mcdf.py.

Quick Start

1. Configure the Experiment

Edit configs/config_mcdf.py and update the key options as needed:

network = "mcdf"  # or "dual_mcdf"
clip_backbone_type = "CLIP_VMamba_B"
clip_backbone_path = "models/clip/VMamba_B_clip.pt"

wang2020_data_path = "../../datasets/CNNSpot"
test_root = "../../datasets/fdmas"

work_dir = "results/MCDF_blur_sig1/"

2. Train

python main.py

The script will:

• build the selected backbone and detector
• train on the configured dataset
• validate during training
• save checkpoints and logs under results/
• run generator-wise custom testing after training finishes

3. Evaluate Only

To run evaluation only, set the following in configs/config_mcdf.py:

only_test_and_save_figs = True
best_ckpt_path = "path/to/best_ap.pth"
resume_txt_path = "results/your_run/tested_generators.txt"

Then run:

python main.py

Important Configuration Options

Model and Backbone

• network: detector type, such as mcdf or dual_mcdf
• clip_backbone_type: visual backbone name
• clip_backbone_path: checkpoint path for the selected backbone
• model_config: detector-specific options

Data

• wang2020_data_path: path to the CNNSpot-style dataset
• test_root: root directory for OOD generator-wise evaluation
• cropSize: training crop size
• loadSize: resize resolution

Augmentation

• blur_prob
• blur_sig
• jpg_prob
• jpg_method
• jpg_qual
• rz_interp

Optimization

• batch_size
• epochs
• opt
• lr
• weight_decay
• sch

Contrastive Training

• use_supcon
• supcon_temperature
• supcon_weight

Outputs

Training outputs are saved under the configured work_dir:

• checkpoints/: model checkpoints
• log/: training logs
• outputs/: saved outputs
• summary/: TensorBoard event files

TensorBoard

tensorboard --logdir=results

Notes

• Some backbones and kernels require a properly configured CUDA environment.
• The repository contains additional research modules from related experiments; the official paper path is main.py with configs/config_mcdf.py.
• Make sure all pretrained checkpoints are available before launching training.