lint

Author: wwwjn

torchtitan/hf_datasets/flux_dataset.py

@@ -18,6 +18,7 @@
 from torch.distributed.checkpoint.stateful import Stateful
 
 from torch.utils.data import IterableDataset
+
 from torchtitan.components.dataloader import ParallelAwareDataloader
 
 from torchtitan.components.tokenizer import BaseTokenizer

torchtitan/models/flux/__init__.py

@@ -4,12 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from torchtitan.components.loss import build_mse_loss
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
-from torchtitan.protocols.train_spec import TrainSpec
 
 from torchtitan.datasets.flux_dataset import build_flux_dataloader
-from torchtitan.components.loss import build_mse_loss
+from torchtitan.protocols.train_spec import TrainSpec
 from .infra.parallelize import parallelize_flux
 from .model.args import FluxModelArgs
 from .model.autoencoder import AutoEncoderParams

torchtitan/models/flux/model/hf_embedder.py

@@ -40,19 +40,15 @@ def __init__(self, version: str, random_init=False, **hf_kwargs):
                     version, **hf_kwargs
                 )
 
-
         self.hf_module = self.hf_module.eval().requires_grad_(False)
         # This is to make sure the encoders works with FSDP
         self.make_parameters_contiguous()
 
     def make_parameters_contiguous(self):
         """Make all non-contiguous parameters contiguous to avoid FSDP issues."""
-        strided_count = 0
         for name, param in self.hf_module.named_parameters():
             if not param.is_contiguous():
-                strided_count += 1
                 param.data = param.data.contiguous()
-        
 
     def forward(self, batch_tokens: Tensor) -> Tensor:
         """

torchtitan/models/flux/tokenizer.py

@@ -11,9 +11,10 @@
 from typing import List
 
 import torch
+from transformers import CLIPTokenizer, T5Tokenizer
+
 from torchtitan.components.tokenizer import BaseTokenizer, HuggingFaceTokenizer
 from torchtitan.config import JobConfig
-from transformers import CLIPTokenizer, T5Tokenizer
 
 
 class FluxTestTokenizer(BaseTokenizer):

torchtitan/models/flux/train.py

@@ -11,8 +11,6 @@
 
 from torchtitan.config import ConfigManager, JobConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed import utils as dist_utils
-from torchtitan.tools.logging import init_logger, logger
-from torchtitan.train import Trainer
 
 from torchtitan.models.flux.infra.parallelize import parallelize_encoders
 from torchtitan.models.flux.model.autoencoder import load_ae
@@ -23,6 +21,8 @@
     preprocess_data,
     unpack_latents,
 )
+from torchtitan.tools.logging import init_logger, logger
+from torchtitan.train import Trainer
 
 
 class FluxTrainer(Trainer):

torchtitan/models/flux/validate.py

@@ -17,15 +17,13 @@
 from torchtitan.components.tokenizer import BaseTokenizer
 from torchtitan.components.validate import Validator
 from torchtitan.config import JobConfig
+from torchtitan.datasets.flux_dataset import build_flux_validation_dataloader
 from torchtitan.distributed import ParallelDims, utils as dist_utils
-from torchtitan.datasets.flux_dataset import (
-    build_flux_validation_dataloader,
-)
-
-from torchtitan.models.flux.tokenizer import build_flux_tokenizer
 from torchtitan.models.flux.model.autoencoder import AutoEncoder
 from torchtitan.models.flux.model.hf_embedder import FluxEmbedder
 from torchtitan.models.flux.sampling import generate_image, save_image
+
+from torchtitan.models.flux.tokenizer import build_flux_tokenizer
 from torchtitan.models.flux.utils import (
     create_position_encoding_for_latents,
     pack_latents,