bridge.diffusion.data.wan.wan_taskencoder

Module Contents

Classes

WanTaskEncoder

Task encoder for Wan dataset. .. attribute:: cookers

Functions

cook	Processes a raw sample dictionary from energon dataset and returns a new dictionary with specific keys.

API

bridge.diffusion.data.wan.wan_taskencoder.cook(sample: dict) → dict

Processes a raw sample dictionary from energon dataset and returns a new dictionary with specific keys.

Parameters:

sample (dict) – The input dictionary containing the raw sample data.

Returns:

A new dictionary containing the processed sample data with the following keys: - All keys from the result of basic_sample_keys(sample) - ‘json’: The contains meta data like resolution, aspect ratio, fps, etc. - ‘pth’: contains video latent tensor - ‘pickle’: contains text embeddings

Return type:

dict

class bridge.diffusion.data.wan.wan_taskencoder.WanTaskEncoder(

*args,

max_frames: int = None,

patch_spatial: int = 2,

patch_temporal: int = 1,

seq_length: int = 1024,

**kwargs,

)

Bases: megatron.bridge.diffusion.data.common.diffusion_task_encoder_with_sp.DiffusionTaskEncoderWithSequencePacking

Task encoder for Wan dataset. .. attribute:: cookers

A list of Cooker objects used for processing.

Type:

list

.. attribute:: patch_spatial

The spatial patch size. Defaults to 2.

Type:

int

.. attribute:: patch_temporal

The temporal patch size. Defaults to 1.

Type:

int

.. attribute:: seq_length

The sequence length. Defaults to 1024.

Type:

int

Initialization

cookers

None

encode_sample(sample: dict) → dict

batch(

samples: List[megatron.bridge.diffusion.data.common.diffusion_sample.DiffusionSample],

) → dict

Return dictionary with data for batch.