Constraint-Aware Flow on CLEVR

This project is a compact research prototype for image generation with explicit scene-structure constraints.

It trains two small 64x64 CLEVR generators:

• Constraint-aware flow matching: learns a velocity field from random noise to an image and also predicts object positions.
• Diffusion baseline: learns DDPM-style noise prediction and uses the same auxiliary object-position/left-right constraint head for comparison.

The key idea is that generated images should not only look plausible, but also satisfy symbolic scene relations such as "object A is left of object B". CLEVR is a good first testbed because every image has ground-truth objects and spatial relationships.

Hardware Target

The default configs are sized for a laptop with:

• Intel i7-13620H
• 16 GB system RAM
• RTX 4060 Laptop GPU with 8 GB VRAM

Defaults use 64x64 images, base-channels: 48, mixed precision, channels-last tensors, TF32, and batch sizes that should fit 8 GB VRAM. num-workers is set to 2 with DataLoader prefetching because that was the fastest measured setting on i7/RTX 4060 system.

Setup

python -m pip install -r requirements.txt

The repo expects this data layout, which is already present here:

CLEVR_v1.0/
  images/train
  images/val
  scenes/CLEVR_train_scenes.json
  scenes/CLEVR_val_scenes.json

Quick Checks

$env:PYTHONDONTWRITEBYTECODE='1'
python test_dataset.py
python -m py_compile data/clevr_dataset.py models/flow/flow_unet.py models/diffusion/diffusion_unet.py

PYTHONDONTWRITEBYTECODE avoids Windows permission issues if stale __pycache__ files are locked.

Train

Train the constraint-aware flow model:

$env:PYTHONDONTWRITEBYTECODE='1'
python training/train_flow.py --config configs/flow.yaml

Maximum-speed measured flow run:

python training/train_flow.py --config configs/flow.yaml --num-workers 2 --prefetch-factor 4 --batch-size 16

Train the diffusion baseline:

$env:PYTHONDONTWRITEBYTECODE='1'
python training/train_diffusion.py --config configs/diffusion.yaml

For a short smoke run:

python training/train_flow.py --config configs/flow.yaml --max-samples 64 --epochs 1 --batch-size 4 --num-workers 0

Resume training:

python training/train_flow.py --config configs/flow.yaml --resume experiments/checkpoints/flow_latest.pt

Sample

python sampling/sample_flow.py --checkpoint experiments/checkpoints/flow_latest.pt
python sampling/sample_diffusion.py --checkpoint experiments/checkpoints/diffusion_latest.pt

Samples are saved under experiments/samples/.

Evaluate Spatial Relations

python evaluation/compare_models.py --max-samples 256

This reports average left/right relation satisfaction and violation magnitude on the validation split.

Results

After 20 epochs of training, evaluating spatial relation satisfaction across 256 validation samples yields the following:

Model	Relation Satisfaction	Average Violation Magnitude
Constraint-Aware Flow	4.17%	0.0514
Diffusion Baseline	2.63%	0.0503

Note: As a compact 64x64 prototype, absolute satisfaction rates are expectedly low, but the explicit constraint-aware Flow model successfully demonstrates a relative improvement (~58% increase) in structural constraint awareness over the standard diffusion baseline.

Sample Outputs

Constraint-Aware Flow Output:

Diffusion Baseline Output:

Notes

This is not a production-scale image generator. It is a small, runnable experiment for testing whether explicit relational losses can improve controllability and sampling efficiency. Scaling the idea would require stronger architectures, larger datasets, text conditioning, better object discovery, and modern samplers.