Vace finetuning

[Tatiana21]

Code for:

1. Creating segmentation masks for Inpainting tasks for a video dataset, based on open-sora-plan dataset format
2. Preprocessing and creating energon dataset for finetuing with T2V, I2V and V2V tasks
3. Finetuning vace, for T2V, I2V and V2V tasks

Refer to annotators/Inpainting/run_batch_process.sh for segmentation.
Refer to example_commands.sh for commands to process datasets and launch training.

[Copilot]

Pull request overview

This pull request introduces comprehensive support for Vace (Video Auto-Context Encoding) finetuning capabilities, including tools for dataset preparation, preprocessing pipelines, and training infrastructure for T2V (Text-to-Video), I2V (Image-to-Video), and V2V (Video-to-Video) tasks. The implementation extends the existing WAN (Wide Attention Network) model architecture with VACE-specific layers and flow-matching training pipelines.

Key Changes

• Added VACE model architecture with context and base layers for video editing tasks
• Implemented flow matching training pipeline with configurable timestep sampling strategies
• Created preprocessing utilities for video/image/mask data and segmentation mask generation
• Added Gemma model family support (Gemma 1.0 and Gemma 2.0) with proper embedding scaling

Reviewed changes

Copilot reviewed 108 out of 229 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/megatron/bridge/models/wan/wan_layer_spec.py	Defines WAN transformer layer specifications including VACE-specific base and context layers with adaptive layer normalization
src/megatron/bridge/models/wan/wan_bridge.py	Implements parameter mapping bridges between HuggingFace and Megatron formats for WAN and VACE models
src/megatron/bridge/models/wan/utils/utils.py	Provides utility functions for grid size calculation, patching/unpatching, and context parallelism operations
src/megatron/bridge/models/wan/utils/preprocessor.py	Implements video and image preprocessing classes with resizing, cropping, and normalization capabilities
src/megatron/bridge/models/wan/rope_utils.py	Implements 3D RoPE (Rotary Position Embeddings) for spatial-temporal attention in video models
src/megatron/bridge/models/wan/modules/vae.py	Defines VAE encoder/decoder architecture with causal 3D convolutions for video latent encoding
src/megatron/bridge/models/wan/modules/tokenizers.py	Provides HuggingFace tokenizer wrapper with text cleaning utilities
src/megatron/bridge/models/wan/modules/t5.py	Implements T5 encoder/decoder models with custom layer normalization and attention mechanisms
src/megatron/bridge/models/wan/flow_matching/time_shift_utils.py	Implements timestep sampling strategies and sigma computation for flow matching training
src/megatron/bridge/models/wan/flow_matching/flow_pipeline.py	Defines training pipeline for flow matching with support for both WAN and VACE models
src/megatron/bridge/models/wan/flow_matching/flow_inference_pipeline.py	Implements inference pipeline with DPM/UniPC solvers and pipeline parallelism support
src/megatron/bridge/models/wan/inference/configs/*.py	Configuration files for different WAN model variants (T2V, I2V, VACE) with size-specific settings
src/megatron/bridge/models/gemma/*.py	Adds complete Gemma model family support with proper embedding scaling and configuration mappings

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Corrected spelling of 'becuase' to 'because' in both comments.

Suggested change

	query = query.contiguous() # important becuase TE attention expects contiguous tensors
	key = key.contiguous() # important becuase TE attention expects contiguous tensors
	query = query.contiguous() # important because TE attention expects contiguous tensors
	key = key.contiguous() # important because TE attention expects contiguous tensors

[Copilot]

Corrected spelling of 'becuase' to 'because' in both comments.

Suggested change

	query = query.contiguous() # important becuase TE attention expects contiguous tensors
	key = key.contiguous() # important becuase TE attention expects contiguous tensors
	query = query.contiguous() # important because TE attention expects contiguous tensors
	key = key.contiguous() # important because TE attention expects contiguous tensors

[Copilot]

Corrected spelling of 'convolusion' to 'convolution' in docstring.

Suggested change

	Causal 3d convolusion.
	Causal 3d convolution.

[Copilot]

Parameter names should follow snake_case convention. These should be input_frames, input_masks, and input_ref_images instead of capitalized versions.

Suggested change

	Input_frames (`list[Tensor]`):
	Input frames for content generation
	Input_masks (`list[Tensor]`):
	Input masks for content generation
	Input_ref_images (`list[Tensor]`):
	input_frames (`list[Tensor]`):
	Input frames for content generation
	input_masks (`list[Tensor]`):
	Input masks for content generation
	input_ref_images (`list[Tensor]`):

[Copilot]

Corrected spelling of 'Thism' to 'This' in docstring.

Suggested change

	Thism model has an identical configuration to GemmaModelProvider2B.
	This model has an identical configuration to GemmaModelProvider2B.