Open GenMask

Unofficial, implementation-oriented reproduction of GenMask: Adapting DiT for Segmentation via Direct Mask Generation.

The project adapts the Wan 2.1 T2V 1.3B diffusion transformer to generate segmentation masks directly in Wan VAE latent space. Qwen2.5-VL-7B provides image-and-text conditioning through a learned 3584 -> 4096 adapter.

This repository is not the official GenMask implementation. Model weights, datasets, and trained checkpoints are not distributed here.

Implemented

• Stage 1 Qwen-to-Wan conditioning adapter training
• Stage 2 full DiT fine-tuning
• temporal image-latent injection with per-frame timesteps
• VAE-free and segmentation-only ablations
• binary-mask and foreground-cutout targets
• one-step and multi-step Euler inference
• multi-GPU RefCOCO, ADE20K, and PASCAL VOC evaluation
• all-expression RefCOCO evaluation and one-pass threshold search

Reproduction Status

The current segmentation-only reproduction was evaluated with all referring expressions and a fixed threshold of 0.15:

Dataset	Split	Reproduction mIoU	Paper full mIoU
RefCOCO	testA	81.30	83.70
RefCOCO	testB	78.10	80.70
RefCOCO+	testA	75.49	80.00
RefCOCO+	testB	67.45	73.10
RefCOCOg	val	74.88	77.20
RefCOCOg	test	76.29	78.20

This is a strict comparison: the reproduction row above excludes generation training data, while the paper's main result uses joint segmentation and generation training. On the validation splits used for the no-generation ablation, the reproduction reaches 78.90 / 71.49 / 74.95 mIoU on RefCOCO / RefCOCO+ / RefCOCOg, compared with 81.00 / 74.20 / 76.70 reported by the paper.

Repository Layout

configs/                 Training configurations
genmask/                 Model, data, loss, and inference implementation
scripts/train_adapter_stage1.py
scripts/train.py
scripts/inference.py
scripts/eval.py
tests/                   Unit and regression tests

Installation

Python 3.10 and CUDA-capable GPUs are expected. The reproduction was tested with PyTorch 2.5.1, CUDA 12.1, Diffusers 0.38.0, Transformers 5.8.0, and Accelerate 1.13.0.

Conda

conda env create -f environment.yml
conda activate genmask
pip install -e . --no-deps

Existing PyTorch environment

Install a CUDA-compatible PyTorch build first, then:

pip install -r requirements.txt
pip install -e . --no-deps

Model Preparation

Place or symlink the frozen base models as follows:

models/
├── Wan2.1-T2V-1.3B/
│   ├── diffusion_pytorch_model.safetensors
│   └── Wan2.1_VAE.pth
└── Qwen2.5-VL-7B-Instruct/

Paths can be changed in YAML or with --override.

The Wan and Qwen checkpoints are subject to their original licenses.

Dataset Layout

The default configs expect:

data/
├── coco/
│   ├── ref/
│   │   ├── refcoco/
│   │   ├── refcoco+/
│   │   └── refcocog/
│   ├── train2014/
│   ├── train2017/
│   ├── val2017/
│   └── annotations/
├── ADE20k/
│   ├── images/{training,validation}/
│   ├── annotations/{training,validation}/
│   └── objectInfo150.txt
├── coco-stuff/
│   ├── images/{train2017,val2017}/
│   └── annotations/{train2017,val2017}/
└── VOCdevkit/VOC2012/
    ├── JPEGImages/
    ├── SegmentationClass/
    ├── SegmentationClassAug/
    └── ImageSets/Segmentation/{trainaug.txt,val.txt}

RefCOCO annotations must use the common refs(...).p and instances.json format. VOC training uses the augmented trainaug split when available.

Stage 1: Train the Conditioning Adapter

Stage 1 freezes Wan DiT, Wan VAE, and Qwen2.5-VL, and trains only the Qwen-to-Wan adapter on COCO caption generation flow matching.

CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
  --config_file scripts/accelerate_config.yaml \
  --num_processes 4 \
  scripts/train_adapter_stage1.py \
  --config configs/stage1_adapter.yaml

Outputs are written to a timestamped directory under output/stage1_adapter_coco/.

Stage 2: Train GenMask

The Stage 2 configs intentionally leave model.adapter_path unset. Pass the checkpoint produced by Stage 1:

CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
  --config_file scripts/accelerate_config.yaml \
  --num_processes 4 \
  scripts/train.py \
  --config configs/train_stage2_v2_seg_only.yaml \
  --override model.adapter_path=output/stage1_adapter_coco/stage1_TIMESTAMP/adapter-1000.pt

Available configurations:

Config	Purpose
train_stage2_v2.yaml	Segmentation and COCO-caption generation, 1:1
train_stage2_v2_seg_only.yaml	Segmentation-only ablation
train_stage2_vae_free.yaml	No image VAE latent injection
train_stage2_v2_foreground_cutout.yaml	Experimental foreground RGB target

Resume training:

accelerate launch \
  --config_file scripts/accelerate_config.yaml \
  --num_processes 4 \
  scripts/train.py \
  --config configs/train_stage2_v2_seg_only.yaml \
  --override model.adapter_path=/path/to/adapter.pt \
  --override training.resume=/path/to/checkpoint.pt

Evaluation

Paper RefCOCO splits

CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
  --config_file scripts/accelerate_config.yaml \
  --num_processes 4 \
  scripts/eval.py \
  --checkpoint /path/to/checkpoint.pt \
  --adapter /path/to/adapter.pt \
  --dataset paper_refcoco \
  --expression_mode all \
  --threshold 0.15 \
  --batch_size 8 \
  --image_size 256 \
  --vae_injection temporal \
  --output_dir output/eval_paper

Validation threshold search

Multiple thresholds reuse one inference pass:

python scripts/eval.py \
  --checkpoint /path/to/checkpoint.pt \
  --adapter /path/to/adapter.pt \
  --dataset all_refcoco \
  --split val \
  --expression_mode all \
  --thresholds 0.02 0.05 0.08 0.10 0.15 0.20 \
  --image_size 256

Use a validation split to select the threshold; do not tune it on test splits.

ADE20K or VOC

python scripts/eval.py \
  --checkpoint /path/to/checkpoint.pt \
  --adapter /path/to/adapter.pt \
  --dataset ade20k \
  --split val \
  --ade_protocol semantic_argmax

semantic_argmax queries only classes present in each ground-truth image. It is therefore an oracle-candidate protocol, not the official unrestricted 150-class ADE20K protocol.

Single-Image Inference

python scripts/inference.py \
  --checkpoint /path/to/checkpoint.pt \
  --adapter /path/to/adapter.pt \
  --image examples/input.jpg \
  --target "the person in a red shirt" \
  --threshold 0.15 \
  --output output/mask.png \
  --overlay output/overlay.png

Important Numerical Conventions

Segmentation target latent:

z_s = (1 - s) * epsilon + s * z0
v_target = epsilon - z0
model_t = (1 - s) * 1000

For temporal injection, the clean image latent is frame 0 with timestep zero; the noisy mask latent is frame 1 with the sampled mask timestep. Only frame 1 is supervised and returned.

Differences from the Paper

• The public generation loader currently uses COCO Captions. The paper reports DiffusionDB, BLIP-3o series, and additional third-party generation data.
• No official GenMask checkpoint or training code was available when this reproduction was developed.
• The Qwen-to-Wan adapter is an implementation choice required to preserve the pretrained Wan 4096-dimensional conditioning interface.
• Foreground-cutout training is an experimental extension, not part of the original method.
• ADE20K/VOC semantic argmax evaluation uses oracle ground-truth candidate classes and must not be compared with unrestricted semantic segmentation.

Testing

python -m unittest discover -s tests -p 'test_*.py'
python -m compileall -q genmask scripts tests

The test suite covers dataset mapping, expression expansion, target construction, gradient accumulation behavior, training logs, portable paths, and trainer configuration validation.

License

Repository code is released under the MIT License. Base models and datasets retain their original licenses and terms of use.

Citation

Please cite the original paper:

@article{yang2026genmask,
  title={GenMask: Adapting DiT for Segmentation via Direct Mask Generation},
  author={Yang, Yuhuan and Zhuang, Xianwei and Cai, Yuxuan and Ma, Chaofan and Bai, Shuai and Yao, Jiangchao and Zhang, Ya and Lin, Junyang and Wang, Yanfeng},
  journal={arXiv preprint arXiv:2603.23906},
  year={2026}
}