Delete AILab_LamaRemover.py
This commit is contained in:
AI Lab
@@ -1,191 +0,0 @@
|
||||
# ComfyUI-RMBG
|
||||
# This custom node for ComfyUI provides functionality for Object removal using Big-Lama model.
|
||||
#
|
||||
# reference from https://github.com/advimman/lama
|
||||
#
|
||||
# This integration script follows GPL-3.0 License.
|
||||
# When using or modifying this code, please respect both the original model licenses
|
||||
# and this integration's license terms.
|
||||
#
|
||||
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
|
||||
|
||||
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps, ImageFilter
|
||||
import folder_paths
|
||||
from comfy.model_management import get_torch_device
|
||||
from torchvision import transforms
|
||||
from huggingface_hub import hf_hub_download
|
||||
import shutil
|
||||
import gc
|
||||
|
||||
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 pil2comfy(image):
|
||||
img_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
|
||||
if len(img_tensor.shape) == 3:
|
||||
img_tensor = img_tensor.unsqueeze(0)
|
||||
return img_tensor
|
||||
|
||||
def pad_image(image, is_mask=False):
|
||||
w, h = image.size
|
||||
if w % 8 != 0:
|
||||
w = w + (8 - w % 8)
|
||||
if h % 8 != 0:
|
||||
h = h + (8 - h % 8)
|
||||
|
||||
fill_color = 0 if is_mask else None
|
||||
padded = Image.new(image.mode, (w, h), color=fill_color)
|
||||
padded.paste(image, (0, 0))
|
||||
return padded
|
||||
|
||||
def cropimage(image, w, h):
|
||||
return image.crop((0, 0, w, h))
|
||||
|
||||
DEVICE = get_torch_device()
|
||||
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
|
||||
|
||||
class AILab_LamaRemover:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
tooltips = {
|
||||
"images": "Input images to be processed",
|
||||
"masks": "Masks defining areas to be removed (white=remove)",
|
||||
"removal_strength": "Strength of the removal effect (higher values increase the effect area)",
|
||||
"edge_smoothness": "Controls edge smoothness (higher values create smoother transitions)"
|
||||
}
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {"tooltip": tooltips["images"]}),
|
||||
"masks": ("MASK", {"tooltip": tooltips["masks"]}),
|
||||
"removal_strength": ("INT", {"default": 230, "min": 0, "max": 255, "step": 1, "display": "slider", "tooltip": tooltips["removal_strength"]}),
|
||||
"edge_smoothness": ("INT", {"default": 8, "min": 0, "max": 20, "step": 1, "display": "slider", "tooltip": tooltips["edge_smoothness"]}),
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "🧪AILab/🧽RMBG"
|
||||
RETURN_NAMES = ("images",)
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "remove_object"
|
||||
|
||||
def __init__(self):
|
||||
self.model = None
|
||||
self.device = DEVICE
|
||||
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "Lama")
|
||||
self.model_path = os.path.join(self.cache_dir, "big-lama.pt")
|
||||
self.to_pil = transforms.ToPILImage()
|
||||
|
||||
def load_model(self):
|
||||
if self.model is not None:
|
||||
return
|
||||
|
||||
if not os.path.exists(self.model_path):
|
||||
self.download_model()
|
||||
|
||||
try:
|
||||
self.model = torch.jit.load(self.model_path, map_location=self.device)
|
||||
except Exception as e:
|
||||
print(f"Can't use comfy device: {str(e)}")
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
self.model = torch.jit.load(self.model_path, map_location=self.device)
|
||||
|
||||
self.model.eval()
|
||||
self.model.to(self.device)
|
||||
|
||||
def download_model(self):
|
||||
print("Downloading Big-Lama model...")
|
||||
os.makedirs(self.cache_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
downloaded_path = hf_hub_download(
|
||||
repo_id="1038lab/Lama",
|
||||
filename="big-lama.pt",
|
||||
local_dir=self.cache_dir,
|
||||
local_dir_use_symlinks=False
|
||||
)
|
||||
|
||||
if os.path.dirname(downloaded_path) != self.cache_dir:
|
||||
shutil.move(downloaded_path, self.model_path)
|
||||
|
||||
print("Big-Lama model downloaded successfully")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Error downloading Big-Lama model: {str(e)}")
|
||||
|
||||
def process_with_model(self, img_tensor, mask_tensor):
|
||||
with torch.inference_mode():
|
||||
img_tensor = img_tensor.to(self.device)
|
||||
mask_tensor = mask_tensor.to(self.device)
|
||||
|
||||
result = self.model(img_tensor, mask_tensor)
|
||||
result_cpu = result[0].cpu()
|
||||
|
||||
del img_tensor
|
||||
del mask_tensor
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return result_cpu
|
||||
|
||||
def remove_object(self, images, masks, removal_strength, edge_smoothness):
|
||||
try:
|
||||
self.load_model()
|
||||
results = []
|
||||
|
||||
for image, mask in zip(images, masks):
|
||||
ori_image = tensor2pil(image)
|
||||
w, h = ori_image.size
|
||||
p_image = pad_image(ori_image)
|
||||
|
||||
mask_np = mask.cpu().numpy()
|
||||
mask_pil = Image.fromarray((mask_np * 255).astype(np.uint8))
|
||||
p_mask = pad_image(mask_pil, is_mask=True)
|
||||
|
||||
if p_mask.size != p_image.size:
|
||||
try:
|
||||
p_mask = p_mask.resize(p_image.size, Image.LANCZOS)
|
||||
except AttributeError:
|
||||
p_mask = p_mask.resize(p_image.size, Image.ANTIALIAS)
|
||||
|
||||
p_mask = ImageOps.invert(p_mask)
|
||||
p_mask = p_mask.filter(ImageFilter.GaussianBlur(radius=edge_smoothness))
|
||||
gray = p_mask.point(lambda x: 0 if x > removal_strength else 255)
|
||||
|
||||
img_tensor = torch.FloatTensor(np.array(p_image)).permute(2, 0, 1).unsqueeze(0) / 255.0
|
||||
mask_tensor = torch.FloatTensor(np.array(gray)).unsqueeze(0).unsqueeze(0) / 255.0
|
||||
|
||||
result = self.process_with_model(img_tensor, mask_tensor)
|
||||
result_img = self.to_pil(result.squeeze())
|
||||
|
||||
if result_img.width > w or result_img.height > h:
|
||||
result_img = cropimage(result_img, w, h)
|
||||
|
||||
result_tensor = pil2comfy(result_img)
|
||||
results.append(result_tensor)
|
||||
|
||||
del result
|
||||
gc.collect()
|
||||
|
||||
return (torch.cat(results, dim=0),)
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(traceback.format_exc())
|
||||
raise RuntimeError(f"Error in object removal: {str(e)}")
|
||||
finally:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AILab_LamaRemover": AILab_LamaRemover,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AILab_LamaRemover": "Lama Remover (RMBG)",
|
||||
}
|
||||
Reference in New Issue
Block a user