Lint and format code properly

Signed-off-by: Neal Pan <[email protected]>

Author: LostnEkko

CHANGELOG.md

@@ -24,7 +24,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Updated CorrDiff training code to support multiple patch iterations to amortize regression cost and usage of `torch.compile`
 - Refactored `physicsnemo/models/diffusion/layers.py` to optimize data type casting workflow, avoiding unnecessary casting under autocast mode
 - Refactored Conv2d to enable fusion of conv2d with bias addition
-- Refactored GroupNorm, UNetBlock, SongUNet, SongUNetPosEmbd to support usage of Apex GroupNorm, fusion of activation with GroupNorm, and AMP workflow.
+- Refactored GroupNorm, UNetBlock, SongUNet, SongUNetPosEmbd to support usage of
+Apex GroupNorm, fusion of activation with GroupNorm, and AMP workflow.
 - Updated SongUNetPosEmbd to avoid unnecessary HtoD Memcpy of `pos_embd`
 - Updated `from_checkpoint` to accommodate usage of Apex GroupNorm
 - Refactored CorrDiff NVTX annotation workflow to be configurable

examples/generative/corrdiff/train.py

@@ -1057,7 +1057,15 @@ def __init__(
         self.hr_mean_conditioning = hr_mean_conditioning
         self.y_mean = None
 
-    def __call__(self, net, img_clean, img_lr, patch_num_per_iter=-1, labels=None, augment_pipe=None):
+    def __call__(
+        self,
+        net,
+        img_clean,
+        img_lr,
+        patch_num_per_iter=-1,
+        labels=None,
+        augment_pipe=None,
+    ):
         """
         Calculate and return the loss for denoising score matching.
 
@@ -1085,27 +1093,29 @@ def __call__(self, net, img_clean, img_lr, patch_num_per_iter=-1, labels=None, a
             A tensor representing the loss calculated based on the network's
             predictions.
         """
-        
-        self.patch_num = patch_num_per_iter if patch_num_per_iter != -1 else self.patch_num
-        
+
+        self.patch_num = (
+            patch_num_per_iter if patch_num_per_iter != -1 else self.patch_num
+        )
+
         rnd_normal = torch.randn([img_clean.shape[0], 1, 1, 1], device=img_clean.device)
         sigma = (rnd_normal * self.P_std + self.P_mean).exp()
         weight = (sigma**2 + self.sigma_data**2) / (sigma * self.sigma_data) ** 2
-        
+
         # augment for conditional generaiton
         img_tot = torch.cat((img_clean, img_lr), dim=1)
         y_tot, augment_labels = (
             augment_pipe(img_tot) if augment_pipe is not None else (img_tot, None)
         )
         y = y_tot[:, : img_clean.shape[1], :, :]
         y_lr = y_tot[:, img_clean.shape[1] :, :, :]
-        
+
         y_lr_res = y_lr.to(memory_format=torch.channels_last)
 
         # global index
         b = y.shape[0]
-        Nx = torch.arange(self.img_shape_x,device=img_clean.device).int()
-        Ny = torch.arange(self.img_shape_y,device=img_clean.device).int()
+        Nx = torch.arange(self.img_shape_x, device=img_clean.device).int()
+        Ny = torch.arange(self.img_shape_y, device=img_clean.device).int()
         grid = torch.stack(torch.meshgrid(Ny, Nx, indexing="ij"), dim=0)[
             None,
         ].expand(b, -1, -1, -1)
@@ -1122,25 +1132,24 @@ def __call__(self, net, img_clean, img_lr, patch_num_per_iter=-1, labels=None, a
             self.y_mean = y_mean
         # else:
         #     y_mean = self.y_mean
-            
+
         y = y - self.y_mean
 
         if self.hr_mean_conditioning:
             y_lr = torch.cat((self.y_mean, y_lr), dim=1)
-        
-        
+
         global_index = None
         # patchified training
         # conditioning: cat(y_mean, y_lr, input_interp, pos_embd), 4+12+100+4
-        
+
         if (
             self.img_shape_x != self.patch_shape_x
             or self.img_shape_y != self.patch_shape_y
         ):
-            
+
             c_in = y_lr.shape[1]
             c_out = y.shape[1]
-            
+
             rnd_normal = torch.randn(
                 [img_clean.shape[0] * self.patch_num, 1, 1, 1], device=img_clean.device
             )
@@ -1179,7 +1188,7 @@ def __call__(self, net, img_clean, img_lr, patch_num_per_iter=-1, labels=None, a
                 dtype=torch.int,
                 device=img_clean.device,
             )
-            
+
             for i in range(self.patch_num):
                 rnd_x = random.randint(0, self.img_shape_x - self.patch_shape_x)
                 rnd_y = random.randint(0, self.img_shape_y - self.patch_shape_y)
@@ -1207,10 +1216,10 @@ def __call__(self, net, img_clean, img_lr, patch_num_per_iter=-1, labels=None, a
                     ),
                     1,
                 )
-                
+
             y = y_new
             y_lr = y_lr_new
-        
+
         latent = y + torch.randn_like(y) * sigma
         D_yn = net(
             latent,
@@ -1220,11 +1229,12 @@ def __call__(self, net, img_clean, img_lr, patch_num_per_iter=-1, labels=None, a
             global_index=global_index,
             augment_labels=augment_labels,
         )
-        
+
         loss = weight * ((D_yn - y) ** 2)
 
         return loss
 
+
 class VELoss_dfsr:
     """
     Loss function for dfsr model, modified from class VELoss.

physicsnemo/metrics/diffusion/loss.py

@@ -23,7 +23,6 @@
     Linear,
     PositionalEmbedding,
     UNetBlock,
-
 )
 from .song_unet import SongUNet, SongUNetPosEmbd, SongUNetPosLtEmbd
 from .dhariwal_unet import DhariwalUNet