Match published source formatting

encoder_blocks.py

@@ -3,19 +3,19 @@
 
 Copyright (2024) Bytedance Ltd. and/or its affiliates
 
-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 
+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 
+    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. 
+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.
 
-Reference: 
+Reference:
     https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/transformer.py
     https://github.com/baofff/U-ViT/blob/main/libs/timm.py
 """
@@ -74,11 +74,11 @@ class TiTokEncoder(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
-        self.width_size = config.dataset.preprocessing.width_size 
-        self.height_size = config.dataset.preprocessing.height_size 
+        self.width_size = config.dataset.preprocessing.width_size
+        self.height_size = config.dataset.preprocessing.height_size
         self.patch_size = config.model.vq_model.vit_enc_patch_size
-        self.grid_size_w = self.width_size // self.patch_size 
-        self.grid_size_h = self.height_size // self.patch_size 
+        self.grid_size_w = self.width_size // self.patch_size
+        self.grid_size_h = self.height_size // self.patch_size
         self.model_size = config.model.vq_model.vit_enc_model_size
         self.num_latent_tokens = config.model.vq_model.num_latent_tokens
         self.token_size = config.model.vq_model.token_size
@@ -103,11 +103,11 @@ class TiTokEncoder(nn.Module):
                 "base": 12,
                 "large": 16,
             }[self.model_size]
-        
+
         self.patch_embed = nn.Conv2d(
             in_channels=3, out_channels=self.width,
               kernel_size=self.patch_size, stride=self.patch_size,padding = (4,2), bias=True)
-        
+
         scale = self.width ** -0.5
         self.class_embedding = nn.Parameter(scale * torch.randn(1, self.width))
         self.positional_embedding = nn.Parameter(
@@ -132,7 +132,7 @@ class TiTokEncoder(nn.Module):
         # class embeddings and positional embeddings
         x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1)
         x = x + self.positional_embedding.to(x.dtype) # shape = [*, grid ** 2 + 1, width]
-        
+
 
         latent_tokens = _expand_token(latent_tokens, x.shape[0]).to(x.dtype)
         latent_tokens = latent_tokens + self.latent_token_positional_embedding.to(x.dtype)
@@ -149,7 +149,7 @@ class TiTokEncoder(nn.Module):
                 x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.transformer[i]), x,use_reentrant=False)
 
         x = x.permute(1, 0, 2)  # LND -> NLD
-        
+
         latent_tokens = x[:, 1+self.grid_size_h*self.grid_size_w:]
         latent_tokens = self.ln_post(latent_tokens)
         # fake 2D shape
@@ -173,7 +173,7 @@ class HW_encoder_2(nn.Module):
     #     # self.conv4 = nn.Conv2d(in_channels , in_channels//4 , kernel_size=3, padding=1)
 
     # def pixel_shuffle(self, x, scale_factor=0.5):
-    #     n, c, h, w = x.size() 
+    #     n, c, h, w = x.size()
     #     new_h = int(h * scale_factor)
     #     new_w = int(w * scale_factor)
     #     x = x.view(n, int(c / (scale_factor ** 2)), new_h, new_w)

quantizer.py

@@ -2,19 +2,19 @@
 
 Copyright (2024) Bytedance Ltd. and/or its affiliates
 
-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 
+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 
+    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 
+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.
 
-Reference: 
+Reference:
     https://github.com/CompVis/taming-transformers/blob/master/taming/modules/vqvae/quantize.py
     https://github.com/google-research/magvit/blob/main/videogvt/models/vqvae.py
     https://github.com/CompVis/latent-diffusion/blob/main/ldm/modules/distributions/distributions.py
@@ -92,7 +92,7 @@ class VectorQuantizer(torch.nn.Module):
                 all_unnormed_z_flattened = gather(unnormed_z_flattened).detach()
                 if all_d.shape[0] != all_unnormed_z_flattened.shape[0]:
                     raise ValueError(
-                        "all_d and all_unnormed_z_flattened have different length" + 
+                        "all_d and all_unnormed_z_flattened have different length" +
                         f"{all_d.shape}, {all_unnormed_z_flattened.shape}")
                 indices = torch.argmin(all_d, dim=0)
                 random_feat = all_unnormed_z_flattened[indices]
@@ -128,7 +128,7 @@ class VectorQuantizer(torch.nn.Module):
         if self.use_l2_norm:
             z_quantized = torch.nn.functional.normalize(z_quantized, dim=-1)
         return z_quantized
-    
+
 
 class DiagonalGaussianDistribution(object):
     @autocast(enabled=False)