lint, ruff

Author: Huy Vu2

dfm/src/megatron/data/wan/wan_mock_energon_datamodule.py

@@ -17,11 +17,12 @@
 from dataclasses import dataclass
 
 import torch
-from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
 from torch.utils.data import DataLoader, Dataset
+from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
 
 from dfm.src.megatron.model.wan.utils import patchify
 
+
 class _MockDataset(Dataset):
     def __init__(self, length: int):
         self.length = max(int(length), 1)
@@ -32,6 +33,7 @@ def __len__(self) -> int:
     def __getitem__(self, idx: int) -> dict:
         return {}
 
+
 def mock_batch(
     F_latents: int,
     H_latents: int,
@@ -92,6 +94,7 @@ def mock_batch(
 
     return batch
 
+
 @dataclass(kw_only=True)
 class WanMockDataModuleConfig(DatasetProvider):
     path: str = ""

dfm/src/megatron/data/wan/wan_taskencoder.py

@@ -142,7 +142,7 @@ def encode_sample(self, sample: dict) -> dict:
             seq_len_q=torch.tensor([seq_len_q], dtype=torch.int32),
             seq_len_q_padded=torch.tensor([seq_len_q_padded], dtype=torch.int32),
             seq_len_kv=torch.tensor([seq_len_kv], dtype=torch.int32),
-            pos_ids=torch.zeros(1, dtype=torch.bfloat16), # dummy pos_ids
+            pos_ids=torch.zeros(1, dtype=torch.bfloat16),  # dummy pos_ids
             video_metadata=video_metadata,
         )
 

dfm/src/megatron/model/wan/flow_matching/flow_pipeline.py

@@ -43,8 +43,8 @@ def training_step(
         logit_std: float = 1.0,
         flow_shift: float = 3.0,
         mix_uniform_ratio: float = 0.1,
-        sigma_min: float = 0.0, # Default: no clamping (pretrain)
-        sigma_max: float = 1.0, # Default: no clamping (pretrain)
+        sigma_min: float = 0.0,  # Default: no clamping (pretrain)
+        sigma_max: float = 1.0,  # Default: no clamping (pretrain)
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
         """
         Performs a single training step using flow matching algorithm.

dfm/src/megatron/model/wan/wan_provider.py

@@ -30,7 +30,7 @@
 
 @dataclass
 class WanModelProvider(TransformerConfig, ModelProviderMixin[VisionModule]):
-    crossattn_emb_size: int = 1536 # cross attention emebedding size after linear projection
+    crossattn_emb_size: int = 1536  # cross attention emebedding size after linear projection
     add_bias_linear: bool = True
     gated_linear_unit: bool = False