From 272fb5a978579eb4d51119b8553f6b6a327a747e Mon Sep 17 00:00:00 2001
From: gorillaframeai <149839341+gorillaframeai@users.noreply.github.com>
Date: Thu, 14 Nov 2024 20:41:19 +0300
Subject: [PATCH] Initial commit

---
 .gitattributes |  2 ++
 LICENSE        | 21 +++++++++++
 __init__.py    |  9 +++++
 gfrbmg2.py     | 98 ++++++++++++++++++++++++++++++++++++++++++++++++++
 readme.md      |  2 ++
 5 files changed, 132 insertions(+)
 create mode 100644 .gitattributes
 create mode 100644 LICENSE
 create mode 100644 __init__.py
 create mode 100644 gfrbmg2.py
 create mode 100644 readme.md

diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..dfe0770
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..980d7bf
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 gorillaframeai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..a7d6d2a
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,9 @@
+from .gfrbmg2 import GFrbmg2
+
+NODE_CLASS_MAPPINGS = {
+    "GFrbmg2": GFrbmg2
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "GFrbmg2": "🐵 GF Remove Background 2.0"
+}
\ No newline at end of file
diff --git a/gfrbmg2.py b/gfrbmg2.py
new file mode 100644
index 0000000..d73598f
--- /dev/null
+++ b/gfrbmg2.py
@@ -0,0 +1,98 @@
+import torch, os
+import torch.nn.functional as F
+import folder_paths
+from PIL import Image
+from transformers import AutoModelForImageSegmentation
+from torchvision import transforms
+from torchvision.transforms.functional import normalize
+import numpy as np
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Добавляем путь к моделям ComfyUI
+folder_paths.add_model_folder_path("rmbg_models", os.path.join(folder_paths.models_dir, "RMBG-2.0"))
+
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+def resize_image(image):
+    image = image.convert('RGB')
+    model_input_size = (1024, 1024)
+    image = image.resize(model_input_size, Image.BILINEAR)
+    return image
+
+class GFrbmg2:
+    def __init__(self):
+        self.model = None
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "invert_mask": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
+    RETURN_NAMES = ("image_rgba", "mask", "image_black")
+    FUNCTION = "remove_background"
+    CATEGORY = "🐵 GorillaFrame/Image"
+  
+    def remove_background(self, image, invert_mask):
+        if self.model is None:
+            self.model = AutoModelForImageSegmentation.from_pretrained(
+                os.path.join(folder_paths.models_dir, "RMBG-2.0"),
+                trust_remote_code=True,
+                local_files_only=True
+            )
+            self.model.to(device)
+            self.model.eval()
+
+        processed_images = []
+        processed_masks = []
+        processed_blacks = []
+
+        for img in image:
+            orig_image = tensor2pil(img)
+            w,h = orig_image.size
+            image = resize_image(orig_image)
+            im_np = np.array(image)
+            im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
+            im_tensor = torch.unsqueeze(im_tensor,0)
+            im_tensor = torch.divide(im_tensor,255.0)
+            im_tensor = normalize(im_tensor,[0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
+            if torch.cuda.is_available():
+                im_tensor=im_tensor.cuda()
+
+            with torch.no_grad():
+                result = self.model(im_tensor)[-1].sigmoid().cpu()
+            
+            result = result[0].squeeze()
+            result = F.interpolate(result.unsqueeze(0).unsqueeze(0), size=(h,w), mode='bilinear').squeeze()
+            
+            if invert_mask:
+                result = 1 - result
+                
+            mask_pil = tensor2pil(result)
+            
+            # RGBA image
+            rgba_image = orig_image.copy()
+            rgba_image.putalpha(mask_pil)
+            
+            # Black background image
+            black_image = Image.new('RGB', orig_image.size, (0, 0, 0))
+            black_image.paste(orig_image, mask=mask_pil)
+
+            processed_images.append(pil2tensor(rgba_image))
+            processed_masks.append(pil2tensor(mask_pil))
+            processed_blacks.append(pil2tensor(black_image))
+
+        new_images = torch.cat(processed_images, dim=0)
+        new_masks = torch.cat(processed_masks, dim=0)
+        new_blacks = torch.cat(processed_blacks, dim=0)
+
+        return new_images, new_masks, new_blacks
\ No newline at end of file
diff --git a/readme.md b/readme.md
new file mode 100644
index 0000000..feaa13f
--- /dev/null
+++ b/readme.md
@@ -0,0 +1,2 @@
+# GF_nodes
+