Upload folder using huggingface_hub

.gitattributes

@@ -497,3 +497,119 @@ results/reflow_ep99_eval_8step/8_rain_flow.png filter=lfs diff=lfs merge=lfs -te
 results/reflow_ep99_eval_8step/8_rain_input.png filter=lfs diff=lfs merge=lfs -text
 results/reflow_ep99_eval_8step/9_rain_flow.png filter=lfs diff=lfs merge=lfs -text
 results/reflow_ep99_eval_8step/9_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/0_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/0_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/10_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/10_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/11_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/11_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/12_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/12_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/13_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/13_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/14_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/14_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/15_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/15_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/16_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/16_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/17_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/17_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/18_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/18_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/19_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/19_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/1_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/1_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/20_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/20_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/21_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/21_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/22_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/22_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/23_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/23_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/24_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/24_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/25_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/25_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/26_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/26_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/27_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/27_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/28_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/28_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/29_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/29_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/2_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/2_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/30_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/30_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/31_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/31_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/32_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/32_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/33_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/33_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/34_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/34_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/35_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/35_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/36_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/36_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/37_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/37_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/38_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/38_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/39_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/39_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/3_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/3_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/40_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/40_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/41_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/41_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/42_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/42_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/43_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/43_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/44_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/44_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/45_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/45_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/46_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/46_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/47_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/47_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/48_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/48_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/49_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/49_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/4_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/4_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/50_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/50_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/51_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/51_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/52_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/52_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/53_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/53_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/54_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/54_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/55_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/55_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/56_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/56_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/57_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/57_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/5_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/5_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/6_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/6_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/7_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/7_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/8_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/8_rain_input.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/9_rain_flow.png filter=lfs diff=lfs merge=lfs -text
+results/reflow_ep99_eval_1step_fixed/9_rain_input.png filter=lfs diff=lfs merge=lfs -text

diff2flow.py

@@ -58,7 +58,7 @@ def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_time
 class VPDiffusionFlow:
     def __init__(self, args, config):
         self.args = args
-        self.flow_mode = getattr(args, 'flow_mode', 'vp')
+        self.flow_mode = getattr(args, "flow_mode", "vp")
         self.config = config
         self.device = config.device
 
@@ -134,9 +134,9 @@ class VPDiffusionFlow:
         # Calculate alpha_bar analytically for linear beta schedule
         scalar_t = t.item() if isinstance(t, torch.Tensor) else t
         scalar_t = max(0.0, min(1.0, scalar_t))
-        
+
         N = self.num_timesteps
-        
+
         # Integral of N * (b0 + (b1-b0)*s) ds from 0 to t
         # = N * [ b0*t + 0.5*(b1-b0)*t^2 ]
         b0 = self.beta_start
@@ -151,6 +151,13 @@ class VPDiffusionFlow:
         w_list = [i for i in range(0, w - output_size + 1, r)]
         return h_list, w_list
 
