MUNI: Multimodal Unified Latent Diffusion for Coherent Any-to-Any Generation

Kyeongmin Yeo*, Yunhong Min*, Minhyuk Sung
KAIST
* Equal contribution

Paper | Project Page

This repository is the official PyTorch implementation of MUNI, an end-to-end multimodal latent diffusion framework for coherent any-to-any generation. MUNI jointly learns modality-specific encoders and decoders with a single shared flow-based prior, supporting both subset-conditioned generation and fully unconditional joint sampling.

“A group of people with bicycles in the grass.” ▶ Play generated audio	“The beach is empty and clear with palm trees.” ▶ Play generated audio
“A man with an acoustic guitar and shirt.” ▶ Play generated audio	“The room has blue curtains and red furniture.” ▶ Play generated audio

Unconditional image-text-audio co-generation: each triplet is sampled jointly from MUNI’s learned prior. See the project page for more results.

Installation

The reference environment uses Python 3.10 and PyTorch 2.5.1 with CUDA 12.4. The code relies on recent PyTorch 2.x APIs and is intended to work with other recent PyTorch/CUDA combinations as well. Install the build appropriate for your system first; see the official Previous PyTorch Versions page for alternatives.

conda create -n muni python=3.10
conda activate muni

# Example: PyTorch 2.5.1 with CUDA 12.4
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \
  --index-url https://download.pytorch.org/whl/cu124

pip install -r requirements.txt

Run the commands below from the repository root.

Set the external roots once per shell, or add them to the environment's activation script:

export DATA_ROOT=/path/to/data
export CKPT_ROOT=/path/to/checkpoints

DATA_ROOT contains datasets and their compact runtime caches. CKPT_ROOT contains frozen foundation models used by the image-text-audio experiment.

Foundation checkpoints

Place the FlowBind foundation assets as follows:

$CKPT_ROOT/flowbind/
├── gemmae_3.pt
├── clap.pt
├── embeddinggemma-300m/
├── gemma-3-1b-pt/
├── stable-diffusion-2-1-unclip-small/
├── audioldm-m-full/
└── lightweight_audio_ae.pth

The image-text-audio setup follows the official FlowBind repository. FlowBind-specific assets, including gemmae_3.pt, are available from its model release; the remaining directories contain the corresponding upstream foundation models.

FFHQ and PolyMNIST evaluation additionally use the frozen visual verifier checkpoints ckpts/ffhq64_verifier.pth and ckpts/polymnist64_verifier.pth, described below. FFHQ FID uses fid_stats/ffhq64_stats.npz by default.

Data preparation

Training and evaluation read compact tensor or Parquet caches rather than large collections of small files. The builders below create the exact runtime layouts expected by the configs. Their shell launchers select the canonical paths; additional options accepted by the underlying Python builder can be appended to the same command.

PolyMNIST

The builder downloads MNIST and stores all construction sources under $DATA_ROOT/poly_mnist_sources. It writes the aligned image cache and labels under $DATA_ROOT/poly_mnist without creating a PNG corpus.

bash scripts/build_polymnist_cache.sh

FFHQ

Place aligned 256px RGB images under $DATA_ROOT/ffhq_256. The canonical end-to-end builder then:

1. predicts 19-class segmentation maps with the FFHQ-Aging DeepLabV3 model;
2. predicts normal maps with the Marigold E2E-FT pipeline and the pinned GonzaloMG/marigold-e2e-ft-normals checkpoint; and
3. packs RGB, segmentation, normals, and FFHQ-Aging labels into the 64px training cache and 256px verifier cache under $DATA_ROOT/ffhq.

Point the two variables below to the corresponding source checkouts. The FFHQ-Aging checkout must also contain ffhq_aging_labels.csv.

export FFHQ_AGING_REPO=/path/to/FFHQ-Aging-Dataset
export DIFFUSION_E2E_ROOT=/path/to/diffusion-e2e-ft

GENERATE_DENSE=1 \
FFHQ_SEGMENTATION_GPUS=0 \
FFHQ_NORMAL_GPUS=0 \
  bash scripts/build_ffhq_cache.sh

The generated 256px maps are retained at $DATA_ROOT/ffhq_sources/ffhq_256_seg and $DATA_ROOT/ffhq_sources/ffhq_256_normal, so interrupted extraction resumes without recomputing completed files. Running the builder later without GENERATE_DENSE=1 only repacks these retained sources.

Image-text-audio

Arrange the sources under $DATA_ROOT/flowbind_raw with the following names. The AudioCaps media must come from the complete private raw bundle used by the experiment; public metadata-only or incomplete mirrors are not substitutes.

$DATA_ROOT/flowbind_raw/
├── laion_coco_aesthetic/
│   ├── metadata.parquet
│   └── images/<sha256>.jpg
├── flickr30k/
│   ├── annotations.csv
│   └── images/<flickr_id>.jpg
├── audiocaps/
│   ├── metadata/{train,val,test}.csv
│   └── audio/<youtube_id>_<start_time>.wav
├── vggsound/
│   ├── metadata.csv
│   └── videos/<youtube_id>_<start:06d>.mp4
└── coco_val2014/
    ├── metadata.jsonl
    ├── captions.json
    └── images/<sample_stem>.jpg

