### VQ-VAE Training Configuration Example (YAML)

Source: https://github.com/apple/ml-4m/blob/main/README_TOKENIZATION.md

Example YAML configuration snippet demonstrating multi-resolution adaptation settings for VQ-VAE training. This allows tokenizers to work at different resolutions by sampling training resolutions within a specified range.

```yaml
input_size_min: 224
input_size_max: 512
resolution_step: 32
```

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

### Install 4M Framework Dependencies

Source: https://github.com/apple/ml-4m/blob/main/README.md

Instructions for setting up the 4M environment using Conda and pip. This process includes cloning the repository, creating a Python 3.9 environment, and installing the package in editable mode.

```bash
git clone https://github.com/apple/ml-4m
cd ml-4m
conda create -n fourm python=3.9 -y
conda activate fourm
pip install --upgrade pip
pip install -e .
```

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

### Train 4M Models with Distributed Training

Source: https://context7.com/apple/ml-4m/llms.txt

Provides command-line instructions for training 4M models using PyTorch torchrun with DDP or FSDP for single and multi-node setups.

```bash
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_4m.py \
  --config cfgs/default/4m/models/main/4m-b_mod7_500b.yaml

OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_4m_fsdp.py \
  --config cfgs/default/4m/models/main/4m-l_mod7_500b.yaml

torchrun --nnodes=8 --nproc_per_node=8 \
  --rdzv_id=4m_training --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  run_training_4m_fsdp.py \
  --config cfgs/default/4m/models/main/4m-xl_mod7_500b.yaml
```

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

### Multi-Guidance Generation Setup (Python)

Source: https://context7.com/apple/ml-4m/llms.txt

Sets up a generation sampler and defines multiple conditional dictionaries for multi-guidance generation. This allows for precise control over outputs by combining different modalities with weighted classifier-free guidance. It utilizes the GenerationSampler and build_chained_generation_schedules from the fourm library.

```python
from fourm.models.generate import GenerationSampler, build_chained_generation_schedules
import copy

# Setup sampler (assume fm model is loaded)
sampler = GenerationSampler(fm)

# Create unconditional base dict
uncond_dict = {}  # Empty conditioning

# Create multiple conditional dicts with different modalities
cond_dict_caption = copy.deepcopy(uncond_dict)
cond_dict_caption = custom_text(cond_dict_caption, 'A dog [S_1]', '[EOS]', 'caption', device, text_tokenizer)

cond_dict_depth = copy.deepcopy(uncond_dict)
cond_dict_depth['tok_depth@224'] = depth_tokens  # Pre-computed depth tokens

# Define weights for multi-guidance: l_uncond + w1*(l_caption - l_uncond) + w2*(l_depth - l_uncond)
cond_dicts = [cond_dict_caption, cond_dict_depth]

# Build schedule with per-condition weights
schedule = build_chained_generation_schedules(
    cond_domains=[],
    target_domains=['tok_rgb@224'],
    tokens_per_target=[196],
    autoregression_schemes=['roar'],
    decoding_steps=[50],
    token_decoding_schedules=['linear'],
    temps=[3.0],
    temp_schedules=['onex:0.5:0.5'],
    cfg_scales=[[2.0, 1.5]],  # List of weights per condition
    cfg_schedules=['constant'],
)

```

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

### Project Setup and Environment Configuration (Python)

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-21.ipynb

This snippet configures the Python environment for the project. It sets the CUDA visible devices, changes the current directory to the project root, and enables automatic reloading of modules.

```python
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
current_folder = globals()['_dh'][0]
os.chdir(os.path.dirname(os.path.abspath(current_folder)))

%load_ext autoreload
%autoreload 2
```

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

### Initialize Target Modalities in Sample Dictionary - Python

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-7.ipynb

Initializes the sample dictionary with an RGB input and prepares empty slots for target modalities. It defines the structure for each modality, including 'tensor', 'input_mask', 'target_mask', and 'decoder_attention_mask'. This setup is essential for the chained generation process, ensuring all necessary components are present before prediction.

```python
batched_sample = {
    'rgb@224': {
        'tensor': img, # Batched tensor
        'input_mask': torch.zeros(1, 196, dtype=torch.bool, device=device), # False = used as input, True = ignored
        'target_mask': torch.ones(1, 196, dtype=torch.bool, device=device), # False = predicted as target, True = ignored
    }
}

# Initialize target modalities
for target_mod, ntoks in zip(target_domains, tokens_per_target):
    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)
    
```

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

### Quick Start RGB-to-all and Text-to-all Generation with Demo4MSampler

Source: https://context7.com/apple/ml-4m/llms.txt

