diff --git a/Doc/TESTING.md b/Doc/TESTING.md index c1cdea4..ba04774 100644 --- a/Doc/TESTING.md +++ b/Doc/TESTING.md @@ -14,8 +14,8 @@ $ComfyPython = "..\..\..\python_embeded\python.exe" & $ComfyPython -m pytest -q & $ComfyPython -m pytest -q --cov=python --cov-report=term-missing -& $ComfyPython -m ruff check __init__.py canvas_node.py python tests -& $ComfyPython -m compileall -q __init__.py canvas_node.py python tests +& $ComfyPython -m ruff check __init__.py python tests +& $ComfyPython -m compileall -q __init__.py python tests ``` The explicit interpreter path is intentional: it keeps tests aligned with the diff --git a/README.md b/README.md index f418f47..3051d40 100644 --- a/README.md +++ b/README.md @@ -241,7 +241,26 @@ Click the gear button next to **Matting** to configure its behavior. Settings ar - **Processing mode**: remove the background, remove the detected foreground, or apply the generated mask to LayerForge's Draw Mask output-area mask. - **Mask threshold**: use `0` for a soft alpha mask or a higher value for a harder cutout. -The model selector separates checkpoints already installed locally from official full-size BiRefNet variants available for download. Downloaded variants are stored in separate subdirectories under `ComfyUI/models/background_removal/` and are validated against ComfyUI's native loader before use. Lite variants are intentionally excluded because they use a different architecture. +The model selector separates checkpoints already installed locally from official full-size BiRefNet variants available for download. Downloaded variants are stored in separate subdirectories under `ComfyUI/models/background_removal/` with descriptive filenames such as `BiRefNet-general.safetensors` and `BiRefNet-portrait.safetensors`, then validated against ComfyUI's native loader before use. Lite variants are intentionally excluded because they use a different architecture. Hugging Face cache folders are ignored in the selector and are not deleted. + +### BiRefNet model guide + +The specialized models are fine-tuned for different tasks. For general photos, start with **General**; choose a specialized variant when its training domain matches the image content. + +| Variant | Recommended use | Official page | +|---|---|---| +| [General](https://huggingface.co/ZhengPeng7/BiRefNet) | Everyday images and general background removal | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet) | +| [High Resolution](https://huggingface.co/ZhengPeng7/BiRefNet_HR) | Large images and detailed edges; higher memory use | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet_HR) | +| [Portrait](https://huggingface.co/ZhengPeng7/BiRefNet-portrait) | People, portraits, hair, and portrait cutouts | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet-portrait) | +| [Matting](https://huggingface.co/ZhengPeng7/BiRefNet-matting) | Soft alpha edges and semi-transparent details | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet-matting) | +| [High Resolution Matting](https://huggingface.co/ZhengPeng7/BiRefNet_HR-matting) | Fine matting details on larger images | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet_HR-matting) | +| [Dynamic](https://huggingface.co/ZhengPeng7/BiRefNet_dynamic) | Varying aspect ratios and input resolutions | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet_dynamic) | +| [Dynamic Matting](https://huggingface.co/ZhengPeng7/BiRefNet_dynamic-matting) | Arbitrary-size images with soft matting edges | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet_dynamic-matting) | +| [HRSOD](https://huggingface.co/ZhengPeng7/BiRefNet-HRSOD) | High-resolution salient-object detection | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet-HRSOD) | +| [DIS5K](https://huggingface.co/ZhengPeng7/BiRefNet-DIS5K) | Dichotomous foreground/background separation | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet-DIS5K) | +| [COD](https://huggingface.co/ZhengPeng7/BiRefNet-COD) | Camouflaged objects that blend into the background | [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet-COD) | + +All links point to the official model repositories. See the [BiRefNet project](https://github.com/ZhengPeng7/BiRefNet) for the implementation, research details, and the complete model collection. --- diff --git a/__init__.py b/__init__.py index e20bdca..322cb70 100644 --- a/__init__.py +++ b/__init__.py @@ -1,16 +1,10 @@ -import sys -import os - -# Add the custom node's directory to the Python path -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - # ComfyUI loads this file as a package. Pytest also sees the repository-level # __init__.py while discovering tests, where relative imports have no package # context and the ComfyUI runtime is not available. Keep that discovery path # inert while preserving the normal package bootstrap. LayerForgeNode = None if __package__: - from .canvas_node import LayerForgeNode + from .python.node import LayerForgeNode LayerForgeNode.setup_routes() diff --git a/canvas_node.py b/canvas_node.py deleted file mode 100644 index 5a6c89f..0000000 --- a/canvas_node.py +++ /dev/null @@ -1,1415 +0,0 @@ -from PIL import Image, ImageOps -import hashlib -import torch -import numpy as np -import folder_paths -from server import PromptServer -from aiohttp import web -import asyncio -import threading -import os -from tqdm import tqdm -from torchvision import transforms -import torch.nn.functional as F -import traceback -import uuid -import time -import base64 -from PIL import Image -import io -import sys -import os - -try: - from python.log_system import create_module_logger - - log = create_module_logger(__name__) - - log.info("Logger initialized for canvas_node") -except ImportError as e: - - print(f"Warning: Logger module not available: {e}") - - class _FallbackLogger: - @staticmethod - def debug(*args, **kwargs): print("[DEBUG]", *args) - - @staticmethod - def info(*args, **kwargs): print("[INFO]", *args) - - @staticmethod - def warning(*args, **kwargs): print("[WARN]", *args) - - @staticmethod - def error(*args, **kwargs): print("[ERROR]", *args) - - @staticmethod - def exception(*args, **kwargs): - print("[ERROR]", *args) - traceback.print_exc() - - @staticmethod - def fatal(*args, **kwargs): print("[FATAL]", *args) - - log = _FallbackLogger() - -torch.set_float32_matmul_precision('high') - - -def _get_comfy_birefnet_loader(): - """Return ComfyUI's native BiRefNet loader when it is available.""" - try: - from comfy.bg_removal_model import load - - return load - except Exception as error: - log.debug(f"Native ComfyUI BiRefNet loader is unavailable: {error}") - return None - - -class LayerForgeNode: - _canvas_data_storage = {} - _storage_lock = threading.Lock() - - _canvas_cache = { - 'image': None, - 'mask': None, - 'data_flow_status': {}, - 'persistent_cache': {}, - 'last_execution_id': None - } - - - _websocket_data = {} - _websocket_listeners = {} - - def __init__(self): - super().__init__() - self.flow_id = str(uuid.uuid4()) - self.node_id = None # Will be set when node is created - - if self.__class__._canvas_cache['persistent_cache']: - self.restore_cache() - - def restore_cache(self): - try: - persistent = self.__class__._canvas_cache['persistent_cache'] - current_execution = self.get_execution_id() - - if current_execution != self.__class__._canvas_cache['last_execution_id']: - log.info(f"New execution detected: {current_execution}") - self.__class__._canvas_cache['image'] = None - self.__class__._canvas_cache['mask'] = None - self.__class__._canvas_cache['last_execution_id'] = current_execution - else: - - if persistent.get('image') is not None: - self.__class__._canvas_cache['image'] = persistent['image'] - log.info("Restored image from persistent cache") - if persistent.get('mask') is not None: - self.__class__._canvas_cache['mask'] = persistent['mask'] - log.info("Restored mask from persistent cache") - except Exception as e: - log.error(f"Error restoring cache: {str(e)}") - - def get_execution_id(self): - - try: - - return str(int(time.time() * 1000)) - except Exception as e: - log.error(f"Error getting execution ID: {str(e)}") - return None - - def update_persistent_cache(self): - - try: - self.__class__._canvas_cache['persistent_cache'] = { - 'image': self.__class__._canvas_cache['image'], - 'mask': self.__class__._canvas_cache['mask'] - } - log.debug("Updated persistent cache") - except Exception as e: - log.error(f"Error updating persistent cache: {str(e)}") - - def track_data_flow(self, stage, status, data_info=None): - - flow_status = { - 'timestamp': time.time(), - 'stage': stage, - 'status': status, - 'data_info': data_info - } - log.debug(f"Data Flow [{self.flow_id}] - Stage: {stage}, Status: {status}") - if data_info: - log.debug(f"Data Info: {data_info}") - - self.__class__._canvas_cache['data_flow_status'][self.flow_id] = flow_status - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "fit_on_add": ("BOOLEAN", {"default": False, "label_on": "Fit on Add/Paste", "label_off": "Default Behavior"}), - "show_preview": ("BOOLEAN", {"default": False, "label_on": "Show Preview", "label_off": "Hide Preview"}), - "auto_refresh_after_generation": ("BOOLEAN", {"default": False, "label_on": "True", "label_off": "False"}), - "trigger": ("INT", {"default": 0, "min": 0, "max": 99999999, "step": 1}), - "node_id": ("STRING", {"default": "0"}), - }, - "optional": { - "input_image": ("IMAGE",), - "input_mask": ("MASK",), - }, - "hidden": { - "prompt": ("PROMPT",), - "unique_id": ("UNIQUE_ID",), - } - } - - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("image", "mask") - FUNCTION = "process_canvas_image" - CATEGORY = "azNodes > LayerForge" - - def add_image_to_canvas(self, input_image): - - try: - - if not isinstance(input_image, torch.Tensor): - raise ValueError("Input image must be a torch.Tensor") - - if input_image.dim() == 4: - input_image = input_image.squeeze(0) - - if input_image.dim() == 3 and input_image.shape[0] in [1, 3]: - input_image = input_image.permute(1, 2, 0) - - return input_image - - except Exception as e: - log.error(f"Error in add_image_to_canvas: {str(e)}") - return None - - def add_mask_to_canvas(self, input_mask, input_image): - - try: - - if not isinstance(input_mask, torch.Tensor): - raise ValueError("Input mask must be a torch.Tensor") - - if input_mask.dim() == 4: - input_mask = input_mask.squeeze(0) - if input_mask.dim() == 3 and input_mask.shape[0] == 1: - input_mask = input_mask.squeeze(0) - - if input_image is not None: - expected_shape = input_image.shape[:2] - if input_mask.shape != expected_shape: - input_mask = F.interpolate( - input_mask.unsqueeze(0).unsqueeze(0), - size=expected_shape, - mode='bilinear', - align_corners=False - ).squeeze() - - return input_mask - - except Exception as e: - log.error(f"Error in add_mask_to_canvas: {str(e)}") - return None - - _processing_lock = threading.Lock() - - def process_canvas_image(self, fit_on_add, show_preview, auto_refresh_after_generation, trigger, node_id, input_image=None, input_mask=None, prompt=None, unique_id=None): - - try: - - if not self.__class__._processing_lock.acquire(blocking=False): - log.warning(f"Process already in progress for node {node_id}, skipping...") - - return self.get_cached_data() - - log.info(f"Lock acquired. Starting process_canvas_image for node_id: {node_id} (fallback unique_id: {unique_id})") - - # Always store fresh input data, even if None, to clear stale data - log.info(f"Storing input data for node {node_id} - Image: {input_image is not None}, Mask: {input_mask is not None}") - - with self.__class__._storage_lock: - input_data = {} - - if input_image is not None: - # Convert image tensor(s) to base64 - handle batch - if isinstance(input_image, torch.Tensor): - # Ensure correct shape [B, H, W, C] - if input_image.dim() == 3: - input_image = input_image.unsqueeze(0) - - batch_size = input_image.shape[0] - log.info(f"Processing batch of {batch_size} image(s)") - - if batch_size == 1: - # Single image - keep backward compatibility - img_np = (input_image.squeeze(0).cpu().numpy() * 255).astype(np.uint8) - pil_img = Image.fromarray(img_np, 'RGB') - - # Convert to base64 - buffered = io.BytesIO() - pil_img.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - input_data['input_image'] = f"data:image/png;base64,{img_str}" - input_data['input_image_width'] = pil_img.width - input_data['input_image_height'] = pil_img.height - log.debug(f"Stored single input image: {pil_img.width}x{pil_img.height}") - else: - # Multiple images - store as array - images_array = [] - for i in range(batch_size): - img_np = (input_image[i].cpu().numpy() * 255).astype(np.uint8) - pil_img = Image.fromarray(img_np, 'RGB') - - # Convert to base64 - buffered = io.BytesIO() - pil_img.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - images_array.append({ - 'data': f"data:image/png;base64,{img_str}", - 'width': pil_img.width, - 'height': pil_img.height - }) - log.debug(f"Stored batch image {i+1}/{batch_size}: {pil_img.width}x{pil_img.height}") - - input_data['input_images_batch'] = images_array - log.info(f"Stored batch of {batch_size} images") - - if input_mask is not None: - # Convert mask tensor to base64 - if isinstance(input_mask, torch.Tensor): - # Ensure correct shape - if input_mask.dim() == 2: - input_mask = input_mask.unsqueeze(0) - if input_mask.dim() == 3 and input_mask.shape[0] == 1: - input_mask = input_mask.squeeze(0) - - # Convert to numpy and then to PIL - mask_np = (input_mask.cpu().numpy() * 255).astype(np.uint8) - pil_mask = Image.fromarray(mask_np, 'L') - - # Convert to base64 - mask_buffered = io.BytesIO() - pil_mask.save(mask_buffered, format="PNG") - mask_str = base64.b64encode(mask_buffered.getvalue()).decode() - input_data['input_mask'] = f"data:image/png;base64,{mask_str}" - log.debug(f"Stored input mask: {pil_mask.width}x{pil_mask.height}") - - input_data['fit_on_add'] = fit_on_add - - # Store in a special key for input data (overwrites any previous data) - self.__class__._canvas_data_storage[f"{node_id}_input"] = input_data - - storage_key = node_id - - processed_image = None - processed_mask = None - - with self.__class__._storage_lock: - canvas_data = self.__class__._canvas_data_storage.pop(storage_key, None) - - if canvas_data: - log.info(f"Canvas data found for node {storage_key} from WebSocket") - if canvas_data.get('image'): - image_data = canvas_data['image'].split(',')[1] - image_bytes = base64.b64decode(image_data) - pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB') - image_array = np.array(pil_image).astype(np.float32) / 255.0 - processed_image = torch.from_numpy(image_array)[None,] - log.debug(f"Image loaded from WebSocket, shape: {processed_image.shape}") - - if canvas_data.get('mask'): - mask_data = canvas_data['mask'].split(',')[1] - mask_bytes = base64.b64decode(mask_data) - pil_mask = Image.open(io.BytesIO(mask_bytes)).convert('L') - mask_array = np.array(pil_mask).astype(np.float32) / 255.0 - processed_mask = torch.from_numpy(mask_array)[None,] - log.debug(f"Mask loaded from WebSocket, shape: {processed_mask.shape}") - else: - log.warning(f"No canvas data found for node {storage_key} in WebSocket cache.") - - if processed_image is None: - log.warning(f"Processed image is still None, creating default blank image.") - processed_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) - if processed_mask is None: - log.warning(f"Processed mask is still None, creating default blank mask.") - processed_mask = torch.zeros((1, 512, 512), dtype=torch.float32) - - log.debug(f"About to return output - Image shape: {processed_image.shape}, Mask shape: {processed_mask.shape}") - - self.update_persistent_cache() - - log.info(f"Successfully returning processed image and mask") - return (processed_image, processed_mask) - - except Exception as e: - log.exception(f"Error in process_canvas_image: {str(e)}") - return (None, None) - - finally: - - if self.__class__._processing_lock.locked(): - self.__class__._processing_lock.release() - log.debug(f"Process completed for node {node_id}, lock released") - - def get_cached_data(self): - return { - 'image': self.__class__._canvas_cache['image'], - 'mask': self.__class__._canvas_cache['mask'] - } - - @classmethod - def api_get_data(cls, node_id): - try: - return { - 'success': True, - 'data': cls._canvas_cache - } - except Exception as e: - return { - 'success': False, - 'error': str(e) - } - - @classmethod - def get_latest_image(cls): - output_dir = folder_paths.get_output_directory() - files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if - os.path.isfile(os.path.join(output_dir, f))] - - image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))] - - if not image_files: - return None - - latest_image_path = max(image_files, key=os.path.getctime) - return latest_image_path - - @classmethod - def get_latest_images(cls, since_timestamp=0): - output_dir = folder_paths.get_output_directory() - files = [] - for f_name in os.listdir(output_dir): - file_path = os.path.join(output_dir, f_name) - if os.path.isfile(file_path) and file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')): - try: - mtime = os.path.getmtime(file_path) - if mtime > since_timestamp: - files.append((mtime, file_path)) - except OSError: - continue - - files.sort(key=lambda x: x[0]) - - return [f[1] for f in files] - - @classmethod - def get_flow_status(cls, flow_id=None): - - if flow_id: - return cls._canvas_cache['data_flow_status'].get(flow_id) - return cls._canvas_cache['data_flow_status'] - - @classmethod - def _cleanup_old_websocket_data(cls): - """Clean up old WebSocket data from invalid nodes or data older than 5 minutes""" - try: - current_time = time.time() - cleanup_threshold = 300 # 5 minutes - - nodes_to_remove = [] - for node_id, data in cls._websocket_data.items(): - - if node_id < 0: - nodes_to_remove.append(node_id) - continue - - if current_time - data.get('timestamp', 0) > cleanup_threshold: - nodes_to_remove.append(node_id) - continue - - for node_id in nodes_to_remove: - del cls._websocket_data[node_id] - log.debug(f"Cleaned up old WebSocket data for node {node_id}") - - if nodes_to_remove: - log.info(f"Cleaned up {len(nodes_to_remove)} old WebSocket entries") - - except Exception as e: - log.error(f"Error during WebSocket cleanup: {str(e)}") - - @classmethod - def setup_routes(cls): - @PromptServer.instance.routes.get("/layerforge/canvas_ws") - async def handle_canvas_websocket(request): - ws = web.WebSocketResponse(max_msg_size=33554432) - await ws.prepare(request) - - async for msg in ws: - if msg.type == web.WSMsgType.TEXT: - try: - data = msg.json() - node_id = data.get('nodeId') - if not node_id: - await ws.send_json({'status': 'error', 'message': 'nodeId is required'}) - continue - - image_data = data.get('image') - mask_data = data.get('mask') - - with cls._storage_lock: - cls._canvas_data_storage[node_id] = { - 'image': image_data, - 'mask': mask_data, - 'timestamp': time.time() - } - - log.info(f"Received canvas data for node {node_id} via WebSocket") - - ack_payload = { - 'type': 'ack', - 'nodeId': node_id, - 'status': 'success' - } - await ws.send_json(ack_payload) - log.debug(f"Sent ACK for node {node_id}") - - except Exception as e: - log.error(f"Error processing WebSocket message: {e}") - await ws.send_json({'status': 'error', 'message': str(e)}) - elif msg.type == web.WSMsgType.ERROR: - log.error(f"WebSocket connection closed with exception {ws.exception()}") - - log.info("WebSocket connection closed") - return ws - - @PromptServer.instance.routes.get("/layerforge/get_input_data/{node_id}") - async def get_input_data(request): - try: - node_id = request.match_info["node_id"] - log.debug(f"Checking for input data for node: {node_id}") - - with cls._storage_lock: - input_key = f"{node_id}_input" - input_data = cls._canvas_data_storage.get(input_key, None) - - if input_data: - log.info(f"Input data found for node {node_id}, sending to frontend") - return web.json_response({ - 'success': True, - 'has_input': True, - 'data': input_data - }) - else: - log.debug(f"No input data found for node {node_id}") - return web.json_response({ - 'success': True, - 'has_input': False - }) - - except Exception as e: - log.error(f"Error in get_input_data: {str(e)}") - return web.json_response({ - 'success': False, - 'error': str(e) - }, status=500) - - @PromptServer.instance.routes.post("/layerforge/clear_input_data/{node_id}") - async def clear_input_data(request): - try: - node_id = request.match_info["node_id"] - log.info(f"Clearing input data for node: {node_id}") - - with cls._storage_lock: - input_key = f"{node_id}_input" - if input_key in cls._canvas_data_storage: - del cls._canvas_data_storage[input_key] - log.info(f"Input data cleared for node {node_id}") - else: - log.debug(f"No input data to clear for node {node_id}") - - return web.json_response({ - 'success': True, - 'message': f'Input data cleared for node {node_id}' - }) - - except Exception as e: - log.error(f"Error in clear_input_data: {str(e)}") - return web.json_response({ - 'success': False, - 'error': str(e) - }, status=500) - - @PromptServer.instance.routes.get("/ycnode/get_canvas_data/{node_id}") - async def get_canvas_data(request): - try: - node_id = request.match_info["node_id"] - log.debug(f"Received request for node: {node_id}") - - cache_data = cls._canvas_cache - log.debug(f"Cache content: {cache_data}") - log.debug(f"Image in cache: {cache_data['image'] is not None}") - - response_data = { - 'success': True, - 'data': { - 'image': None, - 'mask': None - } - } - - if cache_data['image'] is not None: - pil_image = cache_data['image'] - buffered = io.BytesIO() - pil_image.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - response_data['data']['image'] = f"data:image/png;base64,{img_str}" - - if cache_data['mask'] is not None: - pil_mask = cache_data['mask'] - mask_buffer = io.BytesIO() - pil_mask.save(mask_buffer, format="PNG") - mask_str = base64.b64encode(mask_buffer.getvalue()).decode() - response_data['data']['mask'] = f"data:image/png;base64,{mask_str}" - - return web.json_response(response_data) - - except Exception as e: - log.error(f"Error in get_canvas_data: {str(e)}") - return web.json_response({ - 'success': False, - 'error': str(e) - }) - - @PromptServer.instance.routes.get("/layerforge/get-latest-images/{since}") - async def get_latest_images_route(request): - try: - since_timestamp = float(request.match_info.get('since', 0)) - # JS Timestamps are in milliseconds, Python's are in seconds - latest_image_paths = cls.get_latest_images(since_timestamp / 1000.0) - - images_data = [] - for image_path in latest_image_paths: - with open(image_path, "rb") as f: - encoded_string = base64.b64encode(f.read()).decode('utf-8') - images_data.append(f"data:image/png;base64,{encoded_string}") - - return web.json_response({ - 'success': True, - 'images': images_data - }) - except Exception as e: - log.error(f"Error in get_latest_images_route: {str(e)}") - return web.json_response({ - 'success': False, - 'error': str(e) - }, status=500) - - @PromptServer.instance.routes.get("/ycnode/get_latest_image") - async def get_latest_image_route(request): - try: - latest_image_path = cls.get_latest_image() - if latest_image_path: - with open(latest_image_path, "rb") as f: - encoded_string = base64.b64encode(f.read()).decode('utf-8') - return web.json_response({ - 'success': True, - 'image_data': f"data:image/png;base64,{encoded_string}" - }) - else: - return web.json_response({ - 'success': False, - 'error': 'No images found in output directory.' - }, status=404) - except Exception as e: - return web.json_response({ - 'success': False, - 'error': str(e) - }, status=500) - - @PromptServer.instance.routes.post("/ycnode/load_image_from_path") - async def load_image_from_path_route(request): - try: - data = await request.json() - file_path = data.get('file_path') - - if not file_path: - return web.json_response({ - 'success': False, - 'error': 'file_path is required' - }, status=400) - - log.info(f"Attempting to load image from path: {file_path}") - - # Check if file exists and is accessible - if not os.path.exists(file_path): - log.warning(f"File not found: {file_path}") - return web.json_response({ - 'success': False, - 'error': f'File not found: {file_path}' - }, status=404) - - # Check if it's an image file - valid_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp', '.tiff', '.tif', '.ico', '.avif') - if not file_path.lower().endswith(valid_extensions): - return web.json_response({ - 'success': False, - 'error': f'Invalid image file extension. Supported: {valid_extensions}' - }, status=400) - - # Try to load and convert the image - try: - with Image.open(file_path) as img: - # Convert to RGB if necessary - if img.mode != 'RGB': - img = img.convert('RGB') - - # Convert to base64 - buffered = io.BytesIO() - img.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode('utf-8') - - log.info(f"Successfully loaded image from path: {file_path}") - return web.json_response({ - 'success': True, - 'image_data': f"data:image/png;base64,{img_str}", - 'width': img.width, - 'height': img.height - }) - - except Exception as img_error: - log.error(f"Error processing image file {file_path}: {str(img_error)}") - return web.json_response({ - 'success': False, - 'error': f'Error processing image file: {str(img_error)}' - }, status=500) - - except Exception as e: - log.error(f"Error in load_image_from_path_route: {str(e)}") - return web.json_response({ - 'success': False, - 'error': str(e) - }, status=500) - - def store_image(self, image_data): - - if isinstance(image_data, str) and image_data.startswith('data:image'): - image_data = image_data.split(',')[1] - image_bytes = base64.b64decode(image_data) - self.cached_image = Image.open(io.BytesIO(image_bytes)) - else: - self.cached_image = image_data - - def get_cached_image(self): - - if self.cached_image: - buffered = io.BytesIO() - self.cached_image.save(buffered, format="PNG") - img_str = base64.b64encode(buffered.getvalue()).decode() - return f"data:image/png;base64,{img_str}" - return None - - -_BIREFNET_REPOSITORY = "ZhengPeng7/BiRefNet" -_BIREFNET_FILENAME = "model.safetensors" -_BIREFNET_REMOTE_PREFIX = "remote:" -_BIREFNET_REMOTE_DIRECTORY = "layerforge_birefnet" -_BIREFNET_REQUIRED_KEYS = { - "bb.layers.1.blocks.0.attn.relative_position_index", - "bb.layers.2.blocks.17.attn.qkv.weight", -} -_BIREFNET_MODEL_CATALOG = ( - { - "id": "general", - "label": "BiRefNet — General", - "description": "General-purpose foreground segmentation.", - "repo_id": "ZhengPeng7/BiRefNet", - "filename": "model.safetensors", - }, - { - "id": "high_resolution", - "label": "BiRefNet — High Resolution", - "description": "High-resolution foreground segmentation.", - "repo_id": "ZhengPeng7/BiRefNet_HR", - "filename": "model.safetensors", - }, - { - "id": "portrait", - "label": "BiRefNet — Portrait", - "description": "Portrait-focused foreground segmentation.", - "repo_id": "ZhengPeng7/BiRefNet-portrait", - "filename": "model.safetensors", - }, - { - "id": "matting", - "label": "BiRefNet — Matting", - "description": "Matting-focused foreground segmentation.", - "repo_id": "ZhengPeng7/BiRefNet-matting", - "filename": "model.safetensors", - }, - { - "id": "high_resolution_matting", - "label": "BiRefNet — High Resolution Matting", - "description": "High-resolution matting-focused segmentation.", - "repo_id": "ZhengPeng7/BiRefNet_HR-matting", - "filename": "model.safetensors", - }, - { - "id": "dynamic", - "label": "BiRefNet — Dynamic", - "description": "Dynamic-resolution foreground segmentation.", - "repo_id": "ZhengPeng7/BiRefNet_dynamic", - "filename": "model.safetensors", - }, - { - "id": "dynamic_matting", - "label": "BiRefNet — Dynamic Matting", - "description": "Dynamic-resolution matting-focused segmentation.", - "repo_id": "ZhengPeng7/BiRefNet_dynamic-matting", - "filename": "model.safetensors", - }, - { - "id": "hrsod", - "label": "BiRefNet — HRSOD", - "description": "High-resolution salient-object segmentation.", - "repo_id": "ZhengPeng7/BiRefNet-HRSOD", - "filename": "model.safetensors", - }, - { - "id": "dis5k", - "label": "BiRefNet — DIS5K", - "description": "Dichotomous image segmentation model.", - "repo_id": "ZhengPeng7/BiRefNet-DIS5K", - "filename": "model.safetensors", - }, - { - "id": "cod", - "label": "BiRefNet — COD", - "description": "Camouflaged-object segmentation model.", - "repo_id": "ZhengPeng7/BiRefNet-COD", - "filename": "model.safetensors", - }, -) - - -def _get_birefnet_base_paths(): - """Return native ComfyUI and legacy locations that may contain BiRefNet.""" - paths = [] - - get_folder_paths = getattr(folder_paths, "get_folder_paths", None) - if callable(get_folder_paths): - try: - paths.extend(get_folder_paths("background_removal")) - except (KeyError, TypeError): - pass - - comfy_models_dir = getattr(folder_paths, "models_dir", None) - if comfy_models_dir: - paths.extend([ - os.path.join(comfy_models_dir, "background_removal"), - os.path.join(comfy_models_dir, "RMBG", "BiRefNet"), - os.path.join(comfy_models_dir, "BiRefNet"), - ]) - - unique_paths = [] - seen = set() - for path in paths: - normalized = os.path.normcase(os.path.normpath(path)) - if normalized not in seen: - seen.add(normalized) - unique_paths.append(path) - - return unique_paths - - -def _is_native_birefnet_checkpoint(path): - """Check the checkpoint signature without loading all weights into memory.""" - if not os.path.isfile(path) or not path.lower().endswith(".safetensors"): - return False - - try: - from safetensors import safe_open - - with safe_open(path, framework="pt") as checkpoint: - keys = checkpoint.keys() - return _BIREFNET_REQUIRED_KEYS.issubset(keys) - except Exception as error: - log.debug(f"Unable to inspect BiRefNet checkpoint {path}: {error}") - return False - - -def _iter_birefnet_checkpoint_paths(): - """Yield candidate checkpoints from native and legacy model directories.""" - for base_path in _get_birefnet_base_paths(): - if not os.path.isdir(base_path): - continue - - for root, directories, files in os.walk(base_path): - directories[:] = [ - directory for directory in directories - if directory not in {".git", ".no_exist", "__pycache__"} - ] - for filename in sorted(files): - if filename.lower().endswith(".safetensors"): - yield os.path.join(root, filename) - - -def _find_local_birefnet_model(model_path=None): - """Find a full BiRefNet checkpoint accepted by ComfyUI's native loader.""" - candidates = [] - seen = set() - for path in _iter_birefnet_checkpoint_paths(): - normalized = os.path.normcase(os.path.normpath(path)) - if normalized in seen: - continue - seen.add(normalized) - if _is_native_birefnet_checkpoint(path): - candidates.append(path) - - if model_path and model_path != "auto": - requested = os.path.normcase(os.path.normpath(os.path.abspath(model_path))) - for candidate in candidates: - if os.path.normcase(os.path.normpath(os.path.abspath(candidate))) == requested: - return candidate - return None - - if not candidates: - return None - - priority = { - "birefnet.safetensors": 0, - "model.safetensors": 1, - "birefnet-general.safetensors": 2, - "birefnet-hr.safetensors": 3, - } - return min( - candidates, - key=lambda path: (priority.get(os.path.basename(path).lower(), 10), path.lower()), - ) - - -def _get_birefnet_remote_model(model_path): - """Resolve a client-facing remote model identifier from the fixed catalog.""" - if not isinstance(model_path, str) or not model_path.startswith(_BIREFNET_REMOTE_PREFIX): - return None - - model_id = model_path[len(_BIREFNET_REMOTE_PREFIX):] - return next((model for model in _BIREFNET_MODEL_CATALOG if model["id"] == model_id), None) - - -def _get_birefnet_remote_checkpoint_path(model): - """Return the managed local path for a catalog model, if ComfyUI exposes storage.""" - base_paths = _get_birefnet_base_paths() - if not base_paths: - return None - - return os.path.join( - base_paths[0], - _BIREFNET_REMOTE_DIRECTORY, - model["id"], - model["filename"], - ) - - -def _get_birefnet_model_options(): - """Return local checkpoints and official models available for download.""" - local_options = [] - seen = set() - managed_roots = [ - os.path.normcase(os.path.normpath(os.path.join(base_path, _BIREFNET_REMOTE_DIRECTORY))) - for base_path in _get_birefnet_base_paths() - ] - - for path in _iter_birefnet_checkpoint_paths(): - normalized = os.path.normcase(os.path.normpath(os.path.abspath(path))) - if normalized in seen or not _is_native_birefnet_checkpoint(path): - continue - if any(normalized == root or normalized.startswith(root + os.sep) for root in managed_roots): - continue - - seen.add(normalized) - label = os.path.basename(path) - for base_path in _get_birefnet_base_paths(): - try: - relative_path = os.path.relpath(path, base_path) - except ValueError: - continue - if relative_path == os.pardir or relative_path.startswith(os.pardir + os.sep): - continue - label = relative_path.replace(os.sep, "/") - break - - local_options.append({ - "path": path, - "label": label, - "source": "local", - "downloaded": True, - }) - - remote_options = [] - for model in _BIREFNET_MODEL_CATALOG: - checkpoint_path = _get_birefnet_remote_checkpoint_path(model) - remote_options.append({ - "path": f"{_BIREFNET_REMOTE_PREFIX}{model['id']}", - "label": model["label"], - "description": model["description"], - "source": "remote", - "downloaded": bool(checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path)), - }) - - local_options.sort(key=lambda option: option["label"].lower()) - return local_options + remote_options - - -def _get_birefnet_download_dir(): - paths = _get_birefnet_base_paths() - if not paths: - raise RuntimeError("ComfyUI did not expose a background_removal model directory") - - download_dir = paths[0] - os.makedirs(download_dir, exist_ok=True) - return download_dir - - -def _download_birefnet_checkpoint(model=None): - """Download and validate a full BiRefNet checkpoint into ComfyUI's model path.""" - try: - from huggingface_hub import hf_hub_download - except ImportError as error: - raise RuntimeError( - "Automatic BiRefNet download requires the 'huggingface_hub' package. " - "Install the LayerForge requirements or place a compatible checkpoint in " - "ComfyUI/models/background_removal/." - ) from error - - if model is None: - repository = _BIREFNET_REPOSITORY - filename = _BIREFNET_FILENAME - download_dir = _get_birefnet_download_dir() - model_label = "BiRefNet — General" - else: - repository = model["repo_id"] - filename = model["filename"] - download_dir = os.path.join(_get_birefnet_download_dir(), _BIREFNET_REMOTE_DIRECTORY, model["id"]) - os.makedirs(download_dir, exist_ok=True) - model_label = model["label"] - - log.info(f"Downloading {model_label} from Hugging Face into {download_dir}...") - - try: - downloaded_path = hf_hub_download( - repo_id=repository, - filename=filename, - local_dir=download_dir, - local_dir_use_symlinks=False, - ) - except TypeError: - # Older huggingface_hub versions do not accept local_dir_use_symlinks. - downloaded_path = hf_hub_download( - repo_id=repository, - filename=filename, - local_dir=download_dir, - ) - - if not _is_native_birefnet_checkpoint(downloaded_path): - raise RuntimeError( - f"Downloaded {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" - ) - - log.info(f"{model_label} checkpoint is ready at {downloaded_path}") - return downloaded_path - - -def _ensure_birefnet_checkpoint(model_path=None): - remote_model = _get_birefnet_remote_model(model_path) - if remote_model: - checkpoint_path = _get_birefnet_remote_checkpoint_path(remote_model) - if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): - return checkpoint_path - return _download_birefnet_checkpoint(remote_model) - - checkpoint_path = _find_local_birefnet_model(model_path) - if checkpoint_path: - return checkpoint_path - - if model_path and model_path != "auto": - raise RuntimeError("The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.") - - return _download_birefnet_checkpoint() - - -class BiRefNetMatting: - _model_cache = {} - _model_cache_lock = threading.Lock() - - def __init__(self): - self.model = None - self.model_path = None - - def load_model(self, model_path=None): - loader = _get_comfy_birefnet_loader() - if loader is None: - raise RuntimeError( - "This ComfyUI version does not provide the native BiRefNet background-removal loader." - ) - - checkpoint_path = _ensure_birefnet_checkpoint(model_path) - with self._model_cache_lock: - if checkpoint_path not in self._model_cache: - log.info(f"Loading BiRefNet with ComfyUI's native loader from {checkpoint_path}") - model = loader(checkpoint_path) - if model is None: - raise RuntimeError( - f"ComfyUI did not recognize the BiRefNet checkpoint: {checkpoint_path}" - ) - self._model_cache[checkpoint_path] = model - else: - log.debug(f"Using cached native BiRefNet model from {checkpoint_path}") - - self.model = self._model_cache[checkpoint_path] - self.model_path = checkpoint_path - - return self.model - - def preprocess_image(self, image): - if not isinstance(image, torch.Tensor): - image = transforms.ToTensor()(image).unsqueeze(0) - - if image.dim() == 3: - if image.shape[0] in (1, 3, 4): - image = image.movedim(0, -1).unsqueeze(0) - else: - image = image.unsqueeze(0) - elif image.dim() == 4 and image.shape[1] in (1, 3, 4): - image = image.movedim(1, -1) - - if image.dim() != 4 or image.shape[-1] not in (1, 3, 4): - raise ValueError(f"Expected an image tensor in BCHW or BHWC format, got {tuple(image.shape)}") - - if image.shape[-1] == 1: - image = image.expand(-1, -1, -1, 3) - elif image.shape[-1] == 4: - image = image[..., :3] - - return image.to(dtype=torch.float32).contiguous() - - def execute(self, image, model_path, threshold=0.5, refinement=1, mode="remove_background"): - try: - PromptServer.instance.send_sync("matting_status", {"status": "processing"}) - - del refinement # Reserved for future edge refinement controls. - if mode not in {"remove_background", "remove_foreground", "mask_only"}: - raise ValueError(f"Unsupported matting mode: {mode}") - - threshold = float(threshold) - if not 0.0 <= threshold <= 1.0: - raise ValueError("Matting threshold must be between 0 and 1") - - image_tensor = self.preprocess_image(image) - original_size = (image_tensor.shape[1], image_tensor.shape[2]) - log.debug(f"Original size: {original_size}") - self.load_model(model_path) - log.debug(f"Processed image shape: {image_tensor.shape}, dtype: {image_tensor.dtype}") - - with torch.no_grad(): - result = self.model.encode_image(image_tensor) - if result.dim() == 3: - result = result.unsqueeze(1) - elif result.dim() == 2: - result = result.unsqueeze(0).unsqueeze(0) - else: - raise ValueError(f"Unexpected BiRefNet output shape: {tuple(result.shape)}") - - result = result.to(device=image_tensor.device, dtype=torch.float32) - if result.shape[-2:] != original_size: - result = F.interpolate( - result, - size=original_size, - mode="bilinear", - align_corners=False, - ) - result = result.clamp(0.0, 1.0) - log.debug(f"Native BiRefNet output shape: {result.shape}, dtype: {result.dtype}") - - if threshold > 0: - result = (result > threshold).to(dtype=torch.float32) - - alpha_mask = 1.0 - result if mode == "remove_foreground" else result - if mode == "mask_only": - masked_image = alpha_mask.expand(-1, 3, -1, -1) - else: - masked_image = image_tensor.movedim(-1, 1) * alpha_mask - - PromptServer.instance.send_sync("matting_status", {"status": "completed"}) - - return (masked_image, alpha_mask) - - except Exception: - PromptServer.instance.send_sync("matting_status", {"status": "error"}) - raise - - @classmethod - def IS_CHANGED(cls, image, model_path, threshold, refinement, mode="remove_background"): - - m = hashlib.md5() - m.update(str(image).encode()) - m.update(str(model_path).encode()) - m.update(str(threshold).encode()) - m.update(str(refinement).encode()) - m.update(str(mode).encode()) - return m.hexdigest() - - -_matting_lock = None - -@PromptServer.instance.routes.get("/matting/check-model") -async def check_matting_model(request): - """Check if the matting model is available and ready to use""" - try: - if _get_comfy_birefnet_loader() is None: - return web.json_response({ - "available": False, - "reason": "unsupported_comfyui", - "message": "This ComfyUI version does not provide the native BiRefNet background-removal loader." - }) - - model_options = _get_birefnet_model_options() - requested_model = request.query.get("model_path") or "auto" - remote_model = _get_birefnet_remote_model(requested_model) - - if remote_model: - local_model_path = _get_birefnet_remote_checkpoint_path(remote_model) - if local_model_path and _is_native_birefnet_checkpoint(local_model_path): - log.info(f"Selected BiRefNet model is ready at {local_model_path}") - return web.json_response({ - "available": True, - "reason": "ready", - "message": "Selected model is ready to use", - "model_path": local_model_path, - "selected_model": remote_model["label"], - "models": model_options, - }) - - log.info(f"Selected BiRefNet model is not downloaded: {remote_model['label']}") - return web.json_response({ - "available": False, - "reason": "not_downloaded", - "message": f"{remote_model['label']} will be downloaded automatically on first use.", - "model_path": requested_model, - "selected_model": remote_model["label"], - "models": model_options, - }) - - if requested_model != "auto": - local_model_path = _find_local_birefnet_model(requested_model) - if local_model_path: - log.info(f"Selected BiRefNet model is ready at {local_model_path}") - return web.json_response({ - "available": True, - "reason": "ready", - "message": "Selected model is ready to use", - "model_path": local_model_path, - "selected_model": local_model_path, - "models": model_options, - }) - - return web.json_response({ - "available": False, - "reason": "selected_model_unavailable", - "message": "The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.", - "model_path": requested_model, - "models": model_options, - }) - - local_model_path = _find_local_birefnet_model() - - if local_model_path: - log.info(f"BiRefNet model files detected at {local_model_path}") - return web.json_response({ - "available": True, - "reason": "ready", - "message": "Model is ready to use", - "model_path": local_model_path, - "models": model_options, - }) - - searched_paths = _get_birefnet_base_paths() - log.info(f"BiRefNet model not found in any of: {searched_paths}") - return web.json_response({ - "available": False, - "reason": "not_downloaded", - "message": "The BiRefNet checkpoint will be downloaded automatically on first use (requires internet connection).", - "model_path": searched_paths[0] if searched_paths else None, - "models": model_options, - }) - - except Exception as e: - log.error(f"Error checking matting model: {str(e)}") - return web.json_response({ - "available": False, - "reason": "error", - "message": f"Error checking model status: {str(e)}" - }, status=500) - -@PromptServer.instance.routes.post("/matting") -async def matting(request): - global _matting_lock - - if _matting_lock is not None: - log.warning("Matting already in progress, rejecting request") - return web.json_response({ - "error": "Another matting operation is in progress", - "details": "Please wait for the current operation to complete" - }, status=429) - - _matting_lock = True - try: - log.info("Received matting request") - data = await request.json() - - matting_instance = BiRefNetMatting() - - image_tensor, original_alpha = convert_base64_to_tensor(data["image"]) - log.debug(f"Input image shape: {image_tensor.shape}") - - mode = data.get("mode", "remove_background") - model_path = data.get("model_path") or "auto" - - matted_image, alpha_mask = matting_instance.execute( - image_tensor, - model_path, - threshold=data.get("threshold", 0.5), - refinement=data.get("refinement", 1), - mode=mode, - ) - - if mode == "mask_only": - result_image = convert_tensor_to_base64(alpha_mask) - else: - result_image = convert_tensor_to_base64(matted_image, alpha_mask, original_alpha) - result_mask = convert_tensor_to_base64(alpha_mask) - # LayerForge's Draw Mask uses white pixels as the area to remove, - # while the matting result uses white pixels as the area to keep. - draw_mask = 1.0 - alpha_mask - - return web.json_response({ - "matted_image": result_image, - "alpha_mask": result_mask, - "draw_mask": convert_tensor_to_base64(draw_mask), - "mode": mode, - "model_path": matting_instance.model_path, - }) - - except RuntimeError as e: - log.error(f"Runtime error during matting: {e}") - return web.json_response({ - "error": "Matting Model Error", - "details": str(e) - }, status=500) - except Exception as e: - log.exception(f"Error in matting endpoint: {e}") - error_text = str(e).lower() - if any( - marker in error_text - for marker in ("offline", "connection", "timed out", "huggingface", "localentrynotfound") - ): - return web.json_response({ - "error": "Network Connection Error", - "details": "Failed to download the BiRefNet model from Hugging Face. Please check your internet connection." - }, status=400) - - return web.json_response({ - "error": "An unexpected error occurred", - "details": traceback.format_exc() - }, status=500) - finally: - _matting_lock = None - log.debug("Matting lock released") - - -def convert_base64_to_tensor(base64_str): - import base64 - import io - - try: - - img_data = base64.b64decode(base64_str.split(',')[1]) - img = Image.open(io.BytesIO(img_data)) - - has_alpha = img.mode == 'RGBA' - alpha = None - if has_alpha: - alpha = img.split()[3] - - background = Image.new('RGB', img.size, (255, 255, 255)) - background.paste(img, mask=alpha) - img = background - elif img.mode != 'RGB': - img = img.convert('RGB') - - transform = transforms.ToTensor() - img_tensor = transform(img).unsqueeze(0) # [1, C, H, W] - - if has_alpha: - alpha_tensor = transforms.ToTensor()(alpha).unsqueeze(0) # [1, 1, H, W] - return img_tensor, alpha_tensor - - return img_tensor, None - - except Exception as e: - log.error(f"Error in convert_base64_to_tensor: {str(e)}") - raise - - -def convert_tensor_to_base64(tensor, alpha_mask=None, original_alpha=None): - import base64 - import io - - try: - - tensor = tensor.cpu() - - if tensor.dim() == 4: - tensor = tensor.squeeze(0) # 移除batch维度 - if tensor.dim() == 3 and tensor.shape[0] in [1, 3]: - tensor = tensor.permute(1, 2, 0) - - img_array = (tensor.numpy() * 255).astype(np.uint8) - - if alpha_mask is not None and original_alpha is not None: - - alpha_mask = alpha_mask.cpu().squeeze().numpy() - alpha_mask = (alpha_mask * 255).astype(np.uint8) - - original_alpha = original_alpha.cpu().squeeze().numpy() - original_alpha = (original_alpha * 255).astype(np.uint8) - - combined_alpha = np.minimum(alpha_mask, original_alpha) - - img = Image.fromarray(img_array, mode='RGB') - alpha_img = Image.fromarray(combined_alpha, mode='L') - img.putalpha(alpha_img) - else: - - if img_array.shape[-1] == 1: - img_array = img_array.squeeze(-1) - img = Image.fromarray(img_array, mode='L') - else: - img = Image.fromarray(img_array, mode='RGB') - - buffer = io.BytesIO() - img.save(buffer, format='PNG') - img_str = base64.b64encode(buffer.getvalue()).decode() - - return f"data:image/png;base64,{img_str}" - - except Exception as e: - log.error(f"Error in convert_tensor_to_base64: {str(e)}") - log.debug(f"Tensor shape: {tensor.shape}, dtype: {tensor.dtype}") - raise diff --git a/js/CanvasView.js b/js/CanvasView.js index 806da0a..56bf0f0 100644 --- a/js/CanvasView.js +++ b/js/CanvasView.js @@ -270,6 +270,46 @@ async function createCanvasWidget(node, widget, app) { ? settings.modelPath : 'auto'; modelSelect.value = selectedModel; + const modelDetails = document.createElement('div'); + modelDetails.className = 'lf-matting-model-details'; + const modelDescription = document.createElement('p'); + modelDescription.className = 'lf-matting-model-description'; + const modelLinks = document.createElement('div'); + modelLinks.className = 'lf-matting-model-links'; + const createModelLink = (label, url) => { + const link = document.createElement('a'); + link.className = 'lf-matting-model-link'; + link.href = url; + link.target = '_blank'; + link.rel = 'noopener noreferrer'; + link.textContent = label; + return link; + }; + const updateModelDetails = () => { + modelLinks.replaceChildren(); + const selectedOption = modelOptions.find((option) => option.path === modelSelect.value); + if (modelSelect.value === 'auto') { + modelDescription.textContent = 'LayerForge selects the best compatible installed checkpoint. If none is available, the standard General model is downloaded.'; + modelDetails.hidden = false; + return; + } + if (!selectedOption) { + modelDescription.textContent = 'The selected checkpoint is not currently available.'; + modelDetails.hidden = false; + return; + } + modelDescription.textContent = selectedOption.description || 'Installed checkpoint validated by ComfyUI\'s native BiRefNet loader.'; + if (selectedOption.url) { + modelLinks.appendChild(createModelLink('Model page', selectedOption.url)); + } + if (selectedOption.project_url) { + modelLinks.appendChild(createModelLink('BiRefNet project', selectedOption.project_url)); + } + modelDetails.hidden = false; + }; + modelSelect.onchange = updateModelDetails; + modelDetails.append(modelDescription, modelLinks); + updateModelDetails(); const modeSelect = document.createElement('select'); modeSelect.className = 'lf-matting-settings-select'; ['remove_background', 'remove_foreground', 'mask_only'].forEach((mode) => { @@ -303,7 +343,7 @@ async function createCanvasWidget(node, widget, app) { remoteCount > 0 ? `${remoteCount} official model(s) available for download` : '', ].filter(Boolean).join('; '); modelStatus.textContent = `${modelCounts}. ${modelStatusMessage}`; - body.append(createRow('Model', modelSelect, 'Only checkpoints accepted by the native ComfyUI BiRefNet loader are listed.'), createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), modelStatus); + body.append(createRow('Model', modelSelect, 'Only checkpoints accepted by the native ComfyUI BiRefNet loader are listed.'), modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), modelStatus); const actions = document.createElement('div'); actions.className = 'lf-matting-settings-actions'; const resetButton = document.createElement('button'); diff --git a/js/css/canvas_view.css b/js/css/canvas_view.css index 36d404f..d64b455 100644 --- a/js/css/canvas_view.css +++ b/js/css/canvas_view.css @@ -740,6 +740,41 @@ line-height: 1.4; } +.lf-matting-model-details { + display: flex; + flex-direction: column; + gap: 8px; + margin-top: -7px; + padding: 10px 12px; + background: rgba(255, 255, 255, 0.04); + border: 1px solid #4b4b4b; + border-radius: 4px; +} + +.lf-matting-model-description { + margin: 0; + color: #d2d2d2; + font-size: 12px; + line-height: 1.45; +} + +.lf-matting-model-links { + display: flex; + flex-wrap: wrap; + gap: 12px; +} + +.lf-matting-model-link { + color: #76adff; + font-size: 11px; + text-decoration: none; +} + +.lf-matting-model-link:hover { + color: #a9ccff; + text-decoration: underline; +} + .lf-matting-settings-threshold { display: flex; align-items: center; diff --git a/pyproject.toml b/pyproject.toml index af48035..3524a07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,8 +47,6 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F401", "E402", "F811", "I001"] -"canvas_node.py" = ["F401", "F811", "F541", "E731", "I001", "RUF010", "RUF012", "RUF034", "SIM102", "W291"] -"python/logger.py" = ["F401", "I001", "SIM910", "UP008", "UP032"] "tests/*" = ["B011", "I001", "SIM300"] [tool.ruff.lint.isort] diff --git a/python/__init__.py b/python/__init__.py index d129f96..848b8d1 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -1 +1,6 @@ -"""Backend package for the LayerForge ComfyUI custom node.""" +"""Backend package for the LayerForge ComfyUI custom node. + +The runtime modules are intentionally imported explicitly by the entry point: +``node`` owns the ComfyUI node, ``routes`` owns HTTP/WebSocket registration, +``image_utils`` owns tensor conversions, and ``matting`` owns BiRefNet. +""" diff --git a/python/image_utils.py b/python/image_utils.py new file mode 100644 index 0000000..5632360 --- /dev/null +++ b/python/image_utils.py @@ -0,0 +1,74 @@ +"""Image and tensor conversion helpers shared by LayerForge endpoints.""" + +import base64 +import io + +import numpy as np +from PIL import Image +from torchvision import transforms + +from .node import log + + +def convert_base64_to_tensor(base64_str): + """Convert a data-URL image into a BCHW tensor and optional alpha tensor.""" + try: + img_data = base64.b64decode(base64_str.split(",")[1]) + img = Image.open(io.BytesIO(img_data)) + + has_alpha = img.mode == "RGBA" + alpha = None + if has_alpha: + alpha = img.split()[3] + background = Image.new("RGB", img.size, (255, 255, 255)) + background.paste(img, mask=alpha) + img = background + elif img.mode != "RGB": + img = img.convert("RGB") + + img_tensor = transforms.ToTensor()(img).unsqueeze(0) + if has_alpha: + alpha_tensor = transforms.ToTensor()(alpha).unsqueeze(0) + return img_tensor, alpha_tensor + + return img_tensor, None + except Exception as error: + log.error(f"Error in convert_base64_to_tensor: {error}") + raise + + +def convert_tensor_to_base64(tensor, alpha_mask=None, original_alpha=None): + """Convert a tensor to a PNG data URL, optionally preserving alpha.""" + try: + tensor = tensor.cpu() + + if tensor.dim() == 4: + tensor = tensor.squeeze(0) + if tensor.dim() == 3 and tensor.shape[0] in [1, 3]: + tensor = tensor.permute(1, 2, 0) + + img_array = (tensor.numpy() * 255).astype(np.uint8) + + if alpha_mask is not None and original_alpha is not None: + alpha_mask = (alpha_mask.cpu().squeeze().numpy() * 255).astype(np.uint8) + original_alpha = (original_alpha.cpu().squeeze().numpy() * 255).astype(np.uint8) + combined_alpha = np.minimum(alpha_mask, original_alpha) + + img = Image.fromarray(img_array, mode="RGB") + img.putalpha(Image.fromarray(combined_alpha, mode="L")) + elif img_array.shape[-1] == 1: + img = Image.fromarray(img_array.squeeze(-1), mode="L") + else: + img = Image.fromarray(img_array, mode="RGB") + + buffer = io.BytesIO() + img.save(buffer, format="PNG") + img_str = base64.b64encode(buffer.getvalue()).decode() + return f"data:image/png;base64,{img_str}" + except Exception as error: + log.error(f"Error in convert_tensor_to_base64: {error}") + log.debug(f"Tensor shape: {tensor.shape}, dtype: {tensor.dtype}") + raise + + +__all__ = ["convert_base64_to_tensor", "convert_tensor_to_base64"] diff --git a/python/matting.py b/python/matting.py new file mode 100644 index 0000000..7d51269 --- /dev/null +++ b/python/matting.py @@ -0,0 +1,773 @@ +"""BiRefNet model discovery, download, inference, and HTTP endpoints.""" + +import hashlib +import os +import shutil +import threading +import traceback +from typing import ClassVar + +import folder_paths +import torch +import torch.nn.functional as F +from aiohttp import web +from server import PromptServer +from torchvision import transforms + +from .image_utils import convert_base64_to_tensor, convert_tensor_to_base64 +from .node import log + +_BIREFNET_REPOSITORY = "ZhengPeng7/BiRefNet" +_BIREFNET_FILENAME = "model.safetensors" +_BIREFNET_DEFAULT_LOCAL_FILENAME = "BiRefNet-general.safetensors" +_BIREFNET_PROJECT_URL = "https://github.com/ZhengPeng7/BiRefNet" +_BIREFNET_REMOTE_PREFIX = "remote:" +_BIREFNET_REMOTE_DIRECTORY = "layerforge_birefnet" +_BIREFNET_REQUIRED_KEYS = { + "bb.layers.1.blocks.0.attn.relative_position_index", + "bb.layers.2.blocks.17.attn.qkv.weight", +} +_BIREFNET_MODEL_CATALOG = ( + { + "id": "general", + "label": "BiRefNet — General", + "description": "The best starting point for everyday images and general background removal.", + "repo_id": "ZhengPeng7/BiRefNet", + "filename": "model.safetensors", + "local_filename": "BiRefNet-general.safetensors", + }, + { + "id": "high_resolution", + "label": "BiRefNet — High Resolution", + "description": "High-resolution segmentation for detailed edges and larger source images; it uses more memory.", + "repo_id": "ZhengPeng7/BiRefNet_HR", + "filename": "model.safetensors", + "local_filename": "BiRefNet-HR.safetensors", + }, + { + "id": "portrait", + "label": "BiRefNet — Portrait", + "description": "Portrait matting for people, hair, and portrait-focused cutouts.", + "repo_id": "ZhengPeng7/BiRefNet-portrait", + "filename": "model.safetensors", + "local_filename": "BiRefNet-portrait.safetensors", + }, + { + "id": "matting", + "label": "BiRefNet — Matting", + "description": "General matting with a focus on soft alpha edges such as hair and semi-transparent details.", + "repo_id": "ZhengPeng7/BiRefNet-matting", + "filename": "model.safetensors", + "local_filename": "BiRefNet-matting.safetensors", + }, + { + "id": "high_resolution_matting", + "label": "BiRefNet — High Resolution Matting", + "description": "High-resolution general matting for fine details; it requires more memory than the standard matting model.", + "repo_id": "ZhengPeng7/BiRefNet_HR-matting", + "filename": "model.safetensors", + "local_filename": "BiRefNet-HR-matting.safetensors", + }, + { + "id": "dynamic", + "label": "BiRefNet — Dynamic", + "description": "Dynamic-shape segmentation for inputs with varying aspect ratios and resolutions.", + "repo_id": "ZhengPeng7/BiRefNet_dynamic", + "filename": "model.safetensors", + "local_filename": "BiRefNet_dynamic.safetensors", + }, + { + "id": "dynamic_matting", + "label": "BiRefNet — Dynamic Matting", + "description": "Dynamic-shape matting for arbitrary input sizes, with a focus on soft alpha edges.", + "repo_id": "ZhengPeng7/BiRefNet_dynamic-matting", + "filename": "model.safetensors", + "local_filename": "BiRefNet_dynamic-matting.safetensors", + }, + { + "id": "hrsod", + "label": "BiRefNet — HRSOD", + "description": "High-resolution salient-object detection; useful when the main subject should stand out from its surroundings.", + "repo_id": "ZhengPeng7/BiRefNet-HRSOD", + "filename": "model.safetensors", + "local_filename": "BiRefNet-HRSOD.safetensors", + }, + { + "id": "dis5k", + "label": "BiRefNet — DIS5K", + "description": "Dichotomous image segmentation trained for clean foreground/background separation.", + "repo_id": "ZhengPeng7/BiRefNet-DIS5K", + "filename": "model.safetensors", + "local_filename": "BiRefNet-DIS5K.safetensors", + }, + { + "id": "cod", + "label": "BiRefNet — COD", + "description": "Camouflaged-object detection; use it for subjects that blend into their background.", + "repo_id": "ZhengPeng7/BiRefNet-COD", + "filename": "model.safetensors", + "local_filename": "BiRefNet-COD.safetensors", + }, +) + + +def _get_comfy_birefnet_loader(): + """Return ComfyUI's native BiRefNet loader when it is available.""" + try: + from comfy.bg_removal_model import load + + return load + except Exception as error: + log.debug(f"Native ComfyUI BiRefNet loader is unavailable: {error}") + return None + + +def _get_birefnet_base_paths(): + """Return native ComfyUI and legacy locations that may contain BiRefNet.""" + paths = [] + get_folder_paths = getattr(folder_paths, "get_folder_paths", None) + if callable(get_folder_paths): + try: + paths.extend(get_folder_paths("background_removal")) + except (KeyError, TypeError): + pass + + comfy_models_dir = getattr(folder_paths, "models_dir", None) + if comfy_models_dir: + paths.extend( + [ + os.path.join(comfy_models_dir, "background_removal"), + os.path.join(comfy_models_dir, "RMBG", "BiRefNet"), + os.path.join(comfy_models_dir, "BiRefNet"), + ] + ) + + unique_paths = [] + seen = set() + for path in paths: + normalized = os.path.normcase(os.path.normpath(path)) + if normalized not in seen: + seen.add(normalized) + unique_paths.append(path) + return unique_paths + + +def _is_native_birefnet_checkpoint(path): + """Check the checkpoint signature without loading all weights into memory.""" + if not os.path.isfile(path) or not path.lower().endswith(".safetensors"): + return False + + try: + from safetensors import safe_open + + with safe_open(path, framework="pt") as checkpoint: + return _BIREFNET_REQUIRED_KEYS.issubset(checkpoint.keys()) + except Exception as error: + log.debug(f"Unable to inspect BiRefNet checkpoint {path}: {error}") + return False + + +def _iter_birefnet_checkpoint_paths(): + """Yield candidate checkpoints from native and legacy model directories.""" + for base_path in _get_birefnet_base_paths(): + if not os.path.isdir(base_path): + continue + + for root, directories, files in os.walk(base_path): + directories[:] = [ + directory + for directory in directories + if directory not in { + ".git", + ".no_exist", + ".cache", + "__pycache__", + "refs", + "snapshots", + "blobs", + } + and not directory.startswith("models--") + ] + for filename in sorted(files): + if filename.lower().endswith(".safetensors"): + yield os.path.join(root, filename) + + +def _find_local_birefnet_model(model_path=None): + """Find a full BiRefNet checkpoint accepted by ComfyUI's native loader.""" + candidates = [] + seen = set() + for path in _iter_birefnet_checkpoint_paths(): + normalized = os.path.normcase(os.path.normpath(path)) + if normalized in seen: + continue + seen.add(normalized) + if _is_native_birefnet_checkpoint(path): + candidates.append(path) + + if model_path and model_path != "auto": + requested = os.path.normcase(os.path.normpath(os.path.abspath(model_path))) + for candidate in candidates: + if os.path.normcase(os.path.normpath(os.path.abspath(candidate))) == requested: + return candidate + return None + + if not candidates: + return None + + priority = { + "birefnet.safetensors": 0, + "model.safetensors": 1, + "birefnet-general.safetensors": 2, + "birefnet-hr.safetensors": 3, + } + return min( + candidates, + key=lambda path: (priority.get(os.path.basename(path).lower(), 10), path.lower()), + ) + + +def _get_birefnet_remote_model(model_path): + """Resolve a client-facing remote model identifier from the fixed catalog.""" + if not isinstance(model_path, str) or not model_path.startswith(_BIREFNET_REMOTE_PREFIX): + return None + + model_id = model_path[len(_BIREFNET_REMOTE_PREFIX) :] + return next((model for model in _BIREFNET_MODEL_CATALOG if model["id"] == model_id), None) + + +def _get_birefnet_remote_checkpoint_path(model): + """Return the managed local path for a catalog model.""" + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + return os.path.join( + base_paths[0], + _BIREFNET_REMOTE_DIRECTORY, + model["id"], + model["local_filename"], + ) + + +def _get_birefnet_remote_legacy_checkpoint_path(model): + """Return the pre-friendly-name path used by earlier LayerForge builds.""" + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + return os.path.join( + base_paths[0], + _BIREFNET_REMOTE_DIRECTORY, + model["id"], + model["filename"], + ) + + +def _find_existing_birefnet_remote_checkpoint(model): + """Find a managed remote checkpoint and migrate its old generic filename.""" + checkpoint_path = _get_birefnet_remote_checkpoint_path(model) + if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): + return checkpoint_path + + legacy_path = _get_birefnet_remote_legacy_checkpoint_path(model) + if not legacy_path or legacy_path == checkpoint_path or not _is_native_birefnet_checkpoint(legacy_path): + return None + + try: + os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) + os.replace(legacy_path, checkpoint_path) + log.info(f"Renamed BiRefNet checkpoint to the friendly filename {checkpoint_path}") + return checkpoint_path + except OSError as error: + log.warning(f"Unable to rename BiRefNet checkpoint {legacy_path}: {error}") + return legacy_path + + +def _get_birefnet_model_options(): + """Return local checkpoints and official models available for download.""" + _find_existing_birefnet_default_checkpoint() + local_options = [] + seen = set() + managed_roots = [ + os.path.normcase(os.path.normpath(os.path.join(base_path, _BIREFNET_REMOTE_DIRECTORY))) + for base_path in _get_birefnet_base_paths() + ] + + for path in _iter_birefnet_checkpoint_paths(): + normalized = os.path.normcase(os.path.normpath(os.path.abspath(path))) + if normalized in seen or not _is_native_birefnet_checkpoint(path): + continue + if any(normalized == root or normalized.startswith(root + os.sep) for root in managed_roots): + continue + + seen.add(normalized) + label = os.path.basename(path) + for base_path in _get_birefnet_base_paths(): + try: + relative_path = os.path.relpath(path, base_path) + except ValueError: + continue + if relative_path == os.pardir or relative_path.startswith(os.pardir + os.sep): + continue + label = relative_path.replace(os.sep, "/") + break + + local_options.append({ + "path": path, + "label": label, + "source": "local", + "downloaded": True, + }) + + remote_options = [] + for model in _BIREFNET_MODEL_CATALOG: + checkpoint_path = _find_existing_birefnet_remote_checkpoint(model) + remote_options.append({ + "path": f"{_BIREFNET_REMOTE_PREFIX}{model['id']}", + "label": model["label"], + "description": model["description"], + "url": f"https://huggingface.co/{model['repo_id']}", + "project_url": _BIREFNET_PROJECT_URL, + "source": "remote", + "downloaded": bool(checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path)), + }) + + local_options.sort(key=lambda option: option["label"].lower()) + return local_options + remote_options + + +def _get_birefnet_download_dir(): + paths = _get_birefnet_base_paths() + if not paths: + raise RuntimeError("ComfyUI did not expose a background_removal model directory") + + download_dir = paths[0] + os.makedirs(download_dir, exist_ok=True) + return download_dir + + +def _get_birefnet_default_checkpoint_path(): + """Return the friendly path used for the automatic General checkpoint.""" + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + return os.path.join(base_paths[0], _BIREFNET_DEFAULT_LOCAL_FILENAME) + + +def _find_existing_birefnet_default_checkpoint(): + """Find the automatic checkpoint and migrate its old generic filename.""" + checkpoint_path = _get_birefnet_default_checkpoint_path() + if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): + return checkpoint_path + + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + + legacy_path = os.path.join(base_paths[0], _BIREFNET_FILENAME) + if not checkpoint_path or not _is_native_birefnet_checkpoint(legacy_path): + return None + + try: + os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) + os.replace(legacy_path, checkpoint_path) + log.info(f"Renamed the automatic BiRefNet checkpoint to {checkpoint_path}") + return checkpoint_path + except OSError as error: + log.warning(f"Unable to rename the automatic BiRefNet checkpoint {legacy_path}: {error}") + return legacy_path + + +def _download_birefnet_checkpoint(model=None): + """Download and validate a full BiRefNet checkpoint into ComfyUI's model path.""" + try: + from huggingface_hub import hf_hub_download + except ImportError as error: + raise RuntimeError( + "Automatic BiRefNet download requires the 'huggingface_hub' package. " + "Install the LayerForge requirements or place a compatible checkpoint in " + "ComfyUI/models/background_removal/." + ) from error + + if model is None: + repository = _BIREFNET_REPOSITORY + filename = _BIREFNET_FILENAME + download_dir = _get_birefnet_download_dir() + model_label = "BiRefNet — General" + target_path = os.path.join(download_dir, _BIREFNET_DEFAULT_LOCAL_FILENAME) + else: + repository = model["repo_id"] + filename = model["filename"] + download_dir = os.path.join( + _get_birefnet_download_dir(), + _BIREFNET_REMOTE_DIRECTORY, + model["id"], + ) + os.makedirs(download_dir, exist_ok=True) + model_label = model["label"] + target_path = _get_birefnet_remote_checkpoint_path(model) + + log.info(f"Downloading {model_label} from Hugging Face into {download_dir}...") + try: + downloaded_path = hf_hub_download( + repo_id=repository, + filename=filename, + local_dir=download_dir, + local_dir_use_symlinks=False, + ) + except TypeError: + downloaded_path = hf_hub_download( + repo_id=repository, + filename=filename, + local_dir=download_dir, + ) + + if not _is_native_birefnet_checkpoint(downloaded_path): + raise RuntimeError( + f"Downloaded {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" + ) + + if target_path and os.path.normcase(os.path.abspath(downloaded_path)) != os.path.normcase( + os.path.abspath(target_path) + ): + os.makedirs(os.path.dirname(target_path), exist_ok=True) + try: + os.replace(downloaded_path, target_path) + except OSError: + shutil.copy2(downloaded_path, target_path) + downloaded_path = target_path + + if target_path and not _is_native_birefnet_checkpoint(downloaded_path): + raise RuntimeError( + f"Renamed {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" + ) + + log.info(f"{model_label} checkpoint is ready at {downloaded_path}") + return downloaded_path + + +def _ensure_birefnet_checkpoint(model_path=None): + remote_model = _get_birefnet_remote_model(model_path) + if remote_model: + checkpoint_path = _find_existing_birefnet_remote_checkpoint(remote_model) + if checkpoint_path: + return checkpoint_path + return _download_birefnet_checkpoint(remote_model) + + if not model_path or model_path == "auto": + _find_existing_birefnet_default_checkpoint() + + checkpoint_path = _find_local_birefnet_model(model_path) + if checkpoint_path: + return checkpoint_path + + if model_path and model_path != "auto": + raise RuntimeError("The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.") + + return _download_birefnet_checkpoint() + + +class BiRefNetMatting: + """Adapter around ComfyUI's native BiRefNet loader.""" + + _model_cache: ClassVar[dict] = {} + _model_cache_lock = threading.Lock() + + def __init__(self): + self.model = None + self.model_path = None + + def load_model(self, model_path=None): + loader = _get_comfy_birefnet_loader() + if loader is None: + raise RuntimeError( + "This ComfyUI version does not provide the native BiRefNet background-removal loader." + ) + + checkpoint_path = _ensure_birefnet_checkpoint(model_path) + with self._model_cache_lock: + if checkpoint_path not in self._model_cache: + log.info(f"Loading BiRefNet with ComfyUI's native loader from {checkpoint_path}") + model = loader(checkpoint_path) + if model is None: + raise RuntimeError( + f"ComfyUI did not recognize the BiRefNet checkpoint: {checkpoint_path}" + ) + self._model_cache[checkpoint_path] = model + else: + log.debug(f"Using cached native BiRefNet model from {checkpoint_path}") + + self.model = self._model_cache[checkpoint_path] + self.model_path = checkpoint_path + + return self.model + + def preprocess_image(self, image): + if not isinstance(image, torch.Tensor): + image = transforms.ToTensor()(image).unsqueeze(0) + + if image.dim() == 3: + if image.shape[0] in (1, 3, 4): + image = image.movedim(0, -1).unsqueeze(0) + else: + image = image.unsqueeze(0) + elif image.dim() == 4 and image.shape[1] in (1, 3, 4): + image = image.movedim(1, -1) + + if image.dim() != 4 or image.shape[-1] not in (1, 3, 4): + raise ValueError(f"Expected an image tensor in BCHW or BHWC format, got {tuple(image.shape)}") + + if image.shape[-1] == 1: + image = image.expand(-1, -1, -1, 3) + elif image.shape[-1] == 4: + image = image[..., :3] + + return image.to(dtype=torch.float32).contiguous() + + def execute(self, image, model_path, threshold=0.5, refinement=1, mode="remove_background"): + try: + PromptServer.instance.send_sync("matting_status", {"status": "processing"}) + del refinement + if mode not in {"remove_background", "remove_foreground", "mask_only"}: + raise ValueError(f"Unsupported matting mode: {mode}") + + threshold = float(threshold) + if not 0.0 <= threshold <= 1.0: + raise ValueError("Matting threshold must be between 0 and 1") + + image_tensor = self.preprocess_image(image) + original_size = (image_tensor.shape[1], image_tensor.shape[2]) + log.debug(f"Original size: {original_size}") + self.load_model(model_path) + log.debug(f"Processed image shape: {image_tensor.shape}, dtype: {image_tensor.dtype}") + + with torch.no_grad(): + result = self.model.encode_image(image_tensor) + if result.dim() == 3: + result = result.unsqueeze(1) + elif result.dim() == 2: + result = result.unsqueeze(0).unsqueeze(0) + else: + raise ValueError(f"Unexpected BiRefNet output shape: {tuple(result.shape)}") + + result = result.to(device=image_tensor.device, dtype=torch.float32) + if result.shape[-2:] != original_size: + result = F.interpolate( + result, + size=original_size, + mode="bilinear", + align_corners=False, + ) + result = result.clamp(0.0, 1.0) + log.debug(f"Native BiRefNet output shape: {result.shape}, dtype: {result.dtype}") + + if threshold > 0: + result = (result > threshold).to(dtype=torch.float32) + + alpha_mask = 1.0 - result if mode == "remove_foreground" else result + if mode == "mask_only": + masked_image = alpha_mask.expand(-1, 3, -1, -1) + else: + masked_image = image_tensor.movedim(-1, 1) * alpha_mask + + PromptServer.instance.send_sync("matting_status", {"status": "completed"}) + return masked_image, alpha_mask + except Exception: + PromptServer.instance.send_sync("matting_status", {"status": "error"}) + raise + + @classmethod + def IS_CHANGED(cls, image, model_path, threshold, refinement, mode="remove_background"): + digest = hashlib.md5() + digest.update(str(image).encode()) + digest.update(str(model_path).encode()) + digest.update(str(threshold).encode()) + digest.update(str(refinement).encode()) + digest.update(str(mode).encode()) + return digest.hexdigest() + + +_matting_lock = None + + +async def check_matting_model(request): + """Report whether the selected native BiRefNet checkpoint is ready.""" + try: + if _get_comfy_birefnet_loader() is None: + return web.json_response({ + "available": False, + "reason": "unsupported_comfyui", + "message": "This ComfyUI version does not provide the native BiRefNet background-removal loader.", + }) + + model_options = _get_birefnet_model_options() + requested_model = request.query.get("model_path") or "auto" + remote_model = _get_birefnet_remote_model(requested_model) + + if remote_model: + local_model_path = _find_existing_birefnet_remote_checkpoint(remote_model) + if local_model_path: + log.info(f"Selected BiRefNet model is ready at {local_model_path}") + return web.json_response({ + "available": True, + "reason": "ready", + "message": "Selected model is ready to use", + "model_path": local_model_path, + "selected_model": remote_model["label"], + "models": model_options, + }) + + log.info(f"Selected BiRefNet model is not downloaded: {remote_model['label']}") + return web.json_response({ + "available": False, + "reason": "not_downloaded", + "message": f"{remote_model['label']} will be downloaded automatically on first use.", + "model_path": requested_model, + "selected_model": remote_model["label"], + "models": model_options, + }) + + if requested_model != "auto": + local_model_path = _find_local_birefnet_model(requested_model) + if local_model_path: + log.info(f"Selected BiRefNet model is ready at {local_model_path}") + return web.json_response({ + "available": True, + "reason": "ready", + "message": "Selected model is ready to use", + "model_path": local_model_path, + "selected_model": local_model_path, + "models": model_options, + }) + + return web.json_response({ + "available": False, + "reason": "selected_model_unavailable", + "message": "The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.", + "model_path": requested_model, + "models": model_options, + }) + + local_model_path = _find_local_birefnet_model() + if local_model_path: + log.info(f"BiRefNet model files detected at {local_model_path}") + return web.json_response({ + "available": True, + "reason": "ready", + "message": "Model is ready to use", + "model_path": local_model_path, + "models": model_options, + }) + + searched_paths = _get_birefnet_base_paths() + log.info(f"BiRefNet model not found in any of: {searched_paths}") + return web.json_response({ + "available": False, + "reason": "not_downloaded", + "message": "The BiRefNet checkpoint will be downloaded automatically on first use (requires internet connection).", + "model_path": searched_paths[0] if searched_paths else None, + "models": model_options, + }) + except Exception as error: + log.error(f"Error checking matting model: {error}") + return web.json_response({ + "available": False, + "reason": "error", + "message": f"Error checking model status: {error}", + }, status=500) + + +async def matting(request): + global _matting_lock + + if _matting_lock is not None: + log.warning("Matting already in progress, rejecting request") + return web.json_response({ + "error": "Another matting operation is in progress", + "details": "Please wait for the current operation to complete", + }, status=429) + + _matting_lock = True + try: + log.info("Received matting request") + data = await request.json() + matting_instance = BiRefNetMatting() + image_tensor, original_alpha = convert_base64_to_tensor(data["image"]) + log.debug(f"Input image shape: {image_tensor.shape}") + + mode = data.get("mode", "remove_background") + model_path = data.get("model_path") or "auto" + matted_image, alpha_mask = matting_instance.execute( + image_tensor, + model_path, + threshold=data.get("threshold", 0.5), + refinement=data.get("refinement", 1), + mode=mode, + ) + + if mode == "mask_only": + result_image = convert_tensor_to_base64(alpha_mask) + else: + result_image = convert_tensor_to_base64(matted_image, alpha_mask, original_alpha) + result_mask = convert_tensor_to_base64(alpha_mask) + # Draw Mask uses white pixels as the area to remove, while matting uses + # white pixels as the area to keep. + draw_mask = 1.0 - alpha_mask + + return web.json_response({ + "matted_image": result_image, + "alpha_mask": result_mask, + "draw_mask": convert_tensor_to_base64(draw_mask), + "mode": mode, + "model_path": matting_instance.model_path, + }) + except RuntimeError as error: + log.error(f"Runtime error during matting: {error}") + return web.json_response({ + "error": "Matting Model Error", + "details": str(error), + }, status=500) + except Exception as error: + log.exception(f"Error in matting endpoint: {error}") + error_text = str(error).lower() + if any(marker in error_text for marker in ( + "offline", "connection", "timed out", "huggingface", "localentrynotfound" + )): + return web.json_response({ + "error": "Network Connection Error", + "details": "Failed to download the BiRefNet model from Hugging Face. Please check your internet connection.", + }, status=400) + + return web.json_response({ + "error": "An unexpected error occurred", + "details": traceback.format_exc(), + }, status=500) + finally: + _matting_lock = None + log.debug("Matting lock released") + + +def register_matting_routes(): + """Register matting endpoints without import-time decorators.""" + PromptServer.instance.routes.get("/matting/check-model")(check_matting_model) + PromptServer.instance.routes.post("/matting")(matting) + + +__all__ = [ + "_BIREFNET_DEFAULT_LOCAL_FILENAME", + "_BIREFNET_FILENAME", + "_BIREFNET_MODEL_CATALOG", + "_BIREFNET_PROJECT_URL", + "_BIREFNET_REMOTE_PREFIX", + "BiRefNetMatting", + "_download_birefnet_checkpoint", + "_ensure_birefnet_checkpoint", + "_find_existing_birefnet_default_checkpoint", + "_find_existing_birefnet_remote_checkpoint", + "_find_local_birefnet_model", + "_get_birefnet_model_options", + "_get_comfy_birefnet_loader", + "_is_native_birefnet_checkpoint", + "check_matting_model", + "matting", + "register_matting_routes", +] diff --git a/python/node.py b/python/node.py new file mode 100644 index 0000000..25401ea --- /dev/null +++ b/python/node.py @@ -0,0 +1,447 @@ +"""Core LayerForge ComfyUI node and its in-memory canvas state.""" + +import base64 +import io +import os +import threading +import time +import traceback +import uuid +from typing import ClassVar + +import folder_paths +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +try: + from .log_system import create_module_logger + + log = create_module_logger(__name__) + log.info("Logger initialized for LayerForge node") +except ImportError as error: + print(f"Warning: Logger module not available: {error}") + + class _FallbackLogger: + @staticmethod + def debug(*args, **kwargs): + print("[DEBUG]", *args) + + @staticmethod + def info(*args, **kwargs): + print("[INFO]", *args) + + @staticmethod + def warning(*args, **kwargs): + print("[WARN]", *args) + + @staticmethod + def error(*args, **kwargs): + print("[ERROR]", *args) + + @staticmethod + def exception(*args, **kwargs): + print("[ERROR]", *args) + traceback.print_exc() + + @staticmethod + def fatal(*args, **kwargs): + print("[FATAL]", *args) + + log = _FallbackLogger() + + +torch.set_float32_matmul_precision("high") + + +class LayerForgeNode: + """ComfyUI node that stores and processes the LayerForge canvas.""" + + _canvas_data_storage: ClassVar[dict] = {} + _storage_lock = threading.Lock() + _processing_lock = threading.Lock() + _routes_registered: ClassVar[bool] = False + + _canvas_cache: ClassVar[dict] = { + "image": None, + "mask": None, + "data_flow_status": {}, + "persistent_cache": {}, + "last_execution_id": None, + } + + _websocket_data: ClassVar[dict] = {} + _websocket_listeners: ClassVar[dict] = {} + + def __init__(self): + super().__init__() + self.flow_id = str(uuid.uuid4()) + self.node_id = None + + if self.__class__._canvas_cache["persistent_cache"]: + self.restore_cache() + + def restore_cache(self): + try: + persistent = self.__class__._canvas_cache["persistent_cache"] + current_execution = self.get_execution_id() + + if current_execution != self.__class__._canvas_cache["last_execution_id"]: + log.info(f"New execution detected: {current_execution}") + self.__class__._canvas_cache["image"] = None + self.__class__._canvas_cache["mask"] = None + self.__class__._canvas_cache["last_execution_id"] = current_execution + else: + if persistent.get("image") is not None: + self.__class__._canvas_cache["image"] = persistent["image"] + log.info("Restored image from persistent cache") + if persistent.get("mask") is not None: + self.__class__._canvas_cache["mask"] = persistent["mask"] + log.info("Restored mask from persistent cache") + except Exception as error: + log.error(f"Error restoring cache: {error}") + + def get_execution_id(self): + try: + return str(int(time.time() * 1000)) + except Exception as error: + log.error(f"Error getting execution ID: {error}") + return None + + def update_persistent_cache(self): + try: + self.__class__._canvas_cache["persistent_cache"] = { + "image": self.__class__._canvas_cache["image"], + "mask": self.__class__._canvas_cache["mask"], + } + log.debug("Updated persistent cache") + except Exception as error: + log.error(f"Error updating persistent cache: {error}") + + def track_data_flow(self, stage, status, data_info=None): + flow_status = { + "timestamp": time.time(), + "stage": stage, + "status": status, + "data_info": data_info, + } + log.debug(f"Data Flow [{self.flow_id}] - Stage: {stage}, Status: {status}") + if data_info: + log.debug(f"Data Info: {data_info}") + + self.__class__._canvas_cache["data_flow_status"][self.flow_id] = flow_status + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "fit_on_add": ( + "BOOLEAN", + { + "default": False, + "label_on": "Fit on Add/Paste", + "label_off": "Default Behavior", + }, + ), + "show_preview": ( + "BOOLEAN", + {"default": False, "label_on": "Show Preview", "label_off": "Hide Preview"}, + ), + "auto_refresh_after_generation": ( + "BOOLEAN", + {"default": False, "label_on": "True", "label_off": "False"}, + ), + "trigger": ("INT", {"default": 0, "min": 0, "max": 99999999, "step": 1}), + "node_id": ("STRING", {"default": "0"}), + }, + "optional": { + "input_image": ("IMAGE",), + "input_mask": ("MASK",), + }, + "hidden": { + "prompt": ("PROMPT",), + "unique_id": ("UNIQUE_ID",), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "process_canvas_image" + CATEGORY = "azNodes > LayerForge" + + def add_image_to_canvas(self, input_image): + try: + if not isinstance(input_image, torch.Tensor): + raise ValueError("Input image must be a torch.Tensor") + + if input_image.dim() == 4: + input_image = input_image.squeeze(0) + + if input_image.dim() == 3 and input_image.shape[0] in [1, 3]: + input_image = input_image.permute(1, 2, 0) + + return input_image + except Exception as error: + log.error(f"Error in add_image_to_canvas: {error}") + return None + + def add_mask_to_canvas(self, input_mask, input_image): + try: + if not isinstance(input_mask, torch.Tensor): + raise ValueError("Input mask must be a torch.Tensor") + + if input_mask.dim() == 4: + input_mask = input_mask.squeeze(0) + if input_mask.dim() == 3 and input_mask.shape[0] == 1: + input_mask = input_mask.squeeze(0) + + if input_image is not None: + expected_shape = input_image.shape[:2] + if input_mask.shape != expected_shape: + input_mask = F.interpolate( + input_mask.unsqueeze(0).unsqueeze(0), + size=expected_shape, + mode="bilinear", + align_corners=False, + ).squeeze() + + return input_mask + except Exception as error: + log.error(f"Error in add_mask_to_canvas: {error}") + return None + + def process_canvas_image( + self, + fit_on_add, + show_preview, + auto_refresh_after_generation, + trigger, + node_id, + input_image=None, + input_mask=None, + prompt=None, + unique_id=None, + ): + del show_preview, auto_refresh_after_generation, trigger, prompt + + try: + if not self.__class__._processing_lock.acquire(blocking=False): + log.warning(f"Process already in progress for node {node_id}, skipping...") + return self.get_cached_data() + + log.info( + f"Lock acquired. Starting process_canvas_image for node_id: {node_id} " + f"(fallback unique_id: {unique_id})" + ) + log.info( + f"Storing input data for node {node_id} - " + f"Image: {input_image is not None}, Mask: {input_mask is not None}" + ) + + with self.__class__._storage_lock: + input_data = {} + + if input_image is not None and isinstance(input_image, torch.Tensor): + if input_image.dim() == 3: + input_image = input_image.unsqueeze(0) + + batch_size = input_image.shape[0] + log.info(f"Processing batch of {batch_size} image(s)") + + if batch_size == 1: + img_np = (input_image.squeeze(0).cpu().numpy() * 255).astype(np.uint8) + pil_img = Image.fromarray(img_np, "RGB") + buffered = io.BytesIO() + pil_img.save(buffered, format="PNG") + img_str = base64.b64encode(buffered.getvalue()).decode() + input_data["input_image"] = f"data:image/png;base64,{img_str}" + input_data["input_image_width"] = pil_img.width + input_data["input_image_height"] = pil_img.height + log.debug(f"Stored single input image: {pil_img.width}x{pil_img.height}") + else: + images_array = [] + for index in range(batch_size): + img_np = (input_image[index].cpu().numpy() * 255).astype(np.uint8) + pil_img = Image.fromarray(img_np, "RGB") + buffered = io.BytesIO() + pil_img.save(buffered, format="PNG") + img_str = base64.b64encode(buffered.getvalue()).decode() + images_array.append( + { + "data": f"data:image/png;base64,{img_str}", + "width": pil_img.width, + "height": pil_img.height, + } + ) + log.debug( + f"Stored batch image {index + 1}/{batch_size}: " + f"{pil_img.width}x{pil_img.height}" + ) + + input_data["input_images_batch"] = images_array + log.info(f"Stored batch of {batch_size} images") + + if input_mask is not None and isinstance(input_mask, torch.Tensor): + if input_mask.dim() == 2: + input_mask = input_mask.unsqueeze(0) + if input_mask.dim() == 3 and input_mask.shape[0] == 1: + input_mask = input_mask.squeeze(0) + + mask_np = (input_mask.cpu().numpy() * 255).astype(np.uint8) + pil_mask = Image.fromarray(mask_np, "L") + mask_buffered = io.BytesIO() + pil_mask.save(mask_buffered, format="PNG") + mask_str = base64.b64encode(mask_buffered.getvalue()).decode() + input_data["input_mask"] = f"data:image/png;base64,{mask_str}" + log.debug(f"Stored input mask: {pil_mask.width}x{pil_mask.height}") + + input_data["fit_on_add"] = fit_on_add + self.__class__._canvas_data_storage[f"{node_id}_input"] = input_data + + processed_image = None + processed_mask = None + with self.__class__._storage_lock: + canvas_data = self.__class__._canvas_data_storage.pop(node_id, None) + + if canvas_data: + log.info(f"Canvas data found for node {node_id} from WebSocket") + if canvas_data.get("image"): + image_bytes = base64.b64decode(canvas_data["image"].split(",")[1]) + pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + image_array = np.array(pil_image).astype(np.float32) / 255.0 + processed_image = torch.from_numpy(image_array)[None,] + log.debug(f"Image loaded from WebSocket, shape: {processed_image.shape}") + + if canvas_data.get("mask"): + mask_bytes = base64.b64decode(canvas_data["mask"].split(",")[1]) + pil_mask = Image.open(io.BytesIO(mask_bytes)).convert("L") + mask_array = np.array(pil_mask).astype(np.float32) / 255.0 + processed_mask = torch.from_numpy(mask_array)[None,] + log.debug(f"Mask loaded from WebSocket, shape: {processed_mask.shape}") + else: + log.warning(f"No canvas data found for node {node_id} in WebSocket cache.") + + if processed_image is None: + log.warning("Processed image is still None, creating default blank image.") + processed_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) + if processed_mask is None: + log.warning("Processed mask is still None, creating default blank mask.") + processed_mask = torch.zeros((1, 512, 512), dtype=torch.float32) + + log.debug( + f"About to return output - Image shape: {processed_image.shape}, " + f"Mask shape: {processed_mask.shape}" + ) + self.update_persistent_cache() + log.info("Successfully returning processed image and mask") + return processed_image, processed_mask + except Exception as error: + log.exception(f"Error in process_canvas_image: {error}") + return None, None + finally: + if self.__class__._processing_lock.locked(): + self.__class__._processing_lock.release() + log.debug(f"Process completed for node {node_id}, lock released") + + def get_cached_data(self): + return { + "image": self.__class__._canvas_cache["image"], + "mask": self.__class__._canvas_cache["mask"], + } + + @classmethod + def api_get_data(cls, node_id): + del node_id + try: + return {"success": True, "data": cls._canvas_cache} + except Exception as error: + return {"success": False, "error": str(error)} + + @classmethod + def get_latest_image(cls): + output_dir = folder_paths.get_output_directory() + files = [ + os.path.join(output_dir, filename) + for filename in os.listdir(output_dir) + if os.path.isfile(os.path.join(output_dir, filename)) + ] + image_files = [ + path for path in files if path.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif")) + ] + if not image_files: + return None + return max(image_files, key=os.path.getctime) + + @classmethod + def get_latest_images(cls, since_timestamp=0): + output_dir = folder_paths.get_output_directory() + files = [] + for filename in os.listdir(output_dir): + file_path = os.path.join(output_dir, filename) + if os.path.isfile(file_path) and file_path.lower().endswith( + (".png", ".jpg", ".jpeg", ".bmp", ".gif") + ): + try: + mtime = os.path.getmtime(file_path) + if mtime > since_timestamp: + files.append((mtime, file_path)) + except OSError: + continue + + files.sort(key=lambda item: item[0]) + return [item[1] for item in files] + + @classmethod + def get_flow_status(cls, flow_id=None): + if flow_id: + return cls._canvas_cache["data_flow_status"].get(flow_id) + return cls._canvas_cache["data_flow_status"] + + @classmethod + def _cleanup_old_websocket_data(cls): + """Clean up WebSocket data from invalid nodes or entries older than five minutes.""" + try: + current_time = time.time() + nodes_to_remove = [] + for node_id, data in cls._websocket_data.items(): + if node_id < 0 or current_time - data.get("timestamp", 0) > 300: + nodes_to_remove.append(node_id) + + for node_id in nodes_to_remove: + del cls._websocket_data[node_id] + log.debug(f"Cleaned up old WebSocket entry for node {node_id}") + + if nodes_to_remove: + log.info(f"Cleaned up {len(nodes_to_remove)} old WebSocket entries") + except Exception as error: + log.error(f"Error during WebSocket cleanup: {error}") + + @classmethod + def setup_routes(cls): + """Register the node and matting HTTP routes once.""" + if cls._routes_registered: + return + + from .routes import register_routes + + register_routes(cls) + cls._routes_registered = True + + def store_image(self, image_data): + if isinstance(image_data, str) and image_data.startswith("data:image"): + image_bytes = base64.b64decode(image_data.split(",")[1]) + self.cached_image = Image.open(io.BytesIO(image_bytes)) + else: + self.cached_image = image_data + + def get_cached_image(self): + if self.cached_image: + buffered = io.BytesIO() + self.cached_image.save(buffered, format="PNG") + img_str = base64.b64encode(buffered.getvalue()).decode() + return f"data:image/png;base64,{img_str}" + return None + + +__all__ = ["LayerForgeNode", "log"] diff --git a/python/routes.py b/python/routes.py new file mode 100644 index 0000000..48b2eab --- /dev/null +++ b/python/routes.py @@ -0,0 +1,209 @@ +"""LayerForge HTTP and WebSocket route registration.""" + +import base64 +import io +import os +import time + +from aiohttp import web +from PIL import Image +from server import PromptServer + +from .matting import register_matting_routes +from .node import log + + +def register_routes(node_class): + """Register the legacy LayerForge routes and the matting routes.""" + + @PromptServer.instance.routes.get("/layerforge/canvas_ws") + async def handle_canvas_websocket(request): + ws = web.WebSocketResponse(max_msg_size=33554432) + await ws.prepare(request) + + async for msg in ws: + if msg.type == web.WSMsgType.TEXT: + try: + data = msg.json() + node_id = data.get("nodeId") + if not node_id: + await ws.send_json({"status": "error", "message": "nodeId is required"}) + continue + + with node_class._storage_lock: + node_class._canvas_data_storage[node_id] = { + "image": data.get("image"), + "mask": data.get("mask"), + "timestamp": time.time(), + } + + log.info(f"Received canvas data for node {node_id} via WebSocket") + await ws.send_json({"type": "ack", "nodeId": node_id, "status": "success"}) + log.debug(f"Sent ACK for node {node_id}") + except Exception as error: + log.error(f"Error processing WebSocket message: {error}") + await ws.send_json({"status": "error", "message": str(error)}) + elif msg.type == web.WSMsgType.ERROR: + log.error(f"WebSocket connection closed with exception {ws.exception()}") + + log.info("WebSocket connection closed") + return ws + + @PromptServer.instance.routes.get("/layerforge/get_input_data/{node_id}") + async def get_input_data(request): + try: + node_id = request.match_info["node_id"] + log.debug(f"Checking for input data for node: {node_id}") + with node_class._storage_lock: + input_data = node_class._canvas_data_storage.get(f"{node_id}_input") + + if input_data: + log.info(f"Input data found for node {node_id}, sending to frontend") + return web.json_response({"success": True, "has_input": True, "data": input_data}) + + log.debug(f"No input data found for node {node_id}") + return web.json_response({"success": True, "has_input": False}) + except Exception as error: + log.error(f"Error in get_input_data: {error}") + return web.json_response({"success": False, "error": str(error)}, status=500) + + @PromptServer.instance.routes.post("/layerforge/clear_input_data/{node_id}") + async def clear_input_data(request): + try: + node_id = request.match_info["node_id"] + log.info(f"Clearing input data for node: {node_id}") + with node_class._storage_lock: + input_key = f"{node_id}_input" + if input_key in node_class._canvas_data_storage: + del node_class._canvas_data_storage[input_key] + log.info(f"Input data cleared for node {node_id}") + else: + log.debug(f"No input data to clear for node {node_id}") + + return web.json_response({ + "success": True, + "message": f"Input data cleared for node {node_id}", + }) + except Exception as error: + log.error(f"Error in clear_input_data: {error}") + return web.json_response({"success": False, "error": str(error)}, status=500) + + @PromptServer.instance.routes.get("/ycnode/get_canvas_data/{node_id}") + async def get_canvas_data(request): + del request + try: + cache_data = node_class._canvas_cache + response_data = { + "success": True, + "data": {"image": None, "mask": None}, + } + + if cache_data["image"] is not None: + buffered = io.BytesIO() + cache_data["image"].save(buffered, format="PNG") + img_str = base64.b64encode(buffered.getvalue()).decode() + response_data["data"]["image"] = f"data:image/png;base64,{img_str}" + + if cache_data["mask"] is not None: + mask_buffer = io.BytesIO() + cache_data["mask"].save(mask_buffer, format="PNG") + mask_str = base64.b64encode(mask_buffer.getvalue()).decode() + response_data["data"]["mask"] = f"data:image/png;base64,{mask_str}" + + return web.json_response(response_data) + except Exception as error: + log.error(f"Error in get_canvas_data: {error}") + return web.json_response({"success": False, "error": str(error)}) + + @PromptServer.instance.routes.get("/layerforge/get-latest-images/{since}") + async def get_latest_images_route(request): + try: + since_timestamp = float(request.match_info.get("since", 0)) + latest_image_paths = node_class.get_latest_images(since_timestamp / 1000.0) + images_data = [] + for image_path in latest_image_paths: + with open(image_path, "rb") as image_file: + encoded_string = base64.b64encode(image_file.read()).decode("utf-8") + images_data.append(f"data:image/png;base64,{encoded_string}") + + return web.json_response({"success": True, "images": images_data}) + except Exception as error: + log.error(f"Error in get_latest_images_route: {error}") + return web.json_response({"success": False, "error": str(error)}, status=500) + + @PromptServer.instance.routes.get("/ycnode/get_latest_image") + async def get_latest_image_route(request): + del request + try: + latest_image_path = node_class.get_latest_image() + if latest_image_path: + with open(latest_image_path, "rb") as image_file: + encoded_string = base64.b64encode(image_file.read()).decode() + return web.json_response({ + "success": True, + "image_data": f"data:image/png;base64,{encoded_string}", + }) + + return web.json_response({ + "success": False, + "error": "No images found in output directory.", + }, status=404) + except Exception as error: + return web.json_response({"success": False, "error": str(error)}, status=500) + + @PromptServer.instance.routes.post("/ycnode/load_image_from_path") + async def load_image_from_path_route(request): + try: + data = await request.json() + file_path = data.get("file_path") + if not file_path: + return web.json_response({ + "success": False, + "error": "file_path is required", + }, status=400) + + log.info(f"Attempting to load image from path: {file_path}") + if not os.path.exists(file_path): + log.warning(f"File not found: {file_path}") + return web.json_response({ + "success": False, + "error": f"File not found: {file_path}", + }, status=404) + + valid_extensions = ( + ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff", ".tif", ".ico", ".avif" + ) + if not file_path.lower().endswith(valid_extensions): + return web.json_response({ + "success": False, + "error": f"Invalid image file extension. Supported: {valid_extensions}", + }, status=400) + + try: + with Image.open(file_path) as image: + if image.mode != "RGB": + image = image.convert("RGB") + buffered = io.BytesIO() + image.save(buffered, format="PNG") + img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") + log.info(f"Successfully loaded image from path: {file_path}") + return web.json_response({ + "success": True, + "image_data": f"data:image/png;base64,{img_str}", + "width": image.width, + "height": image.height, + }) + except Exception as image_error: + log.error(f"Error processing image file {file_path}: {image_error}") + return web.json_response({ + "success": False, + "error": f"Error processing image file: {image_error}", + }, status=500) + except Exception as error: + log.error(f"Error in load_image_from_path_route: {error}") + return web.json_response({"success": False, "error": str(error)}, status=500) + + register_matting_routes() + + +__all__ = ["register_routes"] diff --git a/src/CanvasView.ts b/src/CanvasView.ts index be247ce..6485e5d 100644 --- a/src/CanvasView.ts +++ b/src/CanvasView.ts @@ -35,6 +35,8 @@ interface MattingModelOption { path: string; label: string; description?: string; + url?: string; + project_url?: string; source?: 'local' | 'remote'; downloaded?: boolean; } @@ -356,6 +358,55 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): : 'auto'; modelSelect.value = selectedModel; + const modelDetails = document.createElement('div'); + modelDetails.className = 'lf-matting-model-details'; + + const modelDescription = document.createElement('p'); + modelDescription.className = 'lf-matting-model-description'; + + const modelLinks = document.createElement('div'); + modelLinks.className = 'lf-matting-model-links'; + + const createModelLink = (label: string, url: string): HTMLAnchorElement => { + const link = document.createElement('a'); + link.className = 'lf-matting-model-link'; + link.href = url; + link.target = '_blank'; + link.rel = 'noopener noreferrer'; + link.textContent = label; + return link; + }; + + const updateModelDetails = (): void => { + modelLinks.replaceChildren(); + const selectedOption = modelOptions.find((option) => option.path === modelSelect.value); + + if (modelSelect.value === 'auto') { + modelDescription.textContent = 'LayerForge selects the best compatible installed checkpoint. If none is available, the standard General model is downloaded.'; + modelDetails.hidden = false; + return; + } + + if (!selectedOption) { + modelDescription.textContent = 'The selected checkpoint is not currently available.'; + modelDetails.hidden = false; + return; + } + + modelDescription.textContent = selectedOption.description || 'Installed checkpoint validated by ComfyUI\'s native BiRefNet loader.'; + if (selectedOption.url) { + modelLinks.appendChild(createModelLink('Model page', selectedOption.url)); + } + if (selectedOption.project_url) { + modelLinks.appendChild(createModelLink('BiRefNet project', selectedOption.project_url)); + } + modelDetails.hidden = false; + }; + + modelSelect.onchange = updateModelDetails; + modelDetails.append(modelDescription, modelLinks); + updateModelDetails(); + const modeSelect = document.createElement('select'); modeSelect.className = 'lf-matting-settings-select'; (['remove_background', 'remove_foreground', 'mask_only'] as MattingMode[]).forEach((mode) => { @@ -398,6 +449,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): body.append( createRow('Model', modelSelect, 'Only checkpoints accepted by the native ComfyUI BiRefNet loader are listed.'), + modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), modelStatus, diff --git a/src/css/canvas_view.css b/src/css/canvas_view.css index 36d404f..d64b455 100644 --- a/src/css/canvas_view.css +++ b/src/css/canvas_view.css @@ -740,6 +740,41 @@ line-height: 1.4; } +.lf-matting-model-details { + display: flex; + flex-direction: column; + gap: 8px; + margin-top: -7px; + padding: 10px 12px; + background: rgba(255, 255, 255, 0.04); + border: 1px solid #4b4b4b; + border-radius: 4px; +} + +.lf-matting-model-description { + margin: 0; + color: #d2d2d2; + font-size: 12px; + line-height: 1.45; +} + +.lf-matting-model-links { + display: flex; + flex-wrap: wrap; + gap: 12px; +} + +.lf-matting-model-link { + color: #76adff; + font-size: 11px; + text-decoration: none; +} + +.lf-matting-model-link:hover { + color: #a9ccff; + text-decoration: underline; +} + .lf-matting-settings-threshold { display: flex; align-items: center; diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index 2d62745..6433616 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -91,7 +91,9 @@ def _import_layerforge(monkeypatch, tmp_path): return SimpleNamespace( entrypoint=package, - canvas_node=sys.modules[f"{package_name}.canvas_node"], + node=sys.modules[f"{package_name}.python.node"], + image_utils=sys.modules[f"{package_name}.python.image_utils"], + matting=sys.modules[f"{package_name}.python.matting"], routes=routes, package_name=package_name, ) @@ -154,7 +156,7 @@ def test_entrypoint_registers_backend_route_contract(layerforge_runtime): def test_tensor_input_normalization_preserves_comfyui_shapes(layerforge_runtime): import torch - node_class = layerforge_runtime.canvas_node.LayerForgeNode + node_class = layerforge_runtime.node.LayerForgeNode node_class._canvas_cache["persistent_cache"] = {} node = node_class() @@ -171,7 +173,7 @@ def test_base64_image_conversion_preserves_rgb_and_alpha(layerforge_runtime): import torch from PIL import Image - node_module = layerforge_runtime.canvas_node + node_module = layerforge_runtime.image_utils image = Image.new("RGBA", (2, 1), (255, 0, 0, 128)) buffer = io.BytesIO() image.save(buffer, format="PNG") @@ -199,7 +201,7 @@ def test_base64_image_conversion_preserves_rgb_and_alpha(layerforge_runtime): def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, monkeypatch): import torch - node_module = layerforge_runtime.canvas_node + node_module = layerforge_runtime.matting calls = {} class NativeBiRefNet: @@ -236,7 +238,7 @@ def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, m def test_matting_adapter_supports_inverted_and_mask_only_modes(layerforge_runtime, monkeypatch): import torch - node_module = layerforge_runtime.canvas_node + node_module = layerforge_runtime.matting class NativeBiRefNet: def encode_image(self, image): @@ -279,7 +281,7 @@ def test_matting_adapter_supports_inverted_and_mask_only_modes(layerforge_runtim def test_matting_model_options_include_downloadable_official_variants(layerforge_runtime, monkeypatch): - node_module = layerforge_runtime.canvas_node + node_module = layerforge_runtime.matting monkeypatch.setattr(node_module, "_iter_birefnet_checkpoint_paths", lambda: iter(())) @@ -289,11 +291,16 @@ def test_matting_model_options_include_downloadable_official_variants(layerforge assert remote_options assert all(option["path"].startswith("remote:") for option in remote_options) assert all(option["downloaded"] is False for option in remote_options) + assert all(option["description"] for option in remote_options) + assert all(option["url"].startswith("https://huggingface.co/") for option in remote_options) + assert all(option["project_url"] == "https://github.com/ZhengPeng7/BiRefNet" for option in remote_options) assert any(option["path"] == "remote:portrait" for option in remote_options) + portrait = next(option for option in node_module._BIREFNET_MODEL_CATALOG if option["id"] == "portrait") + assert portrait["local_filename"] == "BiRefNet-portrait.safetensors" def test_selected_remote_matting_model_is_sent_to_downloader(layerforge_runtime, monkeypatch): - node_module = layerforge_runtime.canvas_node + node_module = layerforge_runtime.matting selected_model = next( model for model in node_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait" ) @@ -313,8 +320,29 @@ def test_selected_remote_matting_model_is_sent_to_downloader(layerforge_runtime, assert downloaded["model"] is selected_model +def test_legacy_automatic_checkpoint_gets_friendly_filename(layerforge_runtime, monkeypatch, tmp_path): + node_module = layerforge_runtime.matting + model_dir = Path(tmp_path) / "models" / "background_removal" + model_dir.mkdir(parents=True) + legacy_path = model_dir / "model.safetensors" + legacy_path.write_bytes(b"legacy checkpoint") + + monkeypatch.setattr( + node_module, + "_is_native_birefnet_checkpoint", + lambda path: Path(path).resolve() == legacy_path.resolve(), + ) + + result = node_module._find_existing_birefnet_default_checkpoint() + friendly_path = model_dir / "BiRefNet-general.safetensors" + + assert Path(result) == friendly_path + assert friendly_path.exists() + assert not legacy_path.exists() + + def test_empty_execution_returns_comfyui_compatible_fallback_tensors(layerforge_runtime): - node_class = layerforge_runtime.canvas_node.LayerForgeNode + node_class = layerforge_runtime.node.LayerForgeNode node_class._canvas_data_storage.clear() node_class._canvas_cache["persistent_cache"] = {} node = node_class()