Then run the complete resumable construction pipeline:

STAGE=all bash scripts/build_flowbind_cache.sh

The pipeline filters invalid records, synthesizes missing paired modalities, extracts normalized 768-dimensional features, and writes the compact Parquet dataset under $DATA_ROOT/flowbind. The required raw directory layout, pinned public sources, private AudioCaps caveat, and per-stage commands are documented in scripts/data_prep/flowbind/README.md.

Evaluation verifiers

FFHQ and PolyMNIST use frozen image classifiers to measure whether generated images agree with their conditioning labels. They are evaluation components: they receive no gradients and do not affect the MUNI training objective. The verifiers use the same compact caches prepared above and can be trained with:

# FFHQ: age, gender, glasses, segmentation, and surface normals
CUDA_VISIBLE_DEVICES=0 python -m verifiers.train \
  --config verifiers/configs/ffhq.yaml

# PolyMNIST: digit and quadrant, shared across m0/m1/m2
CUDA_VISIBLE_DEVICES=0 python -m verifiers.train \
  --config verifiers/configs/polymnist.yaml

These final recipes train on all available examples; their small overlapping validation slices are smoke checks rather than held-out estimates. Use best.pth for downstream scoring and place it at the corresponding path under ckpts/ shown above. The paths are defined once by _verifier_ckpt in configs/ffhq/base.yaml and configs/poly_mnist/base.yaml, then inherited by the task configs for config-driven online evaluation. Pass the same checkpoint explicitly with --verifier-ckpt to the standalone benchmark commands.

Training

Use the common distributed launcher with a config and the desired number of GPUs. The launcher defaults to one process/GPU; pass --gpus N for distributed training. train.batch_size is the per-GPU batch size.

# Image + text + audio: main MUNI configuration
bash scripts/train.sh configs/flowbind/muni_poe.yaml --gpus 2 train.batch_size=512

# PolyMNIST-Quadrant-Labels
bash scripts/train.sh configs/poly_mnist/muni.yaml --gpus 1

# FFHQ-64
bash scripts/train.sh configs/ffhq/ours.yaml --gpus 1

# Controlled Gaussian mixture
python -m scripts.build_gmm_marginals --out-dir output/gmm/marginals
bash scripts/train.sh configs/gmm/muni.yaml --gpus 1

FlowBind provides PoE and HELVAE aggregation variants as configs/flowbind/muni_poe.yaml and configs/flowbind/muni_helvae.yaml. The PoE variant uses MUNI's routed objective and is not the MVAE baseline. Config values can be overridden from the command line, and training can be continued with --resume output/path.

Sampling defaults live with their concrete prior and decoder adapters. Optional evaluation-only overrides use inference.prior_kwargs and inference.decoder_kwargs.{default,<modality>} in the YAML. The shared framework only routes these mappings; option names such as sampler_steps, temperature, and cfg_scale are interpreted by the selected adapter.

Evaluation and inference

Run the config-driven evaluation suite for generation, reconstruction, and translation:

CUDA_VISIBLE_DEVICES=0 python -m scripts.evaluate \
  --config configs/flowbind/muni_poe.yaml \
  --checkpoint output/flowbind/muni_poe/checkpoint-200.pth

Use the unified benchmark entrypoint for the paper protocols:

# Image-text-audio fidelity, alignment, many-to-one, and unconditional metrics
python -m scripts.benchmark flowbind --gpu 0 \
  --config configs/flowbind/muni_poe.yaml \
  --checkpoint output/flowbind/muni_poe/checkpoint-200.pth \
  --eval-config configs/flowbind/eval_benchmark.yaml \
  --tasks all --method MUNI

# PolyMNIST
python -m scripts.benchmark polymnist --gpu 0 \
  --method muni \
  --config configs/poly_mnist/muni.yaml \
  --checkpoint output/poly_mnist/muni/checkpoint-300.pth \
  --verifier-ckpt ckpts/polymnist64_verifier.pth

# FFHQ-64
python -m scripts.benchmark ffhq --gpu 0 \
  --method muni \
  --config configs/ffhq/ours.yaml \
  --checkpoint output/ffhq/ours/checkpoint-2000.pth \
  --verifier-ckpt ckpts/ffhq64_verifier.pth

# Controlled Gaussian mixture
python -m scripts.benchmark gmm --gpu 0 \
  --method muni \
  --config configs/gmm/muni.yaml \
  --checkpoint output/gmm/muni/checkpoint-last.pth

For decoded qualitative image-text-audio samples:

CUDA_VISIBLE_DEVICES=0 python -m scripts.qualitative \
  --config configs/flowbind/muni_poe.yaml \
  --checkpoint output/flowbind/muni_poe/checkpoint-200.pth \
  --eval-config configs/flowbind/eval_benchmark.yaml \
  --output-dir output/flowbind/muni_poe/qualitative

Run python -m scripts.benchmark <flowbind|polymnist|ffhq|gmm> --help for the full protocol-specific options.

Acknowledgements

This codebase builds on the following excellent open-source projects:

• FlowBind
• score_sde_pytorch
• JiT (Just image Transformer)