# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for Object removal using SDMatte model.
#
# reference from https://github.com/vivoCameraResearch/SDMatte
# model: https://huggingface.co/1038lab/SDMatte
#
# 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 sys
import copy
from pathlib import Path

if os.environ.get('SDMATTE_CPU_ONLY', '').lower() in ('1', 'true', 'yes'):
    os.environ['CUDA_VISIBLE_DEVICES'] = ''

import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image, ImageFilter
from torch.hub import download_url_to_file
from torchvision import transforms
from torchvision.transforms import InterpolationMode

import folder_paths

try:
    import comfy.model_management
    COMFY_AVAILABLE = True
except Exception as e:
    print(f"Warning: ComfyUI model management not available: {e}")
    COMFY_AVAILABLE = False
    class MockModelManagement:
        @staticmethod
        def get_torch_device():
            return torch.device('cpu')

    class MockComfy:
        model_management = MockModelManagement()

    comfy = MockComfy()

try:
    from safetensors.torch import load_file
    SAFETENSORS_AVAILABLE = True
except ImportError:
    SAFETENSORS_AVAILABLE = False
    print("Warning: safetensors not available. Will use torch.load for model loading.")

try:
    import diffusers
    import transformers
    DIFFUSERS_AVAILABLE = True
except ImportError:
    DIFFUSERS_AVAILABLE = False
    print("Warning: diffusers/transformers not available. SDMatte functionality will be limited.")

current_dir = Path(__file__).resolve().parent
repo_root = current_dir.parent
sdmatte_path = repo_root / "models" / "SDMatte"
sys.path.insert(0, str(sdmatte_path))

SDMATTE_MODELS = {
    "SDMatte": {
        "model_url": "https://huggingface.co/1038lab/SDMatte/resolve/main/SDMatte.safetensors",
        "filename": "SDMatte.safetensors",
        "repo_id": "1038lab/SDMatte"
    },
    "SDMatte_plus": {
        "model_url": "https://huggingface.co/1038lab/SDMatte/resolve/main/SDMatte_plus.safetensors",
        "filename": "SDMatte_plus.safetensors",
        "repo_id": "1038lab/SDMatte"
    }
}

REQUIRED_COMPONENTS = ["scheduler", "text_encoder", "tokenizer", "unet", "vae"]

def get_or_download_model_file(filename, url, dirname):
    local_path = folder_paths.get_full_path(dirname, filename)
    if local_path:
        return local_path
    folder = os.path.join(folder_paths.models_dir, dirname)
    os.makedirs(folder, exist_ok=True)
    local_path = os.path.join(folder, filename)
    if not os.path.exists(local_path):
        print(f"Downloading {filename} from {url} ...")
        try:
            download_url_to_file(url, local_path)
        except Exception as e:
            raise RuntimeError(f"Failed to download {filename} from {url}: {e}")
    return local_path

def ensure_model_components(model_name):
    model_info = SDMATTE_MODELS[model_name]
    repo_id = model_info["repo_id"]

    components_dir = os.path.join(folder_paths.models_dir, "RMBG", "SDMatte")

    missing_components = []
    for component in REQUIRED_COMPONENTS:
        component_path = os.path.join(components_dir, component)
        if not os.path.exists(component_path) or not os.listdir(component_path):
            missing_components.append(component)

    if missing_components:
        print(f"Downloading missing SDMatte components: {missing_components}")
        base_url = f"https://huggingface.co/{repo_id}/resolve/main"

        for component in missing_components:
            component_dir = os.path.join(components_dir, component)
            os.makedirs(component_dir, exist_ok=True)

            if component == "scheduler":
                files = ["scheduler_config.json"]
            elif component == "text_encoder":
                files = ["config.json"]
            elif component == "tokenizer":
                files = ["merges.txt", "special_tokens_map.json", "tokenizer_config.json", "vocab.json"]
            elif component == "unet":
                files = ["config.json"]
            elif component == "vae":
                files = ["config.json"]

            for file in files:
                file_url = f"{base_url}/{component}/{file}"
                file_path = os.path.join(component_dir, file)
                if not os.path.exists(file_path):
                    try:
                        print(f"  Downloading {component}/{file}...")
                        download_url_to_file(file_url, file_path)
                    except Exception as e:
                        print(f"  Warning: Failed to download {file}: {e}")

    return components_dir

def process_mask(mask_image: Image.Image, invert_output: bool = False,
                mask_blur: int = 0, mask_offset: int = 0) -> Image.Image:
    if invert_output:
        mask_np = np.array(mask_image)
        mask_image = Image.fromarray(255 - mask_np)
    if mask_blur > 0:
        mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
    if mask_offset != 0:
        filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter
        size = abs(mask_offset) * 2 + 1
        for _ in range(abs(mask_offset)):
            mask_image = mask_image.filter(filter_type(size))
    return mask_image

