Remove wandb

tim/utils/misc_utils.py

@@ -1,7 +1,6 @@
 import functools
 import importlib
 import os
-import wandb
 import fsspec
 import numpy as np
 import torch
@@ -13,12 +12,13 @@ from PIL import Image, ImageDraw, ImageFont
 from safetensors.torch import load_file
 from tqdm import tqdm
 
+
 def create_npz_from_sample_folder(sample_dir, num=50_000):
     """
     Builds a single .npz file from a folder of .png samples.
     """
     samples = []
-    imgs = sorted(os.listdir(sample_dir), key=lambda x: int(x.split('.')[0]))
+    imgs = sorted(os.listdir(sample_dir), key=lambda x: int(x.split(".")[0]))
     print(len(imgs))
     assert len(imgs) >= num
     for i in tqdm(range(num), desc="Building .npz file from samples"):
@@ -32,14 +32,13 @@ def create_npz_from_sample_folder(sample_dir, num=50_000):
     print(f"Saved .npz file to {npz_path} [shape={samples.shape}].")
     return npz_path
 
-def init_from_ckpt(
-    model, checkpoint_dir, ignore_keys=None, verbose=False
-) -> None: 
+
+def init_from_ckpt(model, checkpoint_dir, ignore_keys=None, verbose=False) -> None:
     if checkpoint_dir.endswith(".safetensors"):
-        model_state_dict=load_file(checkpoint_dir, device='cpu')
+        model_state_dict = load_file(checkpoint_dir, device="cpu")
     else:
-        model_state_dict=torch.load(checkpoint_dir,  map_location="cpu")
-    model_new_ckpt=dict()
+        model_state_dict = torch.load(checkpoint_dir, map_location="cpu")
+    model_new_ckpt = dict()
     for i in model_state_dict.keys():
         model_new_ckpt[i] = model_state_dict[i]
     keys = list(model_new_ckpt.keys())
@@ -63,14 +62,14 @@ def init_from_ckpt(
 
 
 def get_dtype(str_dtype):
-    if str_dtype == 'fp16':
+    if str_dtype == "fp16":
         return torch.float16
-    elif str_dtype == 'bf16':
+    elif str_dtype == "bf16":
         return torch.bfloat16
     else:
         return torch.float32
-    
-    
+
+
 def disabled_train(self, mode=True):
     """Overwrite model.train with this function to make sure train/eval mode
     does not change anymore."""
@@ -221,12 +220,12 @@ def mean_flat(tensor):
 def count_params(model, verbose=False):
     total_params = sum(p.numel() for p in model.parameters())
     if verbose:
-        print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
+        print(f"{model.__class__.__name__} has {total_params * 1.0e-6:.2f} M params.")
     return total_params
 
 
 def instantiate_from_config(config):
-    if not "target" in config:
+    if "target" not in config:
         if config == "__is_first_stage__":
             return None
         elif config == "__is_unconditional__":
@@ -295,7 +294,8 @@ def load_model_from_config(config, ckpt, verbose=True, freeze=True):
 def format_number(num):
     num = float(num)
     num /= 1000.0
-    return '{:.0f}{}'.format(num, 'k')
+    return "{:.0f}{}".format(num, "k")
+
 
 def get_num_params(model: torch.nn.ModuleList) -> int:
     num_params = sum(p.numel() for p in model.parameters())
@@ -319,13 +319,14 @@ def get_num_flop_per_token(num_params, model_config, seq_len) -> int:
 
     return flop_per_token
 
+
 def get_num_flop_per_sequence_encoder_only(num_params, model_config, seq_len) -> int:
     l, h, q = (
         model_config.n_layers,
         model_config.n_heads,
         model_config.dim // model_config.n_heads,
     )
-    
+
     # 1. 每个自注意力层有2个矩阵乘法在前向传播,4个在反向传播 (6)
     # 2. 每个矩阵乘法执行1次乘法和1次加法 (*2)
     # 3. 双向注意力需要考虑所有token对,所以是t^2而不是t
@@ -351,6 +352,7 @@ def get_peak_flops(device_name: str) -> int:
     else:  # for other GPU types, assume A100
         return 312e12
 
+
 @dataclass(frozen=True)
 class Color:
     black = "\033[30m"
@@ -374,4 +376,4 @@ class NoColor:
     magenta = ""
     cyan = ""
     white = ""
-    reset = ""
+    reset = ""