fix Ruff, lint

Author: Huy Vu2

dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import random
 from abc import ABC, abstractmethod
 from typing import List

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

@@ -16,10 +16,10 @@
 
 from dataclasses import dataclass
 
-from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
+from megatron.bridge.data.utils import DatasetBuildContext
 from torch import int_repr
 
-from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModuleConfig, DiffusionDataModule
+from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModule, DiffusionDataModuleConfig
 from dfm.src.megatron.data.wan.wan_taskencoder import WanTaskEncoder
 
 

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
+
 from dfm.src.megatron.data.common.diffusion_sample import DiffusionSample
 
 

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

@@ -14,17 +14,17 @@
 
 # pylint: disable=C0115,C0116,C0301
 
-from torch._tensor import Tensor
-import torch
-import torch.nn.functional as F
-from megatron.energon.task_encoder.base import stateless
-from megatron.core import parallel_state
 from typing import List
-from megatron.energon import SkipSample
+
 from dfm.src.megatron.data.common.diffusion_task_encoder_with_sp import DiffusionTaskEncoderWithSequencePacking
-from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
 from dfm.src.megatron.data.wan.wan_sample import WanSample
 from dfm.src.megatron.model.wan.utils import grid_sizes_calculation, patchify
+from megatron.core import parallel_state
+from megatron.energon import SkipSample
+from megatron.energon.task_encoder.base import stateless
+from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
+import torch
+import torch.nn.functional as F
 
 
 def cook(sample: dict) -> dict:

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

@@ -69,7 +69,7 @@ def forward(self, n_head, dim_head, cu_seqlens_q_padded, grid_sizes, device):
 
         # Pad freqs_real_i to (padded_seq_len, 1, 1, dim_head) with 0s
         for i, freqs_real_i in enumerate(freqs_real):
-            seq_len_q_padded = cu_seqlens_q_padded[i+1] - cu_seqlens_q_padded[i]
+            seq_len_q_padded = cu_seqlens_q_padded[i + 1] - cu_seqlens_q_padded[i]
             if freqs_real_i.shape[0] < seq_len_q_padded:
                 pad_shape = (seq_len_q_padded - freqs_real_i.shape[0], 1, 1, dim_head)
                 freqs_real_i = torch.cat(

dfm/src/megatron/recipes/wan/wan.py

@@ -180,7 +180,7 @@ def pretrain_config(
             global_batch_size=global_batch_size,
             num_workers=10,
             task_encoder_seq_length=None,
-            packing_buffer_size=131072, # 131,072 = 2^17 tokens, each 5 secs of 832*480 is about 45k tokens
+            packing_buffer_size=131072,  # 131,072 = 2^17 tokens, each 5 secs of 832*480 is about 45k tokens
         )
 
     # Config Container

examples/megatron/recipes/wan/prepare_energon_dataset_wan.py

@@ -18,12 +18,11 @@
 from typing import Dict, List, Optional, Tuple
 
 import cv2
+from diffusers import AutoencoderKLWan
 import numpy as np
 import torch
-import webdataset as wds
-
-from diffusers import AutoencoderKLWan
 from transformers import AutoTokenizer, UMT5EncoderModel
+import webdataset as wds
 
 
 def _map_interpolation(resize_mode: str) -> int:
@@ -412,7 +411,7 @@ def main():
         for index, meta in enumerate(metadata_list):
             video_name = meta["file_name"]
             start_frame = int(meta["start_frame"])  # inclusive
-            end_frame = int(meta["end_frame"])      # inclusive
+            end_frame = int(meta["end_frame"])  # inclusive
             caption_text = meta.get("vila_caption", "")
 
             video_path = str(video_folder / video_name)
@@ -431,7 +430,9 @@ def main():
 
                 # Encode text and video with HF models exactly like automodel
                 text_embed = _encode_text(tokenizer, text_encoder, args.device, caption_text)
-                latents = _encode_video_latents(vae, args.device, video_tensor, deterministic_latents=not args.stochastic)
+                latents = _encode_video_latents(
+                    vae, args.device, video_tensor, deterministic_latents=not args.stochastic
+                )
 
                 # Move to CPU without changing dtype; keep exact values to match automodel outputs
                 text_embed_cpu = text_embed.detach().to(device="cpu")