def apply_background_color(image: Image.Image, mask_image: Image.Image,
                         background: str = "Alpha",
                         background_color: str = "#222222") -> Image.Image:
    rgba_image = image.copy().convert('RGBA')
    rgba_image.putalpha(mask_image.convert('L'))

    if background == "Color":
        def hex_to_rgba(hex_color):
            hex_color = hex_color.lstrip('#')
            r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
            return (r, g, b, 255)
        rgba = hex_to_rgba(background_color)
        bg_image = Image.new('RGBA', image.size, rgba)
        composite_image = Image.alpha_composite(bg_image, rgba_image)
        return composite_image.convert('RGB')
    return rgba_image

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

def tensor2pil(image):
    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

SDMatteCore = None

def _resize_norm_image_bchw(image_bchw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor:
    if image_bchw.shape[1] == 4:
        image_bchw = image_bchw[:, :3, :, :]

    resize = transforms.Resize(size_hw, interpolation=InterpolationMode.BILINEAR, antialias=True)
    norm = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    x = resize(image_bchw)
    x = norm(x)
    return x

def _resize_mask_b1hw(mask_b1hw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor:
    resize = transforms.Resize(size_hw, interpolation=InterpolationMode.BILINEAR, antialias=True)
    return resize(mask_b1hw)

class AILab_SDMatte:
    @classmethod
    def INPUT_TYPES(cls):
        tooltips = {
            "model": "SDMatte model variant: Standard or Plus version",
            "image": "Input image for matting extraction",
            "mask": "Mask: White=foreground, Black=background. If omitted and image has alpha, alpha will be used.",
            "process_res": "Processing resolution: higher = better quality but slower",
            "device": "Auto: smart detection, CPU: force CPU, GPU: force GPU",
            "transparent_object": "Whether input image contains transparent objects",
            "mask_refine": "Enable mask refinement using mask constraints",
            "sensitivity": "Sensitivity for mask constraint (0.1-1.0): higher = more strict",
            "mask_blur": "Blur mask edges (0 = disabled)",
            "mask_offset": "Expand/shrink mask (positive = expand)",
            "invert_output": "Invert the mask output",
            "background": "Background type for output",
            "background_color": "Background color (when not Alpha)",
        }
        return {
            "required": {
                "image": ("IMAGE",),
                "model": (list(SDMATTE_MODELS.keys()), {"default": "SDMatte", "tooltip": tooltips["model"]}),
                "device": (["Auto", "CPU", "GPU"], {"default": "Auto", "tooltip": tooltips["device"]}),
                "process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 8, "tooltip": tooltips["process_res"]}),
            },
            "optional": {
                "mask": ("MASK", {"tooltip": tooltips["mask"]}),
                "transparent_object": ("BOOLEAN", {"default": True, "tooltip": tooltips["transparent_object"]}),
                "mask_refine": ("BOOLEAN", {"default": True, "tooltip": tooltips["mask_refine"]}),
                "sensitivity": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 1.0, "step": 0.1, "tooltip": tooltips["sensitivity"]}),
                "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
                "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
                "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
                "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
                "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
    RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
    FUNCTION = "matting_inference"
    CATEGORY = "🧪AILab/🧽RMBG"

    def __init__(self):
        self.model_cache = {}

    def load_sdmatte_model(self, model_name, device="Auto"):
        cache_key = f"{model_name}_{device}"
        
        current_model_keys = [k for k in self.model_cache.keys() if k.startswith(model_name)]
        if cache_key not in self.model_cache and len(self.model_cache) > 0:
            for key in list(self.model_cache.keys()):
                if key not in current_model_keys:
                    del self.model_cache[key]
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        
        if cache_key not in self.model_cache:
            if not DIFFUSERS_AVAILABLE:
                raise ImportError("diffusers and transformers are required for SDMatte functionality")
        
            global SDMatteCore
            if SDMatteCore is None:
                import sys
                import os
                current_dir = os.path.dirname(__file__)
                if current_dir not in sys.path:
                    sys.path.insert(0, current_dir)
                from SDMatte.modeling.SDMatte.meta_arch import SDMatte as SDMatteCore
        
            model_info = SDMATTE_MODELS[model_name]
        
            model_path = get_or_download_model_file(
                model_info["filename"],
                model_info["model_url"],
                "RMBG/SDMatte"
            )
        
            pretrained_repo = ensure_model_components(model_name)
        
            sdmatte_model = SDMatteCore(
                pretrained_model_name_or_path=pretrained_repo,
                load_weight=False,
                use_aux_input=True,
                aux_input="trimap",
                use_encoder_hidden_states=True,
                use_attention_mask=True,
                add_noise=False,
            )
        
            self._load_model_weights(sdmatte_model, model_path)
        
            device_obj = comfy.model_management.get_torch_device()
            if device == "CPU":
                device_obj = torch.device('cpu')
            elif device == "GPU":
                if not torch.cuda.is_available():
                    print("SDMatte: GPU requested but CUDA not available, falling back to CPU")
                    device_obj = torch.device('cpu')
                else:
                    device_obj = comfy.model_management.get_torch_device()
        
            sdmatte_model.eval()
            sdmatte_model.to(device_obj)
        
            if device_obj.type == 'cuda':
                self._apply_memory_optimizations(sdmatte_model)
        
            self.model_cache[cache_key] = sdmatte_model
        
        return self.model_cache[cache_key]

    def _load_model_weights(self, model, model_path):
        if not SAFETENSORS_AVAILABLE:
            raise ImportError("safetensors is required for SDMatte functionality")

        try:
            state_dict = load_file(model_path)
            model.load_state_dict(state_dict, strict=False)
        except Exception as e:
            if os.path.exists(model_path):
                print(f"[SDMatte] Model file appears corrupted, deleting: {model_path}")
                os.remove(model_path)
            raise RuntimeError(f"Failed to load model weights. File may be corrupted. Please try again to re-download. Error: {e}")

    def _apply_memory_optimizations(self, model):
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass

        try:
            unet = getattr(model, 'unet', None)
            if unet is not None and hasattr(unet, 'set_attn_processor'):
                from diffusers.models.attention_processor import SlicedAttnProcessor
                unet.set_attn_processor(SlicedAttnProcessor(slice_size=1))
        except Exception:
            pass

    def matting_inference(self, image, model, process_res, device="Auto",
                   mask=None, transparent_object=True, mask_refine=True,
                   sensitivity=0.8, mask_blur=0, mask_offset=0,
                   invert_output=False, background="Alpha", background_color="#222222"):
        sdmatte_model = self.load_sdmatte_model(model, device)
        device_obj = comfy.model_management.get_torch_device()
        if device == "CPU":
            device_obj = torch.device('cpu')
        
        batch_size = image.shape[0]
        
        result_masks = []
        result_images = []
        result_mask_images = []
        
        for b in range(batch_size):
            img_pil = tensor2pil(image[b])
            B, H, W = 1, img_pil.height, img_pil.width
            orig_h, orig_w = H, W
        
            img_bchw = image[b:b+1].permute(0, 3, 1, 2).contiguous().to(device_obj)
            img_in = _resize_norm_image_bchw(img_bchw, (int(process_res), int(process_res)))
        
            if mask is not None:
                mask_b1hw = mask[b:b+1].unsqueeze(1).contiguous().to(device_obj)
                mask_for_refine = mask[b:b+1]
            else:
                if image.shape[-1] == 4:
                    alpha = image[b, :, :, 3]
                    mask_b1hw = alpha.unsqueeze(0).unsqueeze(0).contiguous().to(device_obj)
                    mask_for_refine = alpha.unsqueeze(0)
                else:
                    raise ValueError("Mask required: provide a mask or use an image with alpha.")
            
            tri = _resize_mask_b1hw(mask_b1hw, (int(process_res), int(process_res))) * 2 - 1
            data = {"image": img_in,
                    "is_trans": torch.tensor([1 if transparent_object else 0], device=device_obj),
                    "caption": [""],
                    "trimap": tri,
                    "trimap_coords": torch.tensor([[0,0,1,1]], dtype=tri.dtype, device=device_obj)}
        
            with torch.inference_mode():
                if device_obj.type == 'cuda':
                    with torch.autocast(device_type='cuda', dtype=torch.float16):
                        pred_alpha = sdmatte_model(data)
                else:
                    pred_alpha = sdmatte_model(data)
        
            out = transforms.Resize((orig_h, orig_w), interpolation=InterpolationMode.BILINEAR, antialias=True)(pred_alpha)
            out = out.squeeze(1).clamp(0, 1).detach().cpu()
        
            if mask_refine:
                out = self._refine_mask(out, mask_for_refine, sensitivity)
        
            mask_pil = Image.fromarray((out[0].numpy() * 255).astype(np.uint8), mode="L")
        
            mask_image = process_mask(mask_pil, invert_output, mask_blur, mask_offset)
        
            result_image = apply_background_color(img_pil, mask_image, background, background_color)
            if background == "Color":
                result_image = result_image.convert("RGB")
            else:
                result_image = result_image.convert("RGBA")
        
            mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
            mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3)
        
            result_masks.append(mask_tensor)
            result_images.append(pil2tensor(result_image))
            result_mask_images.append(mask_image_vis)
        
        if device_obj.type == 'cuda':
            torch.cuda.empty_cache()
        import gc
        gc.collect()
        
        return (torch.cat(result_images, dim=0), torch.cat(result_masks, dim=0), torch.cat(result_mask_images, dim=0))

    def _refine_mask(self, mask, trimap, constraint):
        trimap_cpu = trimap.cpu()
        foreground_regions = trimap_cpu > constraint
        background_regions = trimap_cpu < (1.0 - constraint)
        unknown_regions = ~(foreground_regions | background_regions)

        refined_mask = mask.clone()
        refined_mask[background_regions] = 0.0
        refined_mask[foreground_regions] = torch.clamp(refined_mask[foreground_regions] * 1.2, 0, 1)

        alpha_threshold = 0.3
        low_confidence = (refined_mask < alpha_threshold) & unknown_regions
        refined_mask[low_confidence] = 0.0

        return refined_mask


NODE_CLASS_MAPPINGS = {
    "AILab_SDMatte": AILab_SDMatte,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "AILab_SDMatte": "SDMatte Matting (RMBG)",
}