positive/odd enforcement (#4)

* positive/odd enforcement

* remove step constraint from sigma

README.md

@@ -15,6 +15,7 @@ The CUDA accelerated nodes can be used in real-time workflows for live video str
 - [ComfyUI-Background-Edit](#comfyui-background-edit)
 - [Install](#install)
   - [Comfy Registry](#comfy-registry)
+  - [ComfyUI-Manager](#comfyui-manager)
   - [Manual](#manual)
 - [Example Real-Time Live Video Workflows](#example-real-time-live-video-workflows)
 - [Example Image Workflows](#example-image-workflows)
@@ -22,6 +23,10 @@ The CUDA accelerated nodes can be used in real-time workflows for live video str
 
 # Install
 
+**Prererquisites**
+
+- Install [comfy-cli](https://docs.comfy.org/comfy-cli/getting-started)
+
 The recommended installation method is to use the Comfy Registry.
 
 ## Comfy Registry
@@ -32,6 +37,14 @@ These nodes can be installed via the [Comfy Registry](https://registry.comfy.org
 comfy node registry-install comfyui-background-edit
 ```
 
+## ComfyUI-Manager
+
+These nodes can be installed via ComfyUI-Manager in the UI or via the CLI:
+
+```
+comfy node install comfyui-background-edit
+```
+
 ## Manual
 
 These nodes can also be installed manually by copying them into your `custom_nodes` folder and then installing dependencies:

nodes/composite.py

@@ -1,7 +1,4 @@
 import torch
-import cvcuda
-import cv2
-import numpy as np
 
 
 class Composite:
@@ -35,41 +32,19 @@ class Composite:
                 "mismatch number of backgrounds, foregrounds and foreground masks"
             )
 
-        if mode == "cpu":
-            inverse_masks = 1.0 - foreground_masks
+        if mode != "cuda" and mode != "cpu":
+            raise Exception("invalid mode")
 
-            fgs = foregrounds * foreground_masks.unsqueeze(-1)
-            bgs = backgrounds * inverse_masks.unsqueeze(-1)
-
-            fgs_np = (fgs * 255.0).clamp(0, 255).to(dtype=torch.uint8).cpu().numpy()
-            bgs_np = (bgs * 255.0).clamp(0, 255).to(dtype=torch.uint8).cpu().numpy()
-
-            results = []
-            for fg, bg in zip(fgs_np, bgs_np):
-                result = cv2.add(fg, bg)
-                results.append(result)
-
-            return (
-                torch.from_numpy(np.stack(results, axis=0).astype(np.float32) / 255.0),
-            )
-        elif mode == "cuda":
+        if mode == "cuda":
             foregrounds = foregrounds.to("cuda")
             backgrounds = backgrounds.to("cuda")
             foreground_masks = foreground_masks.to("cuda")
 
-            fgs = cvcuda.convertto(
-                cvcuda.as_tensor(foregrounds, "NHWC"), np.uint8, scale=255
-            )
-            bgs = cvcuda.convertto(
-                cvcuda.as_tensor(backgrounds, "NHWC"), np.uint8, scale=255
-            )
-            fgmasks = cvcuda.convertto(
-                cvcuda.as_tensor(foreground_masks.unsqueeze(-1), "NHWC"),
-                np.uint8,
-                scale=255,
-            )
-            result = cvcuda.composite(fgs, bgs, fgmasks, 3)
+        inverse_masks = 1.0 - foreground_masks
 
-            return (torch.as_tensor(result.cuda()) / 255.0,)
-        else:
-            raise Exception("invalid mode")
+        fgs = foregrounds * foreground_masks.unsqueeze(-1)
+        bgs = backgrounds * inverse_masks.unsqueeze(-1)
+
+        results = fgs + bgs
+
+        return (results,)

nodes/gaussian_blur.py

@@ -1,7 +1,6 @@
 import torch
-import cvcuda
-import cv2
-import numpy as np
+
+from torchvision.transforms import v2
 
 
 class GaussianBlur:
@@ -14,8 +13,8 @@ class GaussianBlur:
         return {
             "required": {
                 "images": ("IMAGE",),
-                "kernel_size": ("INT", {"default": 61}),
-                "sigma": ("INT", {"default": 5}),
+                "kernel_size": ("INT", {"default": 61, "min": 1, "step": 2}),
+                "sigma": ("INT", {"default": 5, "min": 1, "step": 1}),
                 "mode": (["cuda", "cpu"], {"default": "cuda"}),
             }
         }
@@ -27,29 +26,17 @@ class GaussianBlur:
         sigma: int = 5,
         mode: str = "cuda",
     ):
-        if mode == "cpu":
-            images_np = (
-                (images * 255.0).clamp(0, 255).to(dtype=torch.uint8).cpu().numpy()
-            )
+        if mode != "cuda" and mode != "cpu":
+            raise Exception("invalid mode")
 
-            results = []
-            for image in images_np:
-                result = cv2.GaussianBlur(
-                    image, (kernel_size, kernel_size), sigma, sigma
-                )
-                results.append(result)
-
-            return (
-                torch.from_numpy(np.stack(results, axis=0).astype(np.float32) / 255.0),
-            )
-        elif mode == "cuda":
+        if mode == "cuda":
             images = images.to("cuda")
 
-            result = cvcuda.gaussian(
-                cvcuda.as_tensor(images, "NHWC"),
-                (kernel_size, kernel_size),
-                (sigma, sigma),
-            )
-            return (torch.as_tensor(result.cuda()),)
-        else:
-            raise Exception("invalid mode")
+        gaussian_blur = v2.GaussianBlur((kernel_size, kernel_size), (sigma, sigma))
+        # torchvision expects a BCHW tensor
+        # Convert input BHWC -> BCHW
+        blurred = gaussian_blur(images.permute(0, 3, 1, 2))
+
+        # Comfy expects a BHWC tensor
+        # Convert output BCHW -> BHWC
+        return (blurred.permute(0, 2, 3, 1),)

pyproject.toml

@@ -1,14 +1,9 @@
 [project]
 name = "comfyui-background-edit"
 description = "ComfyUI nodes for editing background of images/videos with CUDA acceleration support."
-version = "1.0.0"
+version = "1.1.0"
 license = { file = "LICENSE" }
-dependencies = [
-  "opencv-python",
-  "numpy",
-  "torch",
-  "https://github.com/CVCUDA/CV-CUDA/releases/download/v0.12.0-beta/cvcuda_cu12-0.12.0b0-cp311-cp311-linux_x86_64.whl",
-]
+dependencies = ["torch", "torchvision"]
 
 [project.urls]
 Repository = "https://github.com/yondonfu/ComfyUI-Background-Edit"

requirements.txt

@@ -1,4 +1,2 @@
-opencv-python
-numpy
 torch
-https://github.com/CVCUDA/CV-CUDA/releases/download/v0.12.0-beta/cvcuda_cu12-0.12.0b0-cp311-cp311-linux_x86_64.whl
+torchvision