style: upgrade the linter (#339)

* style: reformated codes

* style: reformated codes

nunchaku/csrc/sana.h

@@ -25,8 +25,7 @@ public:
         net = std::make_unique<SanaModel>(cfg, bf16 ? Tensor::BF16 : Tensor::FP16, Device::cuda((int)deviceId));
     }
 
-    torch::Tensor forward(
-        torch::Tensor hidden_states,
+    torch::Tensor forward(torch::Tensor hidden_states,
                           torch::Tensor encoder_hidden_states,
                           torch::Tensor timestep,
                           torch::Tensor cu_seqlens_img,
@@ -35,8 +34,7 @@ public:
                           int W,
                           bool pag,
                           bool cfg,
-        bool skip_first_layer = false)
-    {
+                          bool skip_first_layer = false) {
         checkModel();
         CUDADeviceContext ctx(deviceId);
 
@@ -48,16 +46,16 @@ public:
         cu_seqlens_img        = cu_seqlens_img.contiguous();
         cu_seqlens_txt        = cu_seqlens_txt.contiguous();
 
-        Tensor result = net->forward(
-            from_torch(hidden_states),
+        Tensor result = net->forward(from_torch(hidden_states),
                                      from_torch(encoder_hidden_states),
                                      from_torch(timestep),
                                      from_torch(cu_seqlens_img),
                                      from_torch(cu_seqlens_txt),
-            H, W,
-            pag, cfg,
-            skip_first_layer
-        );
+                                     H,
+                                     W,
+                                     pag,
+                                     cfg,
+                                     skip_first_layer);
 
         torch::Tensor output = to_torch(result);
         // Tensor::synchronizeDevice();
@@ -65,8 +63,7 @@ public:
         return output;
     }
 
-    torch::Tensor forward_layer(
-        int64_t idx,
+    torch::Tensor forward_layer(int64_t idx,
                                 torch::Tensor hidden_states,
                                 torch::Tensor encoder_hidden_states,
                                 torch::Tensor timestep,
@@ -75,8 +72,7 @@ public:
                                 int H,
                                 int W,
                                 bool pag,
-        bool cfg)
-    {
+                                bool cfg) {
         checkModel();
         CUDADeviceContext ctx(deviceId);
 
@@ -88,15 +84,15 @@ public:
         cu_seqlens_img        = cu_seqlens_img.contiguous();
         cu_seqlens_txt        = cu_seqlens_txt.contiguous();
 
-        Tensor result = net->transformer_blocks.at(idx)->forward(
-            from_torch(hidden_states),
+        Tensor result = net->transformer_blocks.at(idx)->forward(from_torch(hidden_states),
                                                                  from_torch(encoder_hidden_states),
                                                                  from_torch(timestep),
                                                                  from_torch(cu_seqlens_img),
                                                                  from_torch(cu_seqlens_txt),
-            H, W,
-            pag, cfg
-        );
+                                                                 H,
+                                                                 W,
+                                                                 pag,
+                                                                 cfg);
 
         torch::Tensor output = to_torch(result);
         // Tensor::synchronizeDevice();

nunchaku/csrc/utils.h

@@ -6,34 +6,34 @@
 
 namespace nunchaku::utils {
 
-    void set_cuda_stack_limit(int64_t newval) {
+void set_cuda_stack_limit(int64_t newval) {
     size_t val = 0;
     checkCUDA(cudaDeviceSetLimit(cudaLimitStackSize, (size_t)newval));
     checkCUDA(cudaDeviceGetLimit(&val, cudaLimitStackSize));
     spdlog::debug("Stack={}", val);
-    }
+}
 
-    void disable_memory_auto_release() {
+void disable_memory_auto_release() {
     int device;
     checkCUDA(cudaGetDevice(&device));
     cudaMemPool_t mempool;
     checkCUDA(cudaDeviceGetDefaultMemPool(&mempool, device));
     uint64_t threshold = UINT64_MAX;
     checkCUDA(cudaMemPoolSetAttribute(mempool, cudaMemPoolAttrReleaseThreshold, &threshold));
-    }
+}
 
-    void trim_memory() {
+void trim_memory() {
     int device;
     checkCUDA(cudaGetDevice(&device));
     cudaMemPool_t mempool;
     checkCUDA(cudaDeviceGetDefaultMemPool(&mempool, device));
     size_t bytesToKeep = 0;
     checkCUDA(cudaMemPoolTrimTo(mempool, bytesToKeep));
-    }
+}
 
-    void set_faster_i2f_mode(std::string mode) {
+void set_faster_i2f_mode(std::string mode) {
     spdlog::info("Set fasteri2f mode to {}", mode);
     kernels::set_faster_i2f_mode(mode);
-    }
+}
 
-};
\ No newline at end of file
+}; // namespace nunchaku::utils

nunchaku/lora/flux/__init__.py

 from .diffusers_converter import to_diffusers
 from .nunchaku_converter import convert_to_nunchaku_flux_lowrank_dict, to_nunchaku
 from .utils import is_nunchaku_format
+
+__all__ = ["to_diffusers", "to_nunchaku", "convert_to_nunchaku_flux_lowrank_dict", "is_nunchaku_format"]

nunchaku/lora/flux/nunchaku_converter.py

@@ -7,10 +7,10 @@ import torch
 from safetensors.torch import save_file
 from tqdm import tqdm
 
+from ...utils import filter_state_dict, load_state_dict_in_safetensors
 from .diffusers_converter import to_diffusers
 from .packer import NunchakuWeightPacker
 from .utils import is_nunchaku_format, pad
-from ...utils import filter_state_dict, load_state_dict_in_safetensors
 
 logger = logging.getLogger(__name__)
 

nunchaku/lora/flux/packer.py

 # Copy the packer from https://github.com/mit-han-lab/deepcompressor/
 import torch
 
-from .utils import pad
 from ...utils import ceil_divide
+from .utils import pad
 
 
 class MmaWeightPackerBase:

nunchaku/models/__init__.py

 from .text_encoders.t5_encoder import NunchakuT5EncoderModel
 from .transformers import NunchakuFluxTransformer2dModel, NunchakuSanaTransformer2DModel
+
+__all__ = ["NunchakuFluxTransformer2dModel", "NunchakuSanaTransformer2DModel", "NunchakuT5EncoderModel"]

nunchaku/models/pulid/eva_clip/__init__.py

 from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
 from .factory import create_model_and_transforms
+
+__all__ = ["create_model_and_transforms", "OPENAI_DATASET_MEAN", "OPENAI_DATASET_STD"]

nunchaku/models/pulid/eva_clip/eva_vit_model.py

@@ -14,8 +14,8 @@ try:
 except ImportError:
     from timm.layers import drop_path, to_2tuple, trunc_normal_
 
-from .transformer import PatchDropout
 from .rope import VisionRotaryEmbeddingFast
+from .transformer import PatchDropout
 
 if os.getenv("ENV_TYPE") == "deepspeed":
     try:
@@ -26,7 +26,7 @@ else:
     from torch.utils.checkpoint import checkpoint
 
 try:
-    import xformers
+    import xformers  # noqa: F401
     import xformers.ops as xops
 
     XFORMERS_IS_AVAILBLE = True

nunchaku/models/pulid/eva_clip/factory.py

@@ -9,7 +9,7 @@ from typing import Optional, Tuple, Union
 import torch
 
 from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
-from .model import CLIP, convert_to_custom_text_state_dict, CustomCLIP, get_cast_dtype
+from .model import CLIP, CustomCLIP, convert_to_custom_text_state_dict, get_cast_dtype
 from .pretrained import download_pretrained, get_pretrained_cfg, list_pretrained_tags_by_model
 from .transform import image_transform
 from .utils import resize_clip_pos_embed, resize_eva_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed

nunchaku/models/pulid/eva_clip/hf_model.py

@@ -11,7 +11,7 @@ from torch import TensorType
 
 try:
     import transformers
-    from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer, AutoConfig, PretrainedConfig
+    from transformers import AutoConfig, AutoModel, AutoModelForMaskedLM, AutoTokenizer, PretrainedConfig
 except ImportError:
     transformers = None
 

nunchaku/models/pulid/eva_clip/model.py

@@ -16,9 +16,9 @@ try:
     from .hf_model import HFTextEncoder
 except ImportError:
     HFTextEncoder = None
-from .modified_resnet import ModifiedResNet
 from .eva_vit_model import EVAVisionTransformer
-from .transformer import LayerNorm, QuickGELU, VisionTransformer, TextTransformer
+from .modified_resnet import ModifiedResNet
+from .transformer import LayerNorm, QuickGELU, TextTransformer, VisionTransformer
 
 try:
     from apex.normalization import FusedLayerNorm

nunchaku/models/pulid/eva_clip/modified_resnet.py

@@ -4,8 +4,6 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
-from .utils import freeze_batch_norm_2d
-
 
 class Bottleneck(nn.Module):
     expansion = 4

nunchaku/models/pulid/eva_clip/transformer.py

@@ -2,7 +2,7 @@ import logging
 import math
 import os
 from collections import OrderedDict
-from typing import Callable, Optional, Sequence
+from typing import Callable, Optional
 
 import torch
 from torch import nn
@@ -11,7 +11,7 @@ from torch.nn import functional as F
 try:
     from timm.models.layers import trunc_normal_
 except ImportError:
-    from timm.layers import trunc_normal_
+    from timm.layers import trunc_normal_  # noqa: F401
 
 from .utils import to_2tuple
 

nunchaku/models/pulid/eva_clip/utils.py

@@ -3,7 +3,6 @@ import logging
 import math
 from itertools import repeat
 
-import numpy as np
 import torch
 import torch.nn.functional as F
 from torch import nn as nn

nunchaku/models/pulid/pulid_forward.py

 # Adapted from https://github.com/ToTheBeginning/PuLID
-import torch
+import logging
 from typing import Any, Dict, Optional, Union
+
+import torch
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
-import logging
 
 logger = logging.getLogger(__name__)
 

nunchaku/models/text_encoders/linear.py

@@ -4,8 +4,8 @@
 import torch
 import torch.nn as nn
 
-from .tinychat_utils import ceil_num_groups, convert_to_tinychat_w4x16y16_linear_weight
 from ..._C.ops import gemm_awq, gemv_awq
+from .tinychat_utils import ceil_num_groups, convert_to_tinychat_w4x16y16_linear_weight
 
 __all__ = ["W4Linear"]
 

nunchaku/models/transformers/__init__.py

 from .transformer_flux import NunchakuFluxTransformer2dModel
 from .transformer_sana import NunchakuSanaTransformer2DModel
+
+__all__ = ["NunchakuFluxTransformer2dModel", "NunchakuSanaTransformer2DModel"]

nunchaku/models/transformers/transformer_flux.py

@@ -12,11 +12,12 @@ from packaging.version import Version
 from safetensors.torch import load_file
 from torch import nn
 
-from .utils import NunchakuModelLoaderMixin, pad_tensor
-from ..._C import QuantizedFluxModel, utils as cutils
+from ..._C import QuantizedFluxModel
+from ..._C import utils as cutils
 from ...lora.flux.nunchaku_converter import fuse_vectors, to_nunchaku
 from ...lora.flux.utils import is_nunchaku_format
 from ...utils import get_precision, load_state_dict_in_safetensors
+from .utils import NunchakuModelLoaderMixin, pad_tensor
 
 SVD_RANK = 32
 
@@ -77,7 +78,7 @@ class NunchakuFluxTransformerBlocks(nn.Module):
         self.id_embeddings = id_embeddings
         self.id_weight = id_weight
         self.pulid_ca_idx = 0
-        if self.id_embeddings is not None :
+        if self.id_embeddings is not None:
             self.set_residual_callback()
 
         original_dtype = hidden_states.dtype
@@ -122,13 +123,12 @@ class NunchakuFluxTransformerBlocks(nn.Module):
             rotary_emb_single,
             controlnet_block_samples,
             controlnet_single_block_samples,
-            skip_first_layer
+            skip_first_layer,
         )
 
-        if self.id_embeddings is not None :
+        if self.id_embeddings is not None:
             self.reset_residual_callback()
 
-
         hidden_states = hidden_states.to(original_dtype).to(original_device)
 
         encoder_hidden_states = hidden_states[:, :txt_tokens, ...]
@@ -191,20 +191,25 @@ class NunchakuFluxTransformerBlocks(nn.Module):
         encoder_hidden_states = encoder_hidden_states.to(original_dtype).to(original_device)
 
         return encoder_hidden_states, hidden_states
+
     def set_residual_callback(self):
         id_embeddings = self.id_embeddings
         pulid_ca = self.pulid_ca
         pulid_ca_idx = [self.pulid_ca_idx]
         id_weight = self.id_weight
+
         def callback(hidden_states):
             ip = id_weight * pulid_ca[pulid_ca_idx[0]](id_embeddings, hidden_states.to("cuda"))
             pulid_ca_idx[0] += 1
             return ip
+
         self.callback_holder = callback
         self.m.set_residual_callback(callback)
+
     def reset_residual_callback(self):
         self.callback_holder = None
         self.m.set_residual_callback(None)
+
     def __del__(self):
         self.m.reset()
 
@@ -477,10 +482,7 @@ class NunchakuFluxTransformer2dModel(FluxTransformer2DModel, NunchakuModelLoader
         if len(self._unquantized_part_loras) > 0 or len(unquantized_part_loras) > 0:
             self._unquantized_part_loras = unquantized_part_loras
 
-            self._unquantized_part_sd = {
-                k: v for k, v in self._unquantized_part_sd.items()
-                if "pulid_ca" not in k
-            }
+            self._unquantized_part_sd = {k: v for k, v in self._unquantized_part_sd.items() if "pulid_ca" not in k}
             self._update_unquantized_part_lora_params(1)
 
         quantized_part_vectors = {}

nunchaku/models/transformers/transformer_sana.py

@@ -8,9 +8,10 @@ from safetensors.torch import load_file
 from torch import nn
 from torch.nn import functional as F
 
-from .utils import NunchakuModelLoaderMixin
+from ..._C import QuantizedSanaModel
+from ..._C import utils as cutils
 from ...utils import get_precision
-from ..._C import QuantizedSanaModel, utils as cutils
+from .utils import NunchakuModelLoaderMixin
 
 SVD_RANK = 32
 
@@ -130,9 +131,11 @@ class NunchakuSanaTransformerBlocks(nn.Module):
             .to(original_dtype)
             .to(original_device)
         )
+
     def __del__(self):
         self.m.reset()
 
+
 class NunchakuSanaTransformer2DModel(SanaTransformer2DModel, NunchakuModelLoaderMixin):
     @classmethod
     @utils.validate_hf_hub_args