The Demo4MSampler class simplifies RGB-to-all and text-to-all generation. It automatically handles model and tokenizer loading for a streamlined generation pipeline. Input can be an image URL or a text caption, and outputs can be visualized or converted to PIL images.

```python
from fourm.demo_4M_sampler import Demo4MSampler, img_from_url
import torch

# Initialize sampler (loads 4M-21 XL model + super-resolution + all tokenizers)
sampler = Demo4MSampler(
    fm='EPFL-VILAB/4M-21_XL',              # Base 4M model
    fm_sr='EPFL-VILAB/4M-7-SR_L_CC12M',    # Super-resolution model (optional)
).cuda()

# RGB-to-all generation: Input an image, generate all modalities
img_url = 'https://storage.googleapis.com/four_m_site/images/demo_rgb.png'
img = img_from_url(img_url)  # Returns 1x3x224x224 ImageNet-normalized tensor

preds = sampler({'rgb@224': img.cuda()}, seed=42)
# Returns dict with keys like: 'tok_depth@224', 'tok_normal@224', 'tok_semseg@224',
#                              'caption', 'det', 'tok_clip@224', etc.

# Display results
sampler.plot_modalities(preds, save_path='./output.png')

# Text-to-all generation: Generate from caption
preds_from_text = sampler(
    {'caption': 'A lake house with a boat in front [S_1]'},
    seed=42
)
sampler.plot_modalities(preds_from_text, save_path='./text_to_all.png')

# Convert modalities to PIL images for further processing
pil_images = sampler.modalities_to_pil(preds, resize=512)
for img_pil, name in pil_images:
    img_pil.save(f'./outputs/{name}.png')

```

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

### Advanced Generation Control with GenerationSampler

Source: https://context7.com/apple/ml-4m/llms.txt

The GenerationSampler offers fine-grained control over generation, allowing custom schedules, schemes, and guidance. It requires manual loading of the base model and tokenizers. This example demonstrates setting up a chained generation schedule from caption to CLIP and RGB.

```python
from fourm.models.fm import FM
from fourm.models.generate import (
    GenerationSampler,
    build_chained_generation_schedules,
    init_empty_target_modality,
    init_full_input_modality,
    custom_text
)
from fourm.data.modality_info import MODALITY_INFO
from tokenizers import Tokenizer
import torch

# Load model and create sampler
fm = FM.from_pretrained('EPFL-VILAB/4M-21_XL').cuda().eval()
sampler = GenerationSampler(fm)

# Load text tokenizer
text_tokenizer = Tokenizer.from_file('./fourm/utils/tokenizer/trained/text_tokenizer_4m_wordpiece_30k.json')

# Define generation schedule for caption → CLIP → RGB chain
schedule = build_chained_generation_schedules(
    cond_domains=['caption'],                         # Input modalities
    target_domains=['tok_clip@224', 'tok_rgb@224'],   # Output modalities (in order)
    tokens_per_target=[196, 196],                      # Tokens to generate per modality
    autoregression_schemes=['roar', 'roar'],           # Generation scheme (maskgit, roar, autoregressive)
    decoding_steps=[50, 25],                           # Steps for parallel decoding
    token_decoding_schedules=['linear', 'linear'],     # Token schedule (cosine, linear)
    temps=[5.0, 3.0],                                  # Sampling temperatures
    temp_schedules=['onex:0.5:0.5', 'onex:0.5:0.5'],   # Temp decay schedule
    cfg_scales=[3.0, 2.0],                             # Classifier-free guidance scale
    cfg_schedules=['constant', 'constant'],            # CFG schedule
    cfg_grow_conditioning=True,                        # Add completed modalities to CFG
)

# Prepare input with custom caption
device = 'cuda'
sample = {}
sample = custom_text(
    sample,
    input_text='A serene mountain landscape with snow-capped peaks [S_1]',
    eos_token='[EOS]',
    key='caption',
    device=device,
    text_tokenizer=text_tokenizer
)

# Initialize target modalities
for target_mod, ntoks in [('tok_clip@224', 196), ('tok_rgb@224', 196)]:
    sample = init_empty_target_modality(sample, MODALITY_INFO, target_mod, 1, ntoks, device)

# Generate with custom parameters
with torch.no_grad():
    output = sampler.generate(
        sample,
        schedule,
        text_tokenizer=text_tokenizer,
        top_p=0.8,      # Nucleus sampling threshold
        top_k=0.0,      # Top-k sampling (0 = disabled)
        verbose=True,   # Show progress bar
        seed=42
    )

```

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

### Load 4M Tokenizers from Hugging Face Hub (Python)

Source: https://github.com/apple/ml-4m/blob/main/README.md

