Optimise shared data, generate once (#18)

* Optimise shared data, generate once

* fixup

Author: digitaljohn

nodes.py

@@ -62,42 +62,41 @@ def vignette_image(self, image: torch.Tensor, intensity: float, center_x: float,
         batch_size, height, width, _ = image.shape
         result = torch.zeros_like(image)
 
-        for b in range(batch_size):
-            tensor_image = image[b].numpy()
-
-            # Apply vignette
-            vignette_image = self.apply_vignette(tensor_image, intensity, center_x, center_y)
-
-            tensor = torch.from_numpy(vignette_image).unsqueeze(0)
-            result[b] = tensor
-
-        return (result,)
-
-    def apply_vignette(self, image, vignette_strength, center_x_ratio, center_y_ratio):
-        if vignette_strength == 0:
+        if intensity == 0:
             return image
 
-        height, width, _ = image.shape
-
+        # Generate the vignette for each image in the batch
         # Create linear space but centered around the provided center point ratios
         x = np.linspace(-1, 1, width)
         y = np.linspace(-1, 1, height)
-        X, Y = np.meshgrid(x - (2 * center_x_ratio - 1), y - (2 * center_y_ratio - 1))
+        X, Y = np.meshgrid(x - (2 * center_x - 1), y - (2 * center_y - 1))
 
         # Calculate distances to the furthest corner
         distances_to_corners = [
-            np.sqrt((0 - center_x_ratio) ** 2 + (0 - center_y_ratio) ** 2),
-            np.sqrt((1 - center_x_ratio) ** 2 + (0 - center_y_ratio) ** 2),
-            np.sqrt((0 - center_x_ratio) ** 2 + (1 - center_y_ratio) ** 2),
-            np.sqrt((1 - center_x_ratio) ** 2 + (1 - center_y_ratio) ** 2)
+            np.sqrt((0 - center_x) ** 2 + (0 - center_y) ** 2),
+            np.sqrt((1 - center_x) ** 2 + (0 - center_y) ** 2),
+            np.sqrt((0 - center_x) ** 2 + (1 - center_y) ** 2),
+            np.sqrt((1 - center_x) ** 2 + (1 - center_y) ** 2)
         ]
         max_distance_to_corner = np.max(distances_to_corners)
 
         radius = np.sqrt(X ** 2 + Y ** 2)
         radius = radius / (max_distance_to_corner * np.sqrt(2))  # Normalize radius
-        opacity = np.clip(vignette_strength, 0, 1)
+        opacity = np.clip(intensity, 0, 1)
         vignette = 1 - radius * opacity
 
+        for b in range(batch_size):
+            tensor_image = image[b].numpy()
+
+            # Apply vignette
+            vignette_image = self.apply_vignette(tensor_image, vignette)
+
+            tensor = torch.from_numpy(vignette_image).unsqueeze(0)
+            result[b] = tensor
+
+        return (result,)
+
+    def apply_vignette(self, image, vignette):
         # If image needs to be normalized (0-1 range)
         needs_normalization = image.max() > 1
         if needs_normalization:
@@ -206,6 +205,7 @@ def apply_filmgrain(self, image, gray_scale, grain_type, grain_sat, grain_power,
         out_image = filmgrainer.process(image, scale, src_gamma, 
             grain_power, shadows, highs, grain_type, 
             grain_sat, gray_scale, sharpen, seed)
+        
         return out_image
 
 
@@ -263,37 +263,37 @@ def radialblur_image(self, image: torch.Tensor, blur_strength: float, center_x:
         batch_size, height, width, _ = image.shape
         result = torch.zeros_like(image)
 
+        # Generate the vignette for each image in the batch
+        c_x, c_y = int(width * center_x), int(height * center_y)
+
+        # Calculate distances to all corners from the center
+        distances_to_corners = [
+            np.sqrt((c_x - 0)**2 + (c_y - 0)**2),
+            np.sqrt((c_x - width)**2 + (c_y - 0)**2),
+            np.sqrt((c_x - 0)**2 + (c_y - height)**2),
+            np.sqrt((c_x - width)**2 + (c_y - height)**2)
+        ]
+        max_distance_to_corner = max(distances_to_corners)
+
+        # Create and adjust radial mask
+        X, Y = np.meshgrid(np.arange(width) - c_x, np.arange(height) - c_y)
+        radial_mask = np.sqrt(X**2 + Y**2) / max_distance_to_corner
+
         for b in range(batch_size):
             tensor_image = image[b].numpy()
 
             # Apply blur
-            blur_image = self.apply_radialblur(tensor_image, blur_strength, center_x, center_y, focus_spread, steps)
+            blur_image = self.apply_radialblur(tensor_image, blur_strength, radial_mask, focus_spread, steps)
 
             tensor = torch.from_numpy(blur_image).unsqueeze(0)
             result[b] = tensor
 
         return (result,)
 
-    def apply_radialblur(self, image, blur_strength, center_x_ratio, center_y_ratio, focus_spread, steps):
+    def apply_radialblur(self, image, blur_strength, radial_mask, focus_spread, steps):
         needs_normalization = image.max() > 1
         if needs_normalization:
             image = image.astype(np.float32) / 255
-        
-        height, width = image.shape[:2]
-        center_x, center_y = int(width * center_x_ratio), int(height * center_y_ratio)
-
-        # Calculate distances to all corners from the center
-        distances_to_corners = [
-            np.sqrt((center_x - 0)**2 + (center_y - 0)**2),
-            np.sqrt((center_x - width)**2 + (center_y - 0)**2),
-            np.sqrt((center_x - 0)**2 + (center_y - height)**2),
-            np.sqrt((center_x - width)**2 + (center_y - height)**2)
-        ]
-        max_distance_to_corner = max(distances_to_corners)
-
-        # Create and adjust radial mask
-        X, Y = np.meshgrid(np.arange(width) - center_x, np.arange(height) - center_y)
-        radial_mask = np.sqrt(X**2 + Y**2) / max_distance_to_corner
 
         blurred_images = processing_utils.generate_blurred_images(image, blur_strength, steps, focus_spread)
         final_image = processing_utils.apply_blurred_images(image, blurred_images, radial_mask)