refactor: split canvas_node into python modules and enhance matting UI
This commit is contained in:
Dariusz L
+6
-1
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
+447
@@ -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"]
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user