This code snippet demonstrates how to load different 4M tokenizers from the Hugging Face Hub using the `fourm` library. It shows how to instantiate `DiVAE` for modalities like RGB, depth, and normals, and `VQVAE` for semantic segmentation and CLIP embeddings. Ensure the `fourm` library is installed.

```python
from fourm.vq.vqvae import VQVAE, DiVAE

# 4M-7 modalities
tok_rgb = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_rgb_16k_224-448')
tok_depth = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_depth_8k_224-448')
tok_normal = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_normal_8k_224-448')
tok_semseg = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_semseg_4k_224-448')
tok_clip = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448')
```

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

### Verify CUDA Availability

Source: https://github.com/apple/ml-4m/blob/main/README.md

A simple Python check to ensure that the PyTorch installation has access to CUDA-enabled GPU acceleration, which is required for efficient model inference.

```python
import torch
print(torch.cuda.is_available())
```

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

### Load Modality-Specific Tokenizers

Source: https://context7.com/apple/ml-4m/llms.txt

Shows how to initialize VQVAE and DiVAE tokenizers for various modalities, including image-like data and feature embeddings, and prepare them for GPU inference.

```python
from fourm.vq.vqvae import VQVAE, DiVAE

# Load DiVAE tokenizers (with diffusion decoders for RGB, depth, normals, edges)
tok_rgb = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_rgb_16k_224-448')
tok_depth = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_depth_8k_224-448')
tok_normal = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_normal_8k_224-448')
tok_edge = DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_edge_8k_224-512')

# Load VQVAE tokenizers (for semantic/feature modalities)
tok_semseg = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_semseg_4k_224-448')
tok_clip = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448')
tok_dinov2 = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448')
tok_imagebind = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448')

# Load global feature tokenizers (single embedding vectors)
tok_dinov2_global = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_DINOv2-B14-global_8k_16_224')
tok_imagebind_global = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_ImageBind-H14-global_8k_16_224')

# Load structured modality tokenizers
tok_sam_instance = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_sam-instance_1k_64')
tok_human_poses = VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_human-poses_1k_8')

# Move all tokenizers to GPU
for tok in [tok_rgb, tok_depth, tok_normal, tok_semseg, tok_clip]:
    tok.cuda().eval()
```

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

### VQ-VAE Decoder Training with Frozen Encoder (YAML)

Source: https://github.com/apple/ml-4m/blob/main/README_TOKENIZATION.md

YAML configuration snippet for training a diffusion decoder with a frozen VQ-VAE encoder. This setup is useful for training decoders independently or fine-tuning them without altering the pre-trained encoder.

```yaml
full_ckpt: /path/to/checkpoint.pth
freeze_enc: True # Decoder can be trained from scratch or fine-tuned without the encoder
input_size_enc: 256 # Size of the encoder positional embeddings
```

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

### Configure Chained Generation Schedules

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-7.ipynb

Defines the generation sequence for multiple modalities by setting autoregression schemes, decoding steps, and configuration scales. This setup allows the model to generate intermediate representations like CLIP features before final RGB output.

```python
cond_domains = ['caption', 'det']
target_domains = ['tok_clip@224', 'tok_depth@224', 'tok_normal@224', 'tok_semseg@224', 'tok_rgb@224']
tokens_per_target = [196, 196, 196, 196, 196]
autoregression_schemes = ['roar', 'roar', 'roar', 'roar', 'roar']
decoding_steps = [50, 8, 8, 8, 25]
token_decoding_schedules = ['linear', 'linear', 'linear', 'linear', 'linear']
temps = [5.0, 3.0, 3.0, 3.0, 3.0]
temp_schedules = ['onex:0.5:0.5', 'onex:0.5:0.5', 'onex:0.5:0.5', 'onex:0.5:0.5', 'onex:0.5:0.5']
cfg_scales = [3.0, 2.0, 2.0, 2.0, 2.0]
cfg_schedules = ['constant', 'constant', 'constant', 'constant', 'constant']
cfg_grow_conditioning = True
top_p, top_k = 0.8, 0.0

schedule = build_chained_generation_schedules(
    cond_domains=cond_domains, target_domains=target_domains, tokens_per_target=tokens_per_target, autoregression_schemes=autoregression_schemes, 
    decoding_steps=decoding_steps, token_decoding_schedules=token_decoding_schedules, temps=temps, temp_schedules=temp_schedules,
    cfg_scales=cfg_scales, cfg_schedules=cfg_schedules, cfg_grow_conditioning=cfg_grow_conditioning, 
)
```

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

### Train 4M Model using PyTorch DDP

Source: https://github.com/apple/ml-4m/blob/main/README_TRAINING.md

This bash script demonstrates how to initiate the training of a 4M model using PyTorch Distributed Data Parallel (DDP). It is recommended for B-sized models and requires specifying the number of processes per node and the configuration file.

