Diffusion Model Fine-Tuning with NeMo AutoModel

Introduction

Diffusion models generate images and videos by learning to reverse a noise process, starting from random noise and iteratively refining it into coherent visual output guided by a text prompt. Pretrained diffusion models, such as FLUX.1-dev for images or Wan 2.1 for video, produce impressive general-purpose results, but they know nothing about your particular visual domain, style, or subject matter. Fine-tuning bridges that gap: you adapt the model on your own data so it produces outputs that match your requirements, without the cost of training from scratch.

Under the hood, NeMo AutoModel uses flow matching, a modern generative framework that learns to transform noise into data by regressing a velocity field along straight interpolation paths. It integrates with Hugging Face Diffusers to provide distributed fine-tuning for text-to-image and text-to-video models. This guide walks you through the process end to end, from installation through training and inference, using Wan 2.1 T2V 1.3B as a running example.

Workflow Overview

┌──────────────┐    ┌──────────────┐    ┌──────────────┐    ┌──────────────┐    ┌──────────────┐
│ 1. Install   │--->│ 2. Prepare   │--->│ 3. Configure │--->│  4. Train    │--->│ 5. Generate  │
│              │    │    Data      │    │              │    │              │    │              │
│ uv sync      │    │ Encode to    │    │ YAML recipe  │    │ torchrun     │    │ Run inference│
│ or Docker    │    │ cache files  │    │              │    │              │    │ with ckpt    │
└──────────────┘    └──────────────┘    └──────────────┘    └──────────────┘    └──────────────┘

For model-specific configuration (FLUX.1-dev, HunyuanVideo), see Model-Specific Notes.

Supported Models

All models use FSDP2 for distributed training and flow matching for loss computation.

Install NeMo AutoModel