+    def get_blending_window(self, patch_size):
+        # Hanning window (cosine-based, smooth goes to 0 at edges)
+        # Using periodic=False (symmetric window)
+        w = torch.hann_window(patch_size, periodic=False, device=self.device)
+        w2d = w.unsqueeze(0) * w.unsqueeze(1)
+        return w2d.view(1, 1, patch_size, patch_size)
+
     def get_velocity(self, x, t, x_cond, patch_size=None, r_stride=16):
         # If no patching needed (x fits in patch_size or patch_size None), do normal
         if patch_size is None or (
@@ -163,42 +170,56 @@ class VPDiffusionFlow:
         t_idx = min(int(t * N), N - 1)
         t_input_scalar = t_idx
 
-        # Grid setup
-        h_list, w_list = self.overlapping_grid_indices(x, patch_size, r=r_stride)
+        # --- Padding to handle edges ---
+        # Pad by patch_size // 2 to ensure original edges are covered by window center
+        pad_size = patch_size // 2
+        x_padded = torch.nn.functional.pad(
+            x, (pad_size, pad_size, pad_size, pad_size), mode="reflect"
+        )
+        x_cond_padded = torch.nn.functional.pad(
+            x_cond, (pad_size, pad_size, pad_size, pad_size), mode="reflect"
+        )
+
+        # Grid setup on PADDED image
+        h_list, w_list = self.overlapping_grid_indices(x_padded, patch_size, r=r_stride)
         corners = [(i, j) for i in h_list for j in w_list]
 
+        # Use Weighted Averaging (Hanning Window) to reduce grid artifacts
+        window = self.get_blending_window(patch_size)
+
         # Mask for overlap averaging
-        # Cache this if t doesn't change grid? Grid depends on image size, which is constant during integration.
-        # But for now recompute is safer/easier.
-        x_grid_mask = torch.zeros_like(x, device=self.device)
+        x_grid_mask = torch.zeros_like(x_padded, device=self.device)
         for hi, wi in corners:
-            x_grid_mask[:, :, hi : hi + patch_size, wi : wi + patch_size] += 1
+            x_grid_mask[:, :, hi : hi + patch_size, wi : wi + patch_size] += window
 
         # Accumulate output (epsilon or velocity)
-        output_accum = torch.zeros_like(x, device=self.device)
+        output_accum = torch.zeros_like(x_padded, device=self.device)
 
         # Process in batches
         batch_size = 64  # From restoration.py logic or config
 
         # Prepare params if VP
-        if self.flow_mode == 'vp':
+        if self.flow_mode == "vp":
             beta_discrete = self.get_beta_t(t)
             beta_cont = beta_discrete * N
             ab = self.alphas_cumprod[t_idx]
-        
+
         # Loop over patches
         # NOTE: drift depends on x (noisy) and x_cond (clean/cond).
         for i in range(0, len(corners), batch_size):
             batch_corners = corners[i : i + batch_size]
 
-            # Crop batch
+            # Crop batch from PADDED input
             x_batch = torch.cat(
-                [crop(x, hi, wi, patch_size, patch_size) for (hi, wi) in batch_corners],
+                [
+                    crop(x_padded, hi, wi, patch_size, patch_size)
+                    for (hi, wi) in batch_corners
+                ],
                 dim=0,
             )
             cond_batch = torch.cat(
                 [
-                    crop(x_cond, hi, wi, patch_size, patch_size)
+                    crop(x_cond_padded, hi, wi, patch_size, patch_size)
                     for (hi, wi) in batch_corners
                 ],
                 dim=0,
@@ -208,19 +229,33 @@ class VPDiffusionFlow:
             )
 
             with torch.no_grad():
-                model_output = self.model(torch.cat([cond_batch, x_batch], dim=1), t_batch)
+                model_output = self.model(
+                    torch.cat([cond_batch, x_batch], dim=1), t_batch
+                )
+
+            # Scatter back with window weighting
+            # model_output: [B, C, P, P]
+            weighted_output = model_output * window
 
-            # Scatter back
             for idx, (hi, wi) in enumerate(batch_corners):
                 output_accum[0, :, hi : hi + patch_size, wi : wi + patch_size] += (
-                    model_output[idx]
+                    weighted_output[idx]
                 )
 
         # Average
-        model_output_full = torch.div(output_accum, x_grid_mask)
+        # Add epsilon to mask to avoid division by zero
+        model_output_full = torch.div(output_accum, x_grid_mask + 1e-8)
+
+        # --- Crop back to original size ---
+        # x_padded was padded by pad_size on all sides.
+        # Original is at pad_size : -pad_size
+        if pad_size > 0:
+            model_output_full = model_output_full[
+                :, :, pad_size:-pad_size, pad_size:-pad_size
+            ]
 
         # Compute v
-        if self.flow_mode == 'reflow':
+        if self.flow_mode == "reflow":
             # In Reflow, model output is velocity
             v = model_output_full
         else:
@@ -229,7 +264,7 @@ class VPDiffusionFlow:
             coeff1 = -0.5 * beta_cont
             coeff2 = 0.5 * beta_cont / torch.sqrt(1 - ab)
             v = coeff1 * x + coeff2 * epsilon
-            
+
         return v
 
     def _get_velocity_single(self, x, t, x_cond):
@@ -245,7 +280,7 @@ class VPDiffusionFlow:
         with torch.no_grad():
             model_output = self.model(torch.cat([x_cond, x], dim=1), t_input)
 
-        if self.flow_mode == 'reflow':
+        if self.flow_mode == "reflow":
             return model_output
         else:
             epsilon = model_output
@@ -315,7 +350,11 @@ def main():
         "--rtol", type=float, default=1e-4, help="Relative tolerance for ODE solver"
     )
     parser.add_argument(
-        "--flow_mode", type=str, default="vp", choices=["vp", "reflow"], help="Flow mode: vp (default) or reflow"
+        "--flow_mode",
+        type=str,
+        default="vp",
+        choices=["vp", "reflow"],
+        help="Flow mode: vp (default) or reflow",
     )
     args = parser.parse_args()
 