```bash
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_4m.py \
--config cfgs/default/4m/models/<config>.yaml
```

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

### Train VQVAE Tokenizers for RGB, CLIP, and Semantic Segmentation (Bash)

Source: https://context7.com/apple/ml-4m/llms.txt

Trains custom VQ-VAE tokenizers for different modalities like RGB images, CLIP features, and semantic segmentation masks. It utilizes the run_training_vqvae.py script with various configuration files and data paths. Multi-resolution adaptation is also supported after base training.

```bash
# Train VQVAE on RGB images with CLIP perceptual loss
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_vqvae.py \
  --config cfgs/default/tokenization/vqvae/rgb/ViTB-ViTB_1k_224_CLIPB16-5.0.yaml \
  --data_path /path/to/imagenet/train/ \
  --eval_data_path /path/to/imagenet/val/

# Train VQVAE for CLIP feature tokenization
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_vqvae.py \
  --config cfgs/default/tokenization/vqvae/CLIP-B16/ViTB-ViTB_8k_224.yaml \
  --data_path /path/to/dataset/train/ \
  --eval_data_path /path/to/dataset/val/

# Train VQVAE for semantic segmentation
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_vqvae.py \
  --config cfgs/default/tokenization/vqvae/semseg_coco/ViTB-ViTB_4k_224.yaml \
  --data_path /path/to/coco_semseg/train/ \
  --eval_data_path /path/to/coco_semseg/val/

# Multi-resolution adaptation (after base training)
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_vqvae.py \
  --config cfgs/default/tokenization/vqvae/rgb/ViTB-ViTB_1k_224-448_CLIPB16-5.0.yaml \
  --data_path /path/to/dataset/train/ \
  --eval_data_path /path/to/dataset/val/
```

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

### Initialize input modalities for ML-4M

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-21.ipynb

Prepares input data by initializing modalities based on the provided domain configuration. This step ensures the batched samples are correctly formatted for the generation process.

```python
for cond_mod in cond_domains:
    batched_sample = init_full_input_modality(batched_sample, MODALITY_INFO, cond_mod, device, eos_id=text_tok.token_to_id("[EOS]"))
```

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

### Initialize Modalities and Generate Dense SAM Instances

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-21.ipynb

Prepares the input data structures by initializing target and input modalities, then executes the dense SAM prediction using the sampler. The output is subsequently decoded into a usable format.

```python
batched_sample = {
    'rgb@224': {
        'tensor': img, # Batched tensor
        'input_mask': torch.zeros(1, 196, dtype=torch.bool, device=device), # False = used as input, True = ignored
        'target_mask': torch.ones(1, 196, dtype=torch.bool, device=device), # False = predicted as target, True = ignored
    }
}

# Initialize target modalities
for target_mod, ntoks in zip(target_domains, tokens_per_target):
    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)
    
# Initialize input modalities
for cond_mod in cond_domains:
    batched_sample = init_full_input_modality(batched_sample, MODALITY_INFO, cond_mod, device, eos_id=text_tok.token_to_id("[EOS]"))

# In the method the only target_domain considered in the schedule is sam_instance
out_dict = sampler.generate_sam_dense(
    batched_sample, schedule, text_tokenizer=text_tok, 
    verbose=True, seed=0,
    top_p=top_p, top_k=top_k, batch_size=32,
)
dec_dict = decode_dict(
    out_dict, toks, text_tok, 
    image_size=224, patch_size=16,
    decoding_steps=50
)
```

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

### Load 4M Models from Hugging Face Hub

Source: https://context7.com/apple/ml-4m/llms.txt

Demonstrates how to initialize various pre-trained 4M models, including base, large, XL, text-to-image, and super-resolution variants using the FM class.

```python
from fourm.models.fm import FM

# Load 4M-7 models (7 modalities: RGB, depth, normal, semseg, CLIP, caption, detection)
fm7b_cc12m = FM.from_pretrained('EPFL-VILAB/4M-7_B_CC12M')     # Base model, 198M params
fm7l_cc12m = FM.from_pretrained('EPFL-VILAB/4M-7_L_CC12M')     # Large model, 705M params
fm7xl_cc12m = FM.from_pretrained('EPFL-VILAB/4M-7_XL_CC12M')   # XL model, 2.8B params

# Load 4M-21 models (21 modalities including DINOv2, ImageBind, edges, human poses, etc.)
fm21b = FM.from_pretrained('EPFL-VILAB/4M-21_B')   # Base model
fm21l = FM.from_pretrained('EPFL-VILAB/4M-21_L')   # Large model
fm21xl = FM.from_pretrained('EPFL-VILAB/4M-21_XL') # XL model

# Load specialized text-to-image models (better at text→image generation)
fm7b_t2i = FM.from_pretrained('EPFL-VILAB/4M-7-T2I_B_CC12M')
fm7l_t2i = FM.from_pretrained('EPFL-VILAB/4M-7-T2I_L_CC12M')

# Load super-resolution model (upscales 224→448)
fm7l_sr = FM.from_pretrained('EPFL-VILAB/4M-7-SR_L_CC12M')
```

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