Diffusion fine-tuning needs the diffusion extra plus diffusion-media (OpenCV Diffusion fine-tuning requires the diffusion and diffusion-media extras. The diffusion-media extra provides OpenCV for image and video preprocessing and imageio-ffmpeg for T2V export:

$
uv venv
$
source .venv/bin/activate
$
uv pip install "nemo-automodel[diffusion,diffusion-media]"

Alternatively, if you run into dependency or driver issues, use the prebuilt Docker container:

$
docker pull nvcr.io/nvidia/nemo-automodel:26.06.00
$
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.06.00

$
cd /opt/Automodel && uv pip install ".[diffusion-media]"

For the full set of installation methods, see the installation guide.

Prepare Your Dataset

Diffusion models operate in latent space — a compressed representation of visual data — rather than directly on raw images or videos. To avoid re-encoding data on every training step, the preprocessing pipeline encodes all inputs ahead of time and saves them as .meta or .pt cache files.

Each cache file contains:

• Latent representations produced by a variational autoencoder (VAE) from the raw visual data
• Text embeddings produced by a text encoder from the associated captions or prompts

Fine-tuning then operates entirely on these pre-encoded cache files, which is significantly faster than encoding on the fly.

Preprocess your data using the built-in tool at tools/diffusion/preprocessing_multiprocess.py. The script provides image and video subcommands.

Video preprocessing (using Wan 2.1 as a running example):

$
python -m tools.diffusion.preprocessing_multiprocess video \
>
    --video_dir /data/videos \
>
    --output_dir /cache \
>
    --processor wan \
>
    --resolution_preset 512p \
>
    --caption_format sidecar

Image preprocessing (FLUX):

$
python -m tools.diffusion.preprocessing_multiprocess image \
>
    --image_dir /data/images \
>
    --output_dir /cache \
>
    --processor flux

Video preprocessing (HunyuanVideo):

$
python -m tools.diffusion.preprocessing_multiprocess video \
>
    --video_dir /data/videos \
>
    --output_dir /cache \
>
    --processor hunyuan \
>
    --target_frames 121 \
>
    --caption_format meta_json

Video and audio preprocessing (LTX-2.3):

$
python -m tools.diffusion.preprocessing_multiprocess video \
>
    --video_dir /data/videos \
>
    --output_dir /cache \
>
    --processor ltx2 \
>
    --num_frames 121 \
>
    --output_format pt \
>
    --resolution_preset 512p \
>
    --caption_format sidecar

LTX-2.3 requires an explicit 8n+1 frame count and uses each source video’s audio track. If a clip has no audio stream, preprocessing stores silence of the matching duration.

For the full set of arguments and input format details, see the Diffusion Dataset Preparation guide.

Configure Your Training Recipe

Fine-tuning is driven by two components:

1. A recipe script (e.g., train.py) — the Python entry point that orchestrates the training loop: loading the model, building the dataloader, running forward/backward passes, computing the flow matching loss, checkpointing, and logging.
2. A YAML configuration file — a text file in YAML format that specifies all settings the recipe uses: which model to fine-tune, where the data lives, optimizer hyperparameters, parallelism strategy, etc.
3. A recipe script, for example train.py, is the Python entry point that orchestrates the training loop: loading the model, building the data loader, running forward and backward passes, computing the flow matching loss, checkpointing, and logging.
4. A YAML configuration file is a text file in YAML format that specifies all settings the recipe uses: which model to fine-tune, where the data lives, optimizer hyperparameters, parallelism strategy, and other settings. You customize training by editing this file rather than modifying code, allowing you to scale seamlessly from 1 to hundreds of GPUs.

The following annotated wan2_1_t2v_flow.yaml config explains each section:

1
seed: 42
2
3
# Weights & Biases experiment tracking
4
wandb:
5
  project: wan-t2v-flow-matching
6
  mode: online
7
  name: wan2_1_t2v_fm_v2
8
9
dist_env:
10
  backend: nccl
11
  timeout_minutes: 30
12
13
# Model configuration
14
# pretrained_model_name_or_path: Hugging Face model ID
15
# mode: "finetune" loads pretrained weights and adapts them to your dataset
16
model:
17
  pretrained_model_name_or_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
18
  mode: finetune
19
20
# Training schedule
21
step_scheduler:
22
  global_batch_size: 8       # Effective batch size across all GPUs
23
  local_batch_size: 1        # Per-GPU batch size (gradient accumulation = global/local/num_gpus)
24
  ckpt_every_steps: 1000     # Checkpoint frequency
25
  num_epochs: 100
26
  log_remote_every_steps: 2  # Log metrics every N steps
27
  save_checkpoint_every_epoch: false
28
29
# Data: uses pre-encoded .meta files
30
data:
31
  dataloader:
32
    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
33
    cache_dir: PATH_TO_YOUR_DATA
34
    model_type: wan # "wan" for Wan 2.1, "hunyuan" for HunyuanVideo
35
    base_resolution: [512, 512]
36
    dynamic_batch_size: false
37
    shuffle: true
38
    drop_last: false
39
    num_workers: 0
40
41
# Optimizer
42
optimizer:
43
  _target_: torch.optim.AdamW
44
  lr: 5e-6
45
  weight_decay: 0.01
46
  betas: [0.9, 0.999]
47
48
# Learning rate scheduler
49
lr_scheduler:
50
  lr_decay_style: cosine
51
  lr_warmup_steps: 0
52
  min_lr: 1e-6
53
54
# Flow matching configuration
55
flow_matching:
56
  adapter_type: "simple"          # Model-specific adapter (simple, flux, hunyuan)
57
  adapter_kwargs: {}
58
  timestep_sampling: "uniform"    # How timesteps are sampled during training
59
  logit_mean: 0.0
60
  logit_std: 1.0
61
  flow_shift: 3.0                # Shifts the flow schedule
62
  mix_uniform_ratio: 0.1
63
  sigma_min: 0.0
64
  sigma_max: 1.0
65
  num_train_timesteps: 1000
66
  i2v_prob: 0.3                  # Probability of image-to-video conditioning
67
  use_loss_weighting: true
68
  log_interval: 100
69
  summary_log_interval: 10
70
71
# FSDP2 distributed training
72
fsdp:
73
  tp_size: 1      # Tensor parallelism
74
  cp_size: 1      # Context parallelism
75
  pp_size: 1      # Pipeline parallelism
76
  dp_replicate_size: 1
77
  dp_size: 8      # Data parallelism (number of GPUs)
78
79
# Checkpointing
80
checkpoint:
81
  enabled: true
82
  checkpoint_dir: PATH_TO_YOUR_CKPT_DIR
83
  model_save_format: torch_save
84
  save_consolidated: false
85
  restore_from: null

Config Field Reference

Fine-Tune the Model

Launch fine-tuning with torchrun:

$
torchrun --nproc-per-node=8 \
>
  examples/diffusion/finetune/finetune.py \
>
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml

Adjust --nproc-per-node to match the number of GPUs on your node, and ensure fsdp.dp_size in the YAML matches.

Multi-Node Training

When a single node does not provide enough GPUs or memory for your workload, you can scale training across multiple nodes. NeMo AutoModel handles multi-node distributed training through torchrun rendezvous and FSDP2. The same recipe script works on one node or many.

YAML Configuration Changes

The main change is in the fsdp section. Set dp_size to the total number of GPUs across all nodes, and optionally increase dp_replicate_size for gradient replication across nodes.

For example, to train on 2 nodes with 8 GPUs each (16 GPUs total):

1
fsdp:
2
  tp_size: 1
3
  cp_size: 1
4
  pp_size: 1
5
  dp_replicate_size: 2   # Replicate across 2 nodes for robustness
6
  dp_size: 16             # Total GPUs: 2 nodes × 8 GPUs

A complete multi-node config is provided at wan2_1_t2v_flow_multinode.yaml.

Launch with torchrun

Run the following command on each node, setting NODE_RANK to 0 on the first node, 1 on the second, and so on:

$
export MASTER_ADDR=node0.hostname   # hostname or IP of the first node
$
export MASTER_PORT=29500
$
export NODE_RANK=0                  # 0 on master, 1 on second node, etc.
$
$
torchrun \
>
  --nnodes=2 \
>
  --nproc-per-node=8 \
>
  --node_rank=${NODE_RANK} \
>
  --rdzv_backend=c10d \
>
  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
>
  examples/diffusion/finetune/finetune.py \
>
  -c examples/diffusion/finetune/wan2_1_t2v_flow_multinode.yaml

Model-Specific Notes

Use the following table to select a model for your use case:

Wan 2.1 T2V 1.3B

• Adapter type: simple
• Dataloader: build_video_multiresolution_dataloader with model_type: wan
• Config: wan2_1_t2v_flow.yaml

FLUX.1-dev (Text-to-Image)

• Adapter type: flux
• Dataloader: build_text_to_image_multiresolution_dataloader
• Key differences:
  • Uses pipeline_spec to specify the transformer architecture:

    1
    model:
    2
      pipeline_spec:
    3
        transformer_cls: "FluxTransformer2DModel"
    4
        subfolder: "transformer"
    5
        load_full_pipeline: false
  • Requires guidance_scale in adapter kwargs:

    1
    flow_matching:
    2
      adapter_type: "flux"
    3
      adapter_kwargs:
    4
        guidance_scale: 3.5
    5
        use_guidance_embeds: true
  • Uses logit_normal timestep sampling instead of uniform
• Config: flux_t2i_flow.yaml

HunyuanVideo 1.5

• Adapter type: hunyuan
• Dataloader: build_video_multiresolution_dataloader with model_type: hunyuan
• Key differences:
  • Requires activation_checkpointing: true in FSDP config due to model size
  • Uses condition latents in adapter kwargs:

    1
    flow_matching:
    2
      adapter_type: "hunyuan"
    3
      adapter_kwargs:
    4
        use_condition_latents: true
    5
        default_image_embed_shape: [729, 1152]
  • Uses logit_normal timestep sampling
• Config: hunyuan_t2v_flow.yaml

LTX-2.3

• Adapter type: ltx2
• Dataloader: build_video_multiresolution_dataloader with model_type: ltx2
• Preprocessing: use --processor ltx2, an explicit 8n+1 frame count, and --output_format pt
• Loss: applies a shared noise level to synchronized video and audio latents and adds the weighted audio loss
• Configs: full fine-tuning and LoRA

Generation and Inference

After training is complete, you can use the model to generate images or videos from text prompts. This step is called inference. Unlike training, where the model learns from data, inference is where it produces new outputs.

In diffusion models, generation works by starting from random noise and iteratively denoising it, guided by your text prompt, until a clean image or video emerges.

The generation script (generate.py) handles this: it loads your model weights (pretrained or fine-tuned), configures the diffusion sampler, and produces outputs for one or more prompts.

Single-GPU (Wan 2.1 1.3B):

$
python examples/diffusion/generate/generate.py \
>
  -c examples/diffusion/generate/configs/generate_wan.yaml

Multi-GPU (Wan 2.1 1.3B):

Wan 2.1 supports tensor parallelism for inference, which shards the transformer across GPUs to reduce per-GPU memory. Pass the distributed config using CLI overrides:

$
torchrun --nproc-per-node=8 \
>
  examples/diffusion/generate/generate.py \
>
  -c examples/diffusion/generate/configs/generate_wan.yaml \
>
  --distributed.backend nccl \
>
  --distributed.parallel_scheme.transformer.tp_size 8

With a fine-tuned checkpoint:

$
python examples/diffusion/generate/generate.py \
>
  -c examples/diffusion/generate/configs/generate_wan.yaml \
>
  --model.checkpoint ./checkpoints/step_1000 \
>
  --inference.prompts '["A dog running on a beach"]'

FLUX image generation:

$
python examples/diffusion/generate/generate.py \
>
  -c examples/diffusion/generate/configs/generate_flux.yaml

HunyuanVideo:

$
python examples/diffusion/generate/generate.py \
>
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml

LTX-2.3 video with audio:

$
python examples/diffusion/generate/generate.py \
>
  -c examples/diffusion/generate/configs/generate_ltx2.yaml

Available Generation Configs

Hardware Requirements