generate_reflow_pairs.py

@@ -10,7 +10,17 @@ from diff2flow import VPDiffusionFlow, dict2namespace
 import datasets
 from tqdm import tqdm
 
-def ode_inverse_solve(flow_model, x_data, x_cond, steps=100, method="dopri5", patch_size=64, atol=1e-5, rtol=1e-5):
+
+def ode_inverse_solve(
+    flow_model,
+    x_data,
+    x_cond,
+    steps=100,
+    method="dopri5",
+    patch_size=64,
+    atol=1e-5,
+    rtol=1e-5,
+):
     """
     Solves the ODE from t=0 (data) to t=1 (noise).
     Returns x_1 (noise latent).
@@ -18,7 +28,7 @@ def ode_inverse_solve(flow_model, x_data, x_cond, steps=100, method="dopri5", pa
     # Define the drift function wrapper for torchdiffeq
     # For inversion, we integrate from 0 to 1.
     # The drift v(x, t) is defined for t in [0, 1].
-    
+
     def drift_func(t, x):
         # flow_model.get_velocity expects t in [0, 1]
         # When using torchdiffeq, t will be traversing 0->1.
@@ -26,7 +36,7 @@ def ode_inverse_solve(flow_model, x_data, x_cond, steps=100, method="dopri5", pa
 
     # Time points from 0 to 1
     t_eval = torch.linspace(0.0, 1.0, steps + 1, device=x_data.device)
-    
+
     # Solve
     out = torchdiffeq.odeint(
         drift_func, x_data, t_eval, method=method, atol=atol, rtol=rtol
@@ -34,6 +44,7 @@ def ode_inverse_solve(flow_model, x_data, x_cond, steps=100, method="dopri5", pa
     # Return final state at t=1
     return out[-1]
 
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", type=str, required=True)
@@ -47,40 +58,48 @@ def main():
     parser.add_argument("--method", type=str, default="dopri5")
     parser.add_argument("--atol", type=float, default=1e-5)
     parser.add_argument("--rtol", type=float, default=1e-5)
-    parser.add_argument("--max_images", type=int, default=None, help="Max images to generate (for testing)")
+    parser.add_argument(
+        "--max_images",
+        type=int,
+        default=None,
+        help="Max images to generate (for testing)",
+    )
     args = parser.parse_args()
 
     # Load Config
     with open(os.path.join("configs", args.config), "r") as f:
         config_dict = yaml.safe_load(f)
     config = dict2namespace(config_dict)
-    
+
     if args.data_dir:
         config.data.data_dir = args.data_dir
     if args.dataset:
         config.data.dataset = args.dataset
-        
+
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     config.device = device
-    
+
     # Reproducibility
     torch.manual_seed(args.seed)
     np.random.seed(args.seed)
-    
+
     # Load Model
     print("Initializing VPDiffusionFlow...")
     flow = VPDiffusionFlow(args, config)
     flow.load_ckpt(args.resume)
-    
+
     os.makedirs(args.output_dir, exist_ok=True)
-    
+
     # Load Dataset
     print(f"Loading dataset {config.data.dataset}...")
     DATASET = datasets.__dict__[config.data.dataset](config)
-    
+
     # We use the TRAINING set to generate pairs for training the reflow model
-    train_loader, _ = DATASET.get_loaders(parse_patches=False, validation=config.data.dataset if args.dataset else 'raindrop')
-    
+    train_loader, _ = DATASET.get_loaders(
+        parse_patches=False,
+        validation=config.data.dataset if args.dataset else "raindrop",
+    )
+
     # We want to iterate over training data. Note: get_loaders usually returns (train_loader, val_loader).
     # RainDrop.get_loaders returns (train_loader, val_loader).
     # train_loader usually parses patches = True for original training.
@@ -90,72 +109,75 @@ def main():
     # The original training was likely on patches (RainDropDataset uses patch_size).
     # For Reflow, we should probably train on PATCHES to match the original training distribution and efficiency.
     # So let's use parse_patches=True for the loader to match training setup.
-    
+
     # However, to use `ode_inverse_solve`, we need to follow the flow.
     # If we use patches, we can solve ODE for each patch independently.
     # This is consistent.
-    
+
     # Re-get loaders with parse_patches=True to get training patches
     train_loader, _ = DATASET.get_loaders(parse_patches=True)
-    
+
     print(f"Starting generation of reflow pairs...")
-    
+
     count = 0
-    
+
     # Iterate through training patches
-    for i, (x_batch, img_id) in enumerate(tqdm(train_loader, desc="Generating Reflow Pairs")):
+    for i, (x_batch, img_id) in enumerate(
+        tqdm(train_loader, desc="Generating Reflow Pairs")
+    ):
         # x_batch: [B, N, 6, H, W] if parse_patches=True
         # Flatten B and N to process all patches
         if x_batch.ndim == 5:
             x_batch = x_batch.flatten(start_dim=0, end_dim=1)
-        
-        input_img = x_batch[:, :3, :, :].to(device) # Input (Rainy)
-        gt_img = x_batch[:, 3:, :, :].to(device)    # GT (Clean)
-        
+
+        input_img = x_batch[:, :3, :, :].to(device)  # Input (Rainy)
+        gt_img = x_batch[:, 3:, :, :].to(device)  # GT (Clean)
+
         # Transform data to [-1, 1]
         x_cond = utils.sampling.data_transform(input_img)
         x_data = utils.sampling.data_transform(gt_img)
-        
+
         # Run ODE Inversion: x_data (t=0) -> x_noise (t=1)
         # Note: patch_size argument in ode_inverse_solve usually used for stitching.
         # Here x_data IS a patch (e.g. 64x64 or config size).
         # So we can pass patch_size=None or just let it handle it.
         # Our VPDiffusionFlow.get_velocity handles patching if x > patch_size.
         # Here x is likely small.
-        
+
         with torch.no_grad():
             x_noise = ode_inverse_solve(
-                flow, 
-                x_data, 
-                x_cond, 
-                steps=args.steps, 
-                method=args.method, 
+                flow,
+                x_data,
+                x_cond,
+                steps=args.steps,
+                method=args.method,
                 patch_size=args.patch_size,
-                atol=args.atol, 
-                rtol=args.rtol
+                atol=args.atol,
+                rtol=args.rtol,
             )
-        
+
         # Save pair (x_noise, x_cond, x_data)
         # x_noise is the 'target' input for the reflow model (at t=1)
         # x_data is the 'target' output (at t=0)
         # x_cond is the condition
-        
+
         # We save this batch
         batch_data = {
             "x_noise": x_noise.cpu(),
             "x_data": x_data.cpu(),
-            "x_cond": x_cond.cpu()
+            "x_cond": x_cond.cpu(),
         }
-        
+
         save_path = os.path.join(args.output_dir, f"batch_{i}.pth")
         torch.save(batch_data, save_path)
-        
+
         print(f"Saved batch {i} to {save_path}")
-        
+
         count += input_img.shape[0]
         if args.max_images and count >= args.max_images:
             print(f"Reached max images {args.max_images}")
             break
 
+
 if __name__ == "__main__":
     main()