From cf49422eec8ef39ae50bd0876b98f4ea0fef7a5c Mon Sep 17 00:00:00 2001
From: l-comm <l-comm>
Date: Sat, 11 Jan 2025 18:14:06 +0800
Subject: [PATCH] Give it your best shot on comfy

---
 find_watermarks.py | 10 +++++-----
 nodes.py           | 12 ++++++++----
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/find_watermarks.py b/find_watermarks.py
index 67af1f7..178fbfe 100644
--- a/find_watermarks.py
+++ b/find_watermarks.py
@@ -362,7 +362,7 @@ def rgb_to_grayscale_torch(image: torch.Tensor) -> torch.Tensor:
     return gray
 
 
-def find_watermark_frame_tensor(original_color_torch: torch.Tensor):
+def find_watermark_frame_tensor(original_color_torch: torch.Tensor, result_cutoff=0.9):
     """
     Replaces find_watermark_frame_cv() but uses only PyTorch on GPU.
     original_color_torch: shape [3, H, W], float32, device='cuda'
@@ -419,8 +419,8 @@ def find_watermark_frame_tensor(original_color_torch: torch.Tensor):
     # logger.info(f"original_color_torch 2 {original_color_torch.shape}")
     final_corrected_image_torch = remove_watermark_batch(original_color_torch, best_x, best_y)
     final_result = predictor.predict_tensor(final_corrected_image_torch[0])
-
-    if final_result > 0.9:
+    print(f"final {final_result}")
+    if final_result > result_cutoff:
         return None
 
     return best_x, best_y, float(final_result)
@@ -450,7 +450,7 @@ def min_max_loc_torch(matrix: torch.Tensor):
             (min_x, min_y),
             (max_x, max_y))
 
-def find_watermark_tensor(pixel_values, frame_range=range(0, 24, 5)):
+def find_watermark_tensor(pixel_values, frame_range=range(0, 24, 5), result_cutoff=0.9):
     """
     Replaces find_watermark_tensor but does all processing in PyTorch on GPU.
     
@@ -471,7 +471,7 @@ def find_watermark_tensor(pixel_values, frame_range=range(0, 24, 5)):
         # We directly use the PyTorch tensor
         frame_torch = pixel_values[frame_idx]  # shape [3, H, W]
 
-        res = find_watermark_frame_tensor(frame_torch)
+        res = find_watermark_frame_tensor(frame_torch, result_cutoff)
         if not res:
             logger.error(f"{frame_idx} Failed to find")
             continue
diff --git a/nodes.py b/nodes.py
index 4845982..6dfebe6 100644
--- a/nodes.py
+++ b/nodes.py
@@ -39,7 +39,7 @@ class FindWatermarkNode:
             Tuple[int, int, str]: x-coordinate, y-coordinate, and detection score.
         """
         # Call your existing watermark detection function
-        result = find_watermarks.find_watermark_tensor(images.permute(0, 3, 1, 2))
+        result = find_watermarks.find_watermark_tensor(images.permute(0, 3, 1, 2), result_cutoff=1.0)
 
         if not result:
             # If no watermark found, return default values
@@ -67,6 +67,7 @@ class RemoveWatermarkNode:
             "optional": {
                 "x": ("INT", {"forceInput":True}),
                 "y": ("INT", {"forceInput":True}),
+                "score": ("STRING", {"forceInput":True, "default": ""}),
                 "nudge_x": ("INT", {"default": 0, "min": -999999, "max": 999999, "step": 1}),
                 "nudge_y": ("INT", {"default": 0, "min": -999999, "max": 999999, "step": 1}),
             }
@@ -77,12 +78,13 @@ class RemoveWatermarkNode:
     FUNCTION = "remove_watermark"
     CATEGORY = "Video/Watermarks"
 
-    def remove_watermark(self, images, x=0, y=0, nudge_x=0, nudge_y=0):
+    def remove_watermark(self, images, x=0, y=0, score="", nudge_x=0, nudge_y=0):
         """
         Args:
             images (list of PIL.Image): Batch of input images.
             x (int): x-coordinate of the watermark.
             y (int): y-coordinate of the watermark.
+            score (string): lets the detector tell us if it failed
             nudge_x (int): Additional x nudge value.
             nudge_y (int): Additional y nudge value.
 
@@ -93,6 +95,8 @@ class RemoveWatermarkNode:
         y += nudge_y
 
         # Call your existing watermark removal function
-        corrected_frames = find_watermarks.remove_watermark_batch(images.permute(0, 3, 1, 2), x, y)
+        if score != "Failed":
+            corrected_frames = find_watermarks.remove_watermark_batch(images.permute(0, 3, 1, 2), x, y)
+            images = corrected_frames.permute(0, 2, 3, 1)
 
-        return (corrected_frames.permute(0, 2, 3, 1),)
\ No newline at end of file
+        return (images,)
\ No newline at end of file