> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.datasets.diffusion.collate_fns

Collate functions and dataloader builders for multiresolution diffusion training.

Supports both image and video pipelines via the FlowMatchingPipeline
expected batch format.

## Module Contents

### Functions

| Name                                                                                                                                                         | Description                                                                   |
| ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------- |
| [`_build_multiresolution_dataloader_core`](#nemo_automodel-components-datasets-diffusion-collate_fns-_build_multiresolution_dataloader_core)                 | Internal helper: create sampler + DataLoader from dataset and collate fn.     |
| [`_stack_or_pad_text_tensors`](#nemo_automodel-components-datasets-diffusion-collate_fns-_stack_or_pad_text_tensors)                                         | Stack text tensors, padding variable sequence lengths on the first dimension. |
| [`build_text_to_image_multiresolution_dataloader`](#nemo_automodel-components-datasets-diffusion-collate_fns-build_text_to_image_multiresolution_dataloader) | Build a text-to-image multiresolution dataloader for TrainDiffusionRecipe.    |
| [`build_video_multiresolution_dataloader`](#nemo_automodel-components-datasets-diffusion-collate_fns-build_video_multiresolution_dataloader)                 | Build a multiresolution video dataloader for TrainDiffusionRecipe.            |
| [`collate_fn_production`](#nemo_automodel-components-datasets-diffusion-collate_fns-collate_fn_production)                                                   | Production collate function with verification.                                |
| [`collate_fn_text_to_image`](#nemo_automodel-components-datasets-diffusion-collate_fns-collate_fn_text_to_image)                                             | Text-to-image collate function that transforms multiresolution batch output   |
| [`collate_fn_video`](#nemo_automodel-components-datasets-diffusion-collate_fns-collate_fn_video)                                                             | Video-compatible collate function for multiresolution video training.         |

### Data

[`logger`](#nemo_automodel-components-datasets-diffusion-collate_fns-logger)

### API

```python
nemo_automodel.components.datasets.diffusion.collate_fns._build_multiresolution_dataloader_core(
    dataset,
    collate_fn: typing.Callable,
    batch_size: int,
    dp_rank: int,
    dp_world_size: int,
    base_resolution: typing.Tuple[int, int] = (512, 512),
    drop_last: bool = True,
    shuffle: bool = True,
    dynamic_batch_size: bool = False,
    num_workers: int = 4,
    pin_memory: bool = True,
    prefetch_factor: int = 2
) -> typing.Tuple[torchdata.stateful_dataloader.StatefulDataLoader, nemo_automodel.components.datasets.diffusion.sampler.SequentialBucketSampler]
```

Internal helper: create sampler + DataLoader from dataset and collate fn.

```python
nemo_automodel.components.datasets.diffusion.collate_fns._stack_or_pad_text_tensors(
    tensors: typing.List[torch.Tensor],
    sequence_length_multiple: int = 1
) -> torch.Tensor
```

Stack text tensors, padding variable sequence lengths on the first dimension.

```python
nemo_automodel.components.datasets.diffusion.collate_fns.build_text_to_image_multiresolution_dataloader(
    cache_dir: str,
    train_text_encoder: bool = False,
    batch_size: int = 1,
    dp_rank: int = 0,
    dp_world_size: int = 1,
    base_resolution: typing.Tuple[int, int] = (256, 256),
    drop_last: bool = True,
    shuffle: bool = True,
    dynamic_batch_size: bool = False,
    num_workers: int = 4,
    pin_memory: bool = True,
    prefetch_factor: int = 2
) -> typing.Tuple[torchdata.stateful_dataloader.StatefulDataLoader, nemo_automodel.components.datasets.diffusion.sampler.SequentialBucketSampler]
```

Build a text-to-image multiresolution dataloader for TrainDiffusionRecipe.

This wraps the existing TextToImageDataset and SequentialBucketSampler
with a text-to-image collate function.

**Parameters:**

Directory containing preprocessed cache (metadata.json, shards, and resolution subdirs)

If True, returns tokens instead of embeddings

Batch size per GPU

Data parallel rank

Data parallel world size

Base resolution for dynamic batch sizing

Drop incomplete batches

Shuffle data

Scale batch size by resolution

DataLoader workers

Pin memory for GPU transfer

Prefetch batches per worker

**Returns:** `Tuple[StatefulDataLoader, SequentialBucketSampler]`

Tuple of (DataLoader, SequentialBucketSampler)

```python
nemo_automodel.components.datasets.diffusion.collate_fns.build_video_multiresolution_dataloader(
    cache_dir: str,
    model_type: str = 'wan',
    device: str = 'cpu',
    batch_size: int = 1,
    dp_rank: int = 0,
    dp_world_size: int = 1,
    base_resolution: typing.Tuple[int, int] = (512, 512),
    drop_last: bool = True,
    shuffle: bool = True,
    dynamic_batch_size: bool = False,
    num_workers: int = 2,
    pin_memory: bool = True,
    prefetch_factor: int = 2
) -> typing.Tuple[torchdata.stateful_dataloader.StatefulDataLoader, nemo_automodel.components.datasets.diffusion.sampler.SequentialBucketSampler]
```

Build a multiresolution video dataloader for TrainDiffusionRecipe.

Uses TextToVideoDataset with SequentialBucketSampler for bucket-based
multiresolution video training (e.g. Wan, Hunyuan).

**Parameters:**

Directory containing preprocessed cache (metadata.json + shards + WxH/\*.meta)

Model type ("wan", "hunyuan", etc.)

Device to load tensors to

Batch size per GPU

Data parallel rank

Data parallel world size

Base resolution for dynamic batch sizing

Drop incomplete batches

Shuffle data

Scale batch size by resolution

DataLoader workers

Pin memory for GPU transfer

Prefetch batches per worker

**Returns:** `Tuple[StatefulDataLoader, SequentialBucketSampler]`

Tuple of (DataLoader, SequentialBucketSampler)

```python
nemo_automodel.components.datasets.diffusion.collate_fns.collate_fn_production(
    batch: typing.List[typing.Dict]
) -> typing.Dict
```

Production collate function with verification.

```python
nemo_automodel.components.datasets.diffusion.collate_fns.collate_fn_text_to_image(
    batch: typing.List[typing.Dict]
) -> typing.Dict
```

Text-to-image collate function that transforms multiresolution batch output
to match FlowMatchingPipeline expected format.

**Parameters:**

List of samples from TextToImageDataset

**Returns:** `Dict`

Dict compatible with FlowMatchingPipeline.step()

```python
nemo_automodel.components.datasets.diffusion.collate_fns.collate_fn_video(
    batch: typing.List[typing.Dict],
    model_type: str = 'wan'
) -> typing.Dict
```

Video-compatible collate function for multiresolution video training.

Concatenates video\_latents (5D) and text\_embeddings (3D) along the batch dim,
matching the format expected by FlowMatchingPipeline with SimpleAdapter.

**Parameters:**

List of samples from TextToVideoDataset

Model type for model-specific field handling

**Returns:** `Dict`

Dict compatible with FlowMatchingPipeline.step()

```python
nemo_automodel.components.datasets.diffusion.collate_fns.logger = logging.getLogger(__name__)
```