# 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.")

# Lazy: diffusers/transformers availability is checked on first use to avoid ~3.7s import at startup
DIFFUSERS_AVAILABLE = None  # Tri-state: None = not yet checked

from AILab_utils import get_device, clean_vram, unload_model, get_or_download_model_file, process_mask, apply_background_color, pil2tensor, tensor2pil

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 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

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"]}),
                "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}),
                "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:
            global DIFFUSERS_AVAILABLE
            if DIFFUSERS_AVAILABLE is None:
                try:
                    import diffusers  # noqa: F401
                    import transformers  # noqa: F401
                    DIFFUSERS_AVAILABLE = True
                except ImportError:
                    DIFFUSERS_AVAILABLE = False
            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, unload_model=False, background="Alpha", background_color="#222222"):
        if device == "CPU":
            device_obj = torch.device('cpu')
        elif device == "GPU":
            device_obj = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            device_obj = get_device()

        sdmatte_model = self.load_sdmatte_model(model, device)
        
        batch_size = image.shape[0]
        
        result_masks = []
        result_images = []
        result_mask_images = []
        
        try:
            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:
                    if mask.ndim == 2:
                        mask_single = mask.unsqueeze(0).unsqueeze(0)
                    elif mask.ndim == 3:
                        idx = min(b, mask.shape[0] - 1)
                        mask_single = mask[idx:idx+1].unsqueeze(1)
                    elif mask.ndim == 4:
                        idx = min(b, mask.shape[0] - 1)
                        mask_single = mask[idx:idx+1]
                    else:
                        mask_single = mask.view(1, 1, H, W)
                    mask_b1hw = mask_single.contiguous().to(device_obj)
                    mask_for_refine = mask_b1hw.squeeze(1).cpu()
                elif 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).cpu()
                else:
                    print("[SDMatte INFO] No input mask provided; using full foreground mask as fallback. For best results, connect a MASK from RMBG/BiRefNet/SAM node.")
                    mask_b1hw = torch.ones((1, 1, H, W), dtype=torch.float32, device=device_obj)
                    mask_for_refine = torch.ones((1, H, W), dtype=torch.float32)
                
                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)
            
            return (torch.cat(result_images, dim=0), torch.cat(result_masks, dim=0), torch.cat(result_mask_images, dim=0))
        finally:
            if unload_model:
                self.model_cache.clear()
            clean_vram()

    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)",
}