### Run Batch Generation with run_generation.py (Bash)

Source: https://context7.com/apple/ml-4m/llms.txt

Executes batch generation on datasets using the run_generation.py script. This script supports chained generation with customizable schedules and various configurations for models, data, and generation settings. Examples include Text-to-CLIP-to-RGB generation and generation on COCO and CC12M datasets.

```bash
# Text → CLIP → RGB generation on Parti prompts with super-resolution
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_generation.py \
  -c cfgs/default/generation/models/4m-xl_mod7+sr_4m-l_mod7.yaml \
  -dc cfgs/default/generation/data/parti_3x.yaml \
  -gc cfgs/default/generation/settings_base/T2CR_roar49-25_cfg3_t6-0.5.yaml \
  -src cfgs/default/generation/settings_sr/x2CR_mg8_cfg3_t1const.yaml

# Generate on COCO validation set
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_generation.py \
  -c cfgs/default/generation/models/4m-l_mod21+sr_4m-l_mod7.yaml \
  -dc cfgs/default/generation/data/coco_30k.yaml \
  -gc cfgs/default/generation/settings_base/T2CR_roar49-25_cfg3_t6-0.5.yaml

# Generate on CC12M subset
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_generation.py \
  -c cfgs/default/generation/models/4m-b_mod7+sr_4m-l_mod7.yaml \
  -dc cfgs/default/generation/data/cc12m_30k.yaml \
  -gc cfgs/default/generation/settings_base/T2CR_roar49-25_cfg3_t6-0.5.yaml
```

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

### Train DiVAE Tokenizers for RGB, Depth, and Surface Normals (Bash)

Source: https://context7.com/apple/ml-4m/llms.txt

Trains diffusion-decoder (DiVAE) tokenizers for high-fidelity reconstruction of image-like modalities such as RGB, depth, and surface normals. The script run_training_divae.py is used with specific configurations. It supports training from scratch or on top of a frozen VQVAE encoder.

```bash
# Train DiVAE for RGB images from scratch
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_divae.py \
  --config cfgs/default/tokenization/divae/rgb/ViTB-UNetP4_16k_224_predx0.yaml \
  --data_path /path/to/imagenet/train/ \
  --eval_data_path /path/to/imagenet/val/

# Train diffusion decoder on top of frozen VQVAE encoder
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_divae.py \
  --config cfgs/default/tokenization/divae/depth/ViTB-UNetP4_8k_224_predx0.yaml \
  --full_ckpt /path/to/pretrained_vqvae.pth \
  --freeze_enc True \
  --data_path /path/to/depth/train/ \
  --eval_data_path /path/to/depth/val/

# Train DiVAE for surface normals with multi-resolution
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_divae.py \
  --config cfgs/default/tokenization/divae/normal/ViTB-UNetP4_8k_224-448_predx0.yaml \
  --data_path /path/to/normals/train/ \
  --eval_data_path /path/to/normals/val/
```

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

### Generate with Multiple Guidance Conditions

Source: https://context7.com/apple/ml-4m/llms.txt

Performs multi-guided generation using the sampler's generate_multi_guided method. It accepts unconditional and conditional dictionaries along with a generation schedule and sampling parameters like top_p and seed.

```python
output = sampler.generate_multi_guided(
    uncond_dict,
    cond_dicts,
    schedule,
    text_tokenizer=text_tokenizer,
    top_p=0.8,
    verbose=True,
    seed=42
)
```

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

### Train 4M Model using PyTorch FSDP

Source: https://github.com/apple/ml-4m/blob/main/README_TRAINING.md

This bash script shows how to train a 4M model using PyTorch Fully Sharded Data Parallel (FSDP). FSDP is more memory-efficient and is suitable for L and XL models. It also requires specifying the number of processes per node and the configuration file.

```bash
OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_training_4m_fsdp.py \
--config cfgs/default/4m/models/<config>.yaml
```

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

### Run 4M Multimodal Generation

Source: https://github.com/apple/ml-4m/blob/main/README.md

Demonstrates how to use the Demo4MSampler to generate various modalities from an RGB image input or text caption. The sampler requires a pre-trained model identifier and GPU acceleration.

