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