```python
from fourm.demo_4M_sampler import Demo4MSampler, img_from_url
sampler = Demo4MSampler(fm='EPFL-VILAB/4M-21_XL').cuda()
img = img_from_url('https://storage.googleapis.com/four_m_site/images/demo_rgb.png')
preds = sampler({'rgb@224': img.cuda()}, seed=None)
sampler.plot_modalities(preds, save_path=None)
```

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

### Prepare Masked Inputs and Run Inference

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-7.ipynb

Initializes target modalities and processes caption and bounding box inputs with span masking. It then executes the sampler to generate the output dictionary and decodes the tokens into usable formats.

```python
caption = 'a nice restaurant in santorini at sunset [S_1]'
bboxes = '[S_1] v0=0 v1=250 v2=420 v3=999 potted plant ' \
         'v0=400 v1=750 v2=950 v3=999 dining table ' \
         'v0=350 v1=250 v2=700 v3=350 boat ' \
         'v0=700 v1=720 v2=740 v3=850 bottle [S_2]'

batched_sample = {}

for target_mod, ntoks in zip(target_domains, tokens_per_target):
    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)
    
batched_sample = custom_text(batched_sample, input_text=caption, eos_token='[EOS]', key='caption', device=device, text_tokenizer=text_tok)
batched_sample = custom_text(batched_sample, input_text=bboxes, eos_token='[EOS]', key='det', device=device, text_tokenizer=text_tok)

out_dict = sampler.generate(batched_sample, schedule, text_tokenizer=text_tok, verbose=True, seed=0, top_p=top_p, top_k=top_k)
dec_dict = decode_dict(out_dict, toks, text_tok, image_size=224, patch_size=16, decoding_steps=50)
```

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

### Load Models from Safetensors Checkpoints

Source: https://context7.com/apple/ml-4m/llms.txt

Explains how to manually load a 4M model from a local safetensors file, useful for offline environments or custom model paths.

```python
from fourm.utils import load_safetensors
from fourm.models.fm import FM

# Load checkpoint and config from safetensors file
ckpt, config = load_safetensors('/path/to/4m_model.safetensors')

# Initialize model with loaded config and weights
fm = FM(config=config)
fm.load_state_dict(ckpt)
fm.cuda().eval()

# Model is now ready for generation
print(f"Loaded model with {sum(p.numel() for p in fm.parameters())/1e6:.1f}M parameters")
print(f"Encoder modalities: {fm.encoder_modalities}")
print(f"Decoder modalities: {fm.decoder_modalities}")
```

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

### Importing Necessary Libraries (Python)

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-21.ipynb

This code block imports essential libraries for image processing, numerical operations, deep learning, tokenization, and plotting. It also initializes the device (CUDA or CPU) and disables gradient calculations for inference.

```python
from PIL import Image
import numpy as np
import torch
from torchvision.transforms.functional import center_crop
from tokenizers import Tokenizer
import matplotlib.pyplot as plt

from fourm.models.fm import FM
from fourm.vq.vqvae import VQVAE, DiVAE
from fourm.models.generate import GenerationSampler, build_chained_generation_schedules, init_empty_target_modality, init_full_input_modality, custom_text
from fourm.data.modality_transforms import RGBTransform
from fourm.data.modality_info import MODALITY_INFO
from fourm.utils.plotting_utils import decode_dict, visualize_bboxes, plot_text_in_square

# The flag below controls whether to allow TF32 on matmul. This flag defaults to False in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True
# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True

device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_grad_enabled(False)
```

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

### Configure Generation Schedule for SAM

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-21.ipynb

Defines the generation parameters including domains, autoregression schemes, and temperature schedules required for SAM instance prediction. This configuration is passed to the sampler to control the generation process.

```python
cond_domains = ['rgb@224']
target_domains = ['sam_instance']
tokens_per_target = [256]
autoregression_schemes = ['autoregressive']
decoding_steps = [None]
token_decoding_schedules = [ None]
temps = [0.01]
temp_schedules = ['constant'] * len(target_domains)
cfg_scales = [1.0]
cfg_schedules = ['constant'] * len(target_domains)
cfg_grow_conditioning = True
top_p, top_k = 0.8, 0.0

schedule = build_chained_generation_schedules(
    cond_domains=cond_domains, target_domains=target_domains, tokens_per_target=tokens_per_target, autoregression_schemes=autoregression_schemes, 
    decoding_steps=decoding_steps, token_decoding_schedules=token_decoding_schedules, temps=temps, temp_schedules=temp_schedules,
    cfg_scales=cfg_scales, cfg_schedules=cfg_schedules, cfg_grow_conditioning=cfg_grow_conditioning, 
)
```

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

### Build Chained Generation Schedules for Multi-Modal Prediction (Python)

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-21.ipynb

This Python code demonstrates how to construct a chained generation schedule for predicting various modalities from an initial RGB input. It defines conditioning domains, target domains, and specific parameters for each target, such as token counts, autoregression schemes, decoding steps, temperature, and CFG scales. The `build_chained_generation_schedules` function orchestrates this process, enabling consistent multi-modal output generation.

```python
cond_domains = ['rgb@224']
target_domains = [
    'tok_clip@224', 'tok_dinov2@224', 'tok_imagebind@224', 'tok_depth@224', 
    'tok_normal@224', 'tok_semseg@224', 'tok_canny_edge@224', 'tok_sam_edge@224', 
    'caption', 'det', 'human_poses', 'sam_instance', 'color_palette', 'metadata'
]
tokens_per_target = [196, 256, 256, 196, 196, 196, 196, 196, 256, 256, 275, 256, 23, 40]
autoregression_schemes = [
    'roar', 'roar', 'roar', 'roar', 'roar', 'roar', 'roar', 'roar', 
    'autoregressive', 'autoregressive', 'autoregressive', 'autoregressive', 
    'autoregressive', 'autoregressive'
]
decoding_steps = [1, 1, 1, 1, 1, 1, 1, 1, None, None, None, None, None, None]
token_decoding_schedules = [
    'linear', 'linear', 'linear', 'linear', 'linear', 'linear', 'linear', 'linear', 
    None, None, None, None, None, None
]
temps = [0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.3, 0.7, 0.1, 0.01, 0.1, 0.1]
temp_schedules = ['constant'] * len(target_domains)
cfg_scales = [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
cfg_schedules = ['constant'] * len(target_domains)
cfg_grow_conditioning = True
top_p, top_k = 0.8, 0.0

schedule = build_chained_generation_schedules(
    cond_domains=cond_domains, target_domains=target_domains, tokens_per_target=tokens_per_target, autoregression_schemes=autoregression_schemes, 
    decoding_steps=decoding_steps, token_decoding_schedules=token_decoding_schedules, temps=temps, temp_schedules=temp_schedules,
    cfg_scales=cfg_scales, cfg_schedules=cfg_schedules, cfg_grow_conditioning=cfg_grow_conditioning, 
)

```

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

### Set Project Path and Load Extensions (Python)

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-7.ipynb

This snippet configures the Python environment by setting the current working directory to the project root and enabling automatic reloading of modules. It's essential for running scripts within the project structure.

```python
# Switch path to root of project
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
current_folder = globals()['_dh'][0]
os.chdir(os.path.dirname(os.path.abspath(current_folder)))

%load_ext autoreload
%autoreload 2
```

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

### Load 4M Models from Hugging Face Hub (Python)

Source: https://github.com/apple/ml-4m/blob/main/README.md

Demonstrates how to load pre-trained 4M models directly from the Hugging Face Hub using the FM.from_pretrained method. This is the recommended approach for easy access to various model configurations.

```python
from fourm.models.fm import FM

fm7b_cc12m  = FM.from_pretrained('EPFL-VILAB/4M-7_B_CC12M')
fm7b_coyo   = FM.from_pretrained('EPFL-VILAB/4M-7_B_COYO700M')
fm21b       = FM.from_pretrained('EPFL-VILAB/4M-21_B')

fm7l_cc12m  = FM.from_pretrained('EPFL-VILAB/4M-7_L_CC12M')
fm7l_coyo   = FM.from_pretrained('EPFL-VILAB/4M-7_L_COYO700M')
fm21l       = FM.from_pretrained('EPFL-VILAB/4M-21_L')

fm7xl_cc12m = FM.from_pretrained('EPFL-VILAB/4M-7_XL_CC12M')
fm7xl_coyo  = FM.from_pretrained('EPFL-VILAB/4M-7_XL_COYO700M')
fm21xl      = FM.from_pretrained('EPFL-VILAB/4M-21_XL')
```

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

### Initialize Batched Sample and Target Modalities

Source: https://github.com/apple/ml-4m/blob/main/notebooks/retrieval_4M-21.ipynb

Prepares the input sample dictionary with masks and initializes the target modalities to be used by the generation sampler.

```python
batched_sample = {
    'rgb@224': {
        'tensor': img, # Batched tensor
        'input_mask': torch.zeros(1, 196, dtype=torch.bool, device=device), # False = used as input, True = ignored
        'target_mask': torch.ones(1, 196, dtype=torch.bool, device=device), # False = predicted as target, True = ignored
    }
}

# Initialize target modalities
for target_mod, ntoks in zip(target_domains, tokens_per_target):
    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)
```

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

### Initialize Environment and Dependencies

Source: https://github.com/apple/ml-4m/blob/main/notebooks/retrieval_4M-21.ipynb

Configures the runtime environment by setting CUDA device visibility, adjusting working directories, and importing necessary PyTorch and 4M framework modules. It also sets performance flags for TF32 operations on NVIDIA GPUs.

```python
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
current_folder = globals()['_dh'][0]
os.chdir(os.path.dirname(os.path.abspath(current_folder)))

%load_ext autoreload
%autoreload 2

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_grad_enabled(False)
```

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

### Configure and execute super-resolution generation

Source: https://github.com/apple/ml-4m/blob/main/notebooks/generation_4M-7.ipynb

Defines the generation schedule for super-resolution, initializes input and target modalities, and runs the sampler to generate high-resolution tokens. It relies on pre-defined modality information and a text tokenizer.

```python
cond_domains_sr = ['caption', 'det', 'tok_clip@224', 'tok_depth@224', 'tok_normal@224', 'tok_semseg@224', 'tok_rgb@224']
target_domains_sr = ['tok_clip@448', 'tok_depth@448', 'tok_normal@448', 'tok_semseg@448', 'tok_rgb@448']
tokens_per_target_sr = [784, 784, 784, 784, 784]
autoregression_schemes_sr = ['maskgit', 'maskgit', 'maskgit', 'maskgit', 'maskgit']
decoding_steps_sr = [8, 8, 8, 8, 8]
token_decoding_schedules_sr = ['cosine', 'cosine', 'cosine', 'cosine', 'cosine']
temps_sr = [1.0, 1.0, 1.0, 1.0, 1.0]
temp_schedules_sr = ['constant', 'constant', 'constant', 'constant', 'constant']
cfg_scales_sr = [2.0, 2.0, 2.0, 2.0, 2.0]
cfg_schedules_sr = ['constant', 'constant', 'constant', 'constant', 'constant']
cfg_grow_conditioning_sr = True
top_p_sr, top_k_sr = 0.0, 0.0

schedule_sr = build_chained_generation_schedules(
    cond_domains=cond_domains_sr, target_domains=target_domains_sr, tokens_per_target=tokens_per_target_sr, autoregression_schemes=autoregression_schemes_sr, 
    decoding_steps=decoding_steps_sr, token_decoding_schedules=token_decoding_schedules_sr, temps=temps_sr, temp_schedules=temp_schedules_sr,
    cfg_scales=cfg_scales_sr, cfg_schedules=cfg_schedules_sr, cfg_grow_conditioning=cfg_grow_conditioning_sr, 
)

# Initialize input modalities
for cond_mod in cond_domains_sr:
    batched_sample = init_full_input_modality(out_dict, MODALITY_INFO, cond_mod, device, eos_id=text_tok.token_to_id("[EOS]"))

# Initialize target modalities
for target_mod, ntoks in zip(target_domains_sr, tokens_per_target_sr):
    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)

out_dict_sr = sampler_sr.generate(
    batched_sample, schedule_sr, text_tokenizer=text_tok, 
    verbose=True, seed=1,
    top_p=top_p, top_k=top_k,
)
dec_dict_sr = decode_dict(
    out_dict_sr, toks, text_tok, 
    image_size=448, patch_size=16,
    decoding_steps=50
)
```

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

### Load 4M Model Checkpoints Manually (Python)

Source: https://github.com/apple/ml-4m/blob/main/README.md

Provides instructions for loading 4M model checkpoints manually by first downloading the safetensors files. This method involves loading the checkpoint and configuration separately and then initializing the FM model.

```python
from fourm.utils import load_safetensors
from fourm.models.fm import FM

ckpt, config = load_safetensors('/path/to/checkpoint.safetensors')
fm = FM(config=config)
fm.load_state_dict(ckpt)
```

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

### Encode and Decode Images with 4M Tokenizers

Source: https://context7.com/apple/ml-4m/llms.txt

Demonstrates how to convert images to tokens using 4M tokenizers and reconstruct them back into images. It covers both individual encoding/decoding steps and the shorthand autoencode method.

```python
with torch.no_grad():
    quant, code_loss, tokens = tok_rgb.encode(img_tensor)
    print(f"Token shape: {tokens.shape}, Token range: [{tokens.min()}, {tokens.max()}]")

tokens = tok_rgb.tokenize(img_tensor)
tokens_flat = rearrange(tokens, 'b h w -> b (h w)')

with torch.no_grad():
    reconstructed = tok_rgb.decode_tokens(tokens, timesteps=50, image_size=224, verbose=True)

rec_img = (reconstructed * 0.5 + 0.5).clamp(0, 1)
rec_pil = Image.fromarray((rec_img[0].permute(1,2,0).cpu().numpy() * 255).astype('uint8'))
rec_pil.save('reconstructed.png')

with torch.no_grad():
    autoencoded = tok_rgb.autoencode(img_tensor, timesteps=50, verbose=True)
```