diff --git a/__init__.py b/__init__.py index c5bff2e..5f537f7 100644 --- a/__init__.py +++ b/__init__.py @@ -32,6 +32,7 @@ from .nodes.cache_samples import * # Import API endpoints to register routes from . import api from .utility.rgthree_api import routes_model_info +from .nodes import trigger_ref_refresh # Register PrepareRefs backend trigger endpoint NODE_CONFIG = {} diff --git a/api.py b/api.py index 594f526..6e27537 100644 --- a/api.py +++ b/api.py @@ -26,6 +26,17 @@ def get_bg_folder_path(): return str(bg_path) +def get_ref_folder_path(): + """Get the path to the ref folder where reference images are stored""" + import pathlib + current_dir = pathlib.Path(__file__).parent + ref_path = current_dir / "web" / "power_spline_editor" / "ref" + + # Create the directory if it doesn't exist + ref_path.mkdir(parents=True, exist_ok=True) + + return str(ref_path) + @PromptServer.instance.routes.post("/wanvideowrapper_qq/save_ref_image") async def save_ref_image(request): """API endpoint to save ref_image to the bg folder""" @@ -33,15 +44,15 @@ async def save_ref_image(request): post = await request.post() image_data = post.get("image") image_name = post.get("name", "bg_image.png") - + if not image_data: return web.json_response({"error": "No image data provided"}, status=400) - + # Derive filename safely safe_name = os.path.basename(image_name) or "bg_image.png" if "." not in safe_name: safe_name += ".png" - + # Decode base64 image data if image_data.startswith('data:image'): # Remove data URL prefix (e.g., "data:image/png;base64,") @@ -50,10 +61,10 @@ async def save_ref_image(request): else: encoded = image_data image_format = safe_name.split(".")[-1] - + # Decode the base64 string image_bytes = base64.b64decode(encoded) - + # Verify it's a valid image try: img = Image.open(BytesIO(image_bytes)) @@ -61,11 +72,11 @@ async def save_ref_image(request): img = Image.open(BytesIO(image_bytes)) # Reopen after verify except Exception as e: return web.json_response({"error": f"Invalid image data: {str(e)}"}, status=400) - + # Get the bg folder path bg_folder = get_bg_folder_path() ref_image_path = os.path.join(bg_folder, safe_name) - + # Decide format (prefer PNG to preserve alpha) fmt = (image_format or "png").upper() save_kwargs = {} @@ -78,9 +89,9 @@ async def save_ref_image(request): if img.mode not in ("RGBA", "LA"): img = img.convert("RGBA") fmt = "PNG" - + img.save(ref_image_path, format=fmt, **save_kwargs) - + return web.json_response({"success": True, "path": ref_image_path, "message": "Reference image saved successfully"}) except Exception as e: @@ -89,6 +100,86 @@ async def save_ref_image(request): traceback.print_exc() return web.json_response({"error": str(e)}, status=500) +@PromptServer.instance.routes.post("/wanvideowrapper_qq/save_prepare_refs_images") +async def save_prepare_refs_images(request): + """API endpoint to save ALL images from PrepareRefs node to the REF folder""" + try: + data = await request.json() + bg_image_data = data.get("bg_image") + ref_images_data = data.get("ref_images", []) + + print(f"[API save_prepare_refs_images] === Received request ===") + print(f"[API save_prepare_refs_images] bg_image present: {bg_image_data is not None}") + print(f"[API save_prepare_refs_images] ref_images count: {len(ref_images_data) if isinstance(ref_images_data, list) else 0}") + + saved_paths = { + "bg_image_path": None, + "ref_images_paths": [] + } + + # Get ref folder (ALL images go to ref folder for PrepareRefs!) + ref_folder = get_ref_folder_path() + print(f"[API save_prepare_refs_images] Using ref folder: {ref_folder}") + + # Save background image to REF folder (not bg folder!) + if bg_image_data: + try: + print(f"[API save_prepare_refs_images] Saving bg_image to REF folder...") + # Decode base64 image data + if bg_image_data.startswith('data:image'): + header, encoded = bg_image_data.split(',', 1) + else: + encoded = bg_image_data + + image_bytes = base64.b64decode(encoded) + img = Image.open(BytesIO(image_bytes)) + + # Save to REF folder as bg_image.png + bg_image_path = os.path.join(ref_folder, "bg_image.png") + img.save(bg_image_path, format="PNG") + saved_paths["bg_image_path"] = "ref/bg_image.png" + print(f"[API save_prepare_refs_images] ✓ Saved bg_image to {bg_image_path}") + except Exception as e: + print(f"[API save_prepare_refs_images] ERROR saving bg_image: {e}") + import traceback + traceback.print_exc() + + # Save ref images to ref folder + if ref_images_data and isinstance(ref_images_data, list): + for idx, ref_image_data in enumerate(ref_images_data): + try: + print(f"[API save_prepare_refs_images] Saving ref_image {idx} to REF folder...") + # Decode base64 image data + if ref_image_data.startswith('data:image'): + header, encoded = ref_image_data.split(',', 1) + else: + encoded = ref_image_data + + image_bytes = base64.b64decode(encoded) + img = Image.open(BytesIO(image_bytes)) + + # Save to ref folder + ref_image_path = os.path.join(ref_folder, f"ref_image_{idx}.png") + img.save(ref_image_path, format="PNG") + saved_paths["ref_images_paths"].append(f"ref/ref_image_{idx}.png") + print(f"[API save_prepare_refs_images] ✓ Saved ref_image_{idx} to {ref_image_path}") + except Exception as e: + print(f"[API save_prepare_refs_images] ERROR saving ref_image {idx}: {e}") + import traceback + traceback.print_exc() + + print(f"[API save_prepare_refs_images] === COMPLETE ===") + print(f"[API save_prepare_refs_images] Saved bg_image: {saved_paths['bg_image_path']}") + print(f"[API save_prepare_refs_images] Saved {len(saved_paths['ref_images_paths'])} ref images") + + return web.json_response({"success": True, "paths": saved_paths}) + + except Exception as e: + print(f"[API save_prepare_refs_images] ERROR: {e}") + import traceback + traceback.print_exc() + return web.json_response({"error": str(e)}, status=500) + def get_lora_full_path(lora_name): """Get the full path to a LoRA file""" try: diff --git a/nodes/prepare_refs.py b/nodes/prepare_refs.py index 672ec80..b8bc5d4 100644 --- a/nodes/prepare_refs.py +++ b/nodes/prepare_refs.py @@ -7,14 +7,521 @@ import numpy as np from PIL import Image, ImageDraw from scipy.ndimage import gaussian_filter import cv2 +from typing import List, Tuple, Optional, Dict, Any +# ------------------------- +# Module-level constants +# ------------------------- +DEFAULT_BG_CHANNELS = 3 +MASK_FEATHER_SIGMA = 1.5 # gaussian sigma used to feather masks +MASK_THRESHOLD = 0.01 # threshold to consider a mask pixel "active" + + +# ------------------------- +# Utility helpers +# ------------------------- +def ensure_dir(path: Path) -> None: + """Ensure a directory exists.""" + path.mkdir(parents=True, exist_ok=True) + + +def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image: + """ + Convert a torch image tensor to PIL Image. + Accepts HWC or CHW or single-channel (grayscale) tensors. + Expects float tensors in 0..1 or uint8 0..255. + """ + if torch.is_floating_point(image_tensor): + img = torch.clamp(image_tensor, 0.0, 1.0) + img = (img * 255.0).byte() + else: + img = image_tensor.byte() + + # Convert to CHW if HWC provided with channels last + if img.ndim == 3 and img.shape[2] in (1, 3, 4): + img = img.permute(2, 0, 1) + + # torchvision transform expects CHW and byte or float + pil = transforms.ToPILImage()(img) + return pil + + +def pil_to_tensor(img: Image.Image) -> torch.Tensor: + """Convert a PIL Image to a torch float tensor in 0..1 HWC format.""" + tensor = transforms.ToTensor()(img) # CHW float 0..1 + tensor = tensor.permute(1, 2, 0) # HWC + return tensor + + +def clamp_float_tensor(t: torch.Tensor) -> torch.Tensor: + """Clamp float tensor to 0..1.""" + if torch.is_floating_point(t): + return torch.clamp(t, 0.0, 1.0) + return t + + +def safe_device(tensor: torch.Tensor) -> str: + """Return a short string describing tensor's device for debug.""" + return str(tensor.device) if isinstance(tensor, torch.Tensor) else "n/a" + + +# ------------------------- +# Parsing / validation +# ------------------------- +def parse_ref_layer_data_from_prompt(prompt: Optional[Dict], unique_id: Optional[str]) -> List[Dict[str, Any]]: + """ + Extract ref_layer_data from the serialized prompt structure that the front-end provides. + Supports multiple possible serialization shapes (list/dict/value/__value__). + """ + if not prompt or not unique_id or str(unique_id) not in prompt: + return [] + + node_data = prompt[str(unique_id)] + inputs = node_data.get("inputs", {}) if isinstance(node_data, dict) else {} + raw_data = inputs.get("ref_layer_data") + + if raw_data is None: + return [] + + # Normalize different container shapes into a list of dicts + if isinstance(raw_data, list): + return raw_data + if isinstance(raw_data, dict): + # prefer '__value__' then 'value' + if "__value__" in raw_data and isinstance(raw_data["__value__"], list): + return raw_data["__value__"] + if "value" in raw_data: + val = raw_data["value"] + if isinstance(val, list): + return val + if isinstance(val, (tuple,)): + return list(val) + # single value + return [val] if val else [] + # fallback: wrap dict + return [raw_data] + + # fallback: unknown type + return [] + + +def filter_layers_with_shapes(ref_layer_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Validate and filter out layers that do not contain lasso shapes or have no additivePaths. + Only layers that are dicts and contain 'lassoShape' with 'additivePaths' will be returned. + """ + valid_layers = [] + if not isinstance(ref_layer_data, list): + return valid_layers + + for idx, layer in enumerate(ref_layer_data): + if not isinstance(layer, dict): + print(f"[PrepareRefs WARNING] Layer {idx} is not a dict: {type(layer)} - skipping") + continue + + lasso = layer.get("lassoShape") + if not lasso: + print(f"[PrepareRefs WARNING] Layer {idx} ({layer.get('name','unknown')}) missing lassoShape - skipping") + continue + + additive_paths = lasso.get("additivePaths", []) + if not additive_paths: + print(f"[PrepareRefs WARNING] Layer {idx} ({layer.get('name','unknown')}) additivePaths empty - skipping") + continue + + # optional: only include layers explicitly turned "on" if present + if "on" in layer and not layer.get("on", False): + print(f"[PrepareRefs INFO] Layer {idx} ({layer.get('name','unknown')}) 'on' flag is False - skipping") + continue + + valid_layers.append(layer) + + if len(valid_layers) < len(ref_layer_data): + print(f"[PrepareRefs INFO] Filtered layers: {len(ref_layer_data)} → {len(valid_layers)} (removed {len(ref_layer_data)-len(valid_layers)} empty/invalid)") + return valid_layers + + +# ------------------------- +# Mask & image creation +# ------------------------- +def create_mask_from_additive_paths(additive_paths: List[List[Dict[str, float]]], width: int, height: int, + feather_sigma: float = MASK_FEATHER_SIGMA) -> np.ndarray: + """ + Build a feathered mask (numpy float32 HxW) from an array of additive paths. + Each path is a list of {'x': float, 'y': float} where coordinates are normalized 0..1. + """ + mask_img = Image.new("L", (width, height), 0) + draw = ImageDraw.Draw(mask_img) + + for path in additive_paths: + if not path or not isinstance(path, list): + continue + pts = [] + for p in path: + if isinstance(p, dict) and "x" in p and "y" in p: + x = int(p["x"] * width) + y = int(p["y"] * height) + pts.append((x, y)) + if len(pts) >= 3: + draw.polygon(pts, fill=255) + + mask_np = np.array(mask_img).astype(np.float32) / 255.0 + if feather_sigma and feather_sigma > 0: + mask_np = gaussian_filter(mask_np, sigma=feather_sigma) + + mask_np = np.clip(mask_np, 0.0, 1.0) + return mask_np + + +def create_image_and_mask_from_layer(base_image: torch.Tensor, + layer: Dict[str, Any], + width: int, + height: int, + export_alpha: bool) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Given a base image tensor (B,H,W,C), create a per-layer image tensor and mask tensor. + Returns: + image_tensor: HWC float tensor in 0..1 (RGB or RGBA depending on export_alpha) + mask_tensor: HW float tensor in 0..1 (feathered) + """ + try: + lasso_shape = layer.get("lassoShape", {}) + additive_paths = lasso_shape.get("additivePaths", []) + if not additive_paths: + return None, None + + # Construct feathered mask (numpy) and convert to torch + mask_np = create_mask_from_additive_paths(additive_paths, width, height) + mask_tensor = torch.from_numpy(mask_np).float().to(base_image.device) + + # Use the first image in the base batch as source + source_image = base_image[0] # [H,W,C] + + if export_alpha: + # Original RGB channels plus mask as alpha + rgb = source_image[..., :3] + img_tensor = torch.cat((rgb, mask_tensor.unsqueeze(-1)), dim=-1) + else: + img_tensor = source_image[..., :3] + + return img_tensor, mask_tensor + except Exception as e: + print(f"[PrepareRefs ERROR] create_image_and_mask_from_layer failed: {e}") + return None, None + + +def extract_lasso_shapes_as_images(base_image: torch.Tensor, + width: int, + height: int, + ref_layers: List[Dict[str, Any]], + export_alpha: bool) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + For each valid layer in ref_layers, create an image and mask tensor. + Returns stacked tensors: images [N, H, W, C], masks [N, H, W] or (None, None) if none found. + """ + if not ref_layers: + return None, None + + images = [] + masks = [] + for layer in ref_layers: + img, m = create_image_and_mask_from_layer(base_image, layer, width, height, export_alpha) + if img is not None and m is not None: + images.append(img) + masks.append(m) + + if not images: + return None, None + + images_tensor = torch.stack(images, dim=0) + masks_tensor = torch.stack(masks, dim=0) + return images_tensor, masks_tensor + + +# ------------------------- +# Mask combining & preview +# ------------------------- +def combine_masks_union(ref_masks: torch.Tensor) -> torch.Tensor: + """ + Combine multiple mask tensors (N,H,W) into a single HxW mask using max (union). + Result is float32 0..1. + """ + if ref_masks is None or ref_masks.numel() == 0: + return None + return torch.clamp(torch.max(ref_masks, dim=0)[0], 0.0, 1.0) + + +def save_combined_mask_preview(combined_mask: torch.Tensor, save_path: Path) -> None: + """ + Save a preview of the combined mask as an RGB PNG for debugging/previewing. + combined_mask is a float tensor HxW with values 0..1 + """ + try: + pil = tensor_to_pil(combined_mask.cpu()) + pil = pil.convert("RGB") + ensure_dir(save_path.parent) + pil.save(str(save_path)) + except Exception as e: + print(f"[PrepareRefs ERROR] save_combined_mask_preview failed: {e}") + + +def apply_mask_to_bg_with_cv2_preview(combined_mask: torch.Tensor, bg_image_path: Path, out_path: Path) -> None: + """ + If a bg_image exists on disk, apply the combined mask and save a masked bg preview using cv2. + This mirrors the original behavior that produced bg_image_masked.png for frontend debugging. + """ + if not bg_image_path.exists(): + return + + try: + bg_np = cv2.imread(str(bg_image_path)) + if bg_np is None: + return + + mask_np = (combined_mask.cpu().numpy() * 255).astype(np.uint8) + if mask_np.shape != (bg_np.shape[0], bg_np.shape[1]): + mask_np = cv2.resize(mask_np, (bg_np.shape[1], bg_np.shape[0])) + + mask_3c = cv2.cvtColor(mask_np, cv2.COLOR_GRAY2BGR) + masked = cv2.bitwise_and(bg_np, mask_3c) + ensure_dir(out_path.parent) + cv2.imwrite(str(out_path), masked) + except Exception as e: + print(f"[PrepareRefs ERROR] apply_mask_to_bg_with_cv2_preview failed: {e}") + + +# ------------------------- +# Inpainting +# ------------------------- +def inpaint_background_torch(image_tensor: torch.Tensor, mask_tensor: torch.Tensor) -> torch.Tensor: + """ + Inpaint image_tensor using OpenCV inpainting based on mask_tensor. + image_tensor: [1,H,W,C] float 0..1 (RGB) + mask_tensor: HxW float 0..1 + Returns inpainted [1,H,W,C] tensor float 0..1 + """ + if image_tensor is None: + return image_tensor + + # Expect B=1; if not, operate on first + img_np = image_tensor[0].mul(255).byte().cpu().numpy() # HWC uint8 + # ensure image is contiguous with channels last + if img_np.ndim != 3 or img_np.shape[2] not in (3, 4): + # try to permute if CHW provided + try: + img_np = np.transpose(img_np, (1, 2, 0)) + except Exception: + pass + + mask_np = (mask_tensor.cpu().numpy() > MASK_THRESHOLD).astype(np.uint8) * 255 + + try: + img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + inpainted_bgr = cv2.inpaint(img_bgr, mask_np, 3, cv2.INPAINT_TELEA) + inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB) + out_tensor = torch.from_numpy(inpainted_rgb).float().div(255.0).unsqueeze(0).to(image_tensor.device) + return out_tensor + except Exception as e: + print(f"[PrepareRefs ERROR] inpaint_background_torch failed: {e}") + return image_tensor + + +# ------------------------- +# Bounding box / squaring helpers +# ------------------------- +def find_bounding_box(image_tensor: torch.Tensor, mask_tensor: Optional[torch.Tensor] = None) -> Tuple[int, int, int, int]: + """ + Find the bounding box of non-transparent/active pixels. + Returns min_x, min_y, max_x, max_y. + """ + if mask_tensor is not None: + mask = mask_tensor > MASK_THRESHOLD + elif image_tensor.shape[-1] == 4: + mask = image_tensor[..., 3] > MASK_THRESHOLD + else: + rgb_sum = torch.sum(image_tensor[..., :3], dim=-1) + mask = rgb_sum > MASK_THRESHOLD + + if not mask.any(): + return 0, 0, image_tensor.shape[1], image_tensor.shape[0] + + ys, xs = torch.where(mask) + min_x = int(xs.min().item()) + max_x = int(xs.max().item()) + min_y = int(ys.min().item()) + max_y = int(ys.max().item()) + return min_x, min_y, max_x, max_y + + +def create_square_canvas(max_dim: int, channels: int = 4) -> torch.Tensor: + """Create a zero-initialized square canvas tensor HxWxC (float32).""" + return torch.zeros((max_dim, max_dim, channels), dtype=torch.float32) + + +def place_image_in_square(source_image: torch.Tensor, + target_canvas: torch.Tensor, + bbox: Tuple[int, int, int, int], + has_alpha: bool = True) -> torch.Tensor: + """ + Center the visible content from source_image (HWC) inside target_canvas (square HWC). + Handles alpha if present. + """ + min_x, min_y, max_x, max_y = bbox + src_w = max_x - min_x + 1 + src_h = max_y - min_y + 1 + visible = source_image[min_y:max_y + 1, min_x:max_x + 1, :] + + canvas_size = target_canvas.shape[0] + offset_x = (canvas_size - src_w) // 2 + offset_y = (canvas_size - src_h) // 2 + end_x = offset_x + src_w + end_y = offset_y + src_h + + if has_alpha and visible.shape[-1] == 4: + target_canvas[offset_y:end_y, offset_x:end_x, :] = visible + else: + target_canvas[offset_y:end_y, offset_x:end_x, :visible.shape[-1]] = visible + return target_canvas + + +def scale_and_center_in_square(source_image: torch.Tensor, + bbox: Tuple[int, int, int, int], + square_size: int = 768, + has_alpha: bool = True) -> torch.Tensor: + """ + Crop to bounding box, scale to fill as much of square_size as possible (maintaining aspect ratio), + and center in a square canvas. + + Args: + source_image: Input image tensor (H, W, C) + bbox: Bounding box (min_x, min_y, max_x, max_y) + square_size: Target square size (default 768) + has_alpha: Whether image has alpha channel + + Returns: + Square canvas with scaled and centered image + """ + min_x, min_y, max_x, max_y = bbox + src_w = max_x - min_x + 1 + src_h = max_y - min_y + 1 + + # Crop to bounding box + visible = source_image[min_y:max_y + 1, min_x:max_x + 1, :] + + # Calculate scale factor to fill square as much as possible + scale = min(square_size / src_w, square_size / src_h) + + new_w = int(src_w * scale) + new_h = int(src_h * scale) + + # Resize the cropped content using bilinear interpolation + # Convert to CHW format for interpolation + visible_chw = visible.permute(2, 0, 1).unsqueeze(0) # (1, C, H, W) + resized_chw = F.interpolate(visible_chw, size=(new_h, new_w), mode='bilinear', align_corners=False) + resized = resized_chw.squeeze(0).permute(1, 2, 0) # Back to HWC + + # Create square canvas + channels = visible.shape[-1] + canvas = torch.zeros((square_size, square_size, channels), dtype=torch.float32) + + # Center the scaled content + offset_x = (square_size - new_w) // 2 + offset_y = (square_size - new_h) // 2 + + canvas[offset_y:offset_y + new_h, offset_x:offset_x + new_w, :] = resized + + return canvas + + +# ------------------------- +# File saving helpers +# ------------------------- +def save_bg_preview(bg_image: torch.Tensor, out_folder: Path) -> None: + """ + Save bg_image tensor (B,H,W,C) to out_folder/bg_image.png for UI preview. + """ + try: + if bg_image.device != torch.device("cpu"): + bg_image = bg_image.cpu() + + img_tensor = bg_image[0] + # convert to CHW if necessary + if img_tensor.dim() == 3 and img_tensor.shape[0] != 3 and img_tensor.shape[2] == 3: + img_tensor = img_tensor.permute(2, 0, 1) + elif img_tensor.dim() == 2: + img_tensor = img_tensor.unsqueeze(0).repeat(3, 1, 1) + + img_tensor = clamp_float_tensor(img_tensor) + pil = tensor_to_pil(img_tensor) + ensure_dir(out_folder) + pil.save(str(out_folder / "bg_image.png"), format="PNG") + except Exception as e: + print(f"[PrepareRefs ERROR] save_bg_preview failed: {e}") + + +def save_image_to_ref_folder(image: Image.Image, layer_name: str, ref_folder: Path) -> Optional[str]: + """ + Save a PIL image to ref_folder/.png and return its relative path (ref/...). + """ + try: + ensure_dir(ref_folder) + if image.mode not in ("RGB", "RGBA"): + image = image.convert("RGB") + out_path = ref_folder / f"{layer_name}.png" + image.save(str(out_path), format="PNG") + return f"ref/{layer_name}.png" + except Exception as e: + print(f"[PrepareRefs ERROR] save_image_to_ref_folder failed for {layer_name}: {e}") + return None + + +def save_masks_to_folder(masks: torch.Tensor, ref_layer_data: List[Dict[str, Any]], ref_folder: Path) -> None: + """ + Save mask tensors (N,H,W) to /_mask.png + """ + try: + ensure_dir(ref_folder) + if masks.device != torch.device("cpu"): + masks_cpu = masks.cpu() + else: + masks_cpu = masks + + for idx in range(masks_cpu.shape[0]): + if idx < len(ref_layer_data): + layer_name = ref_layer_data[idx].get("name", f"ref_{idx + 1}") + else: + layer_name = f"ref_{idx + 1}" + + mask_tensor = masks_cpu[idx] + if mask_tensor.ndim == 2: + mask_np = (mask_tensor.numpy() * 255).astype("uint8") + elif mask_tensor.ndim == 3 and mask_tensor.shape[2] == 1: + mask_np = (mask_tensor[:, :, 0].numpy() * 255).astype("uint8") + else: + # Unexpected shape: try to pick first channel + try: + mask_np = (mask_tensor[0, :, :].numpy() * 255).astype("uint8") + except Exception: + print(f"[PrepareRefs WARNING] Unexpected mask dims for {layer_name}: {mask_tensor.shape}") + continue + + pil = Image.fromarray(mask_np, mode="L") + pil.save(str(ref_folder / f"{layer_name}_mask.png"), format="PNG") + except Exception as e: + print(f"[PrepareRefs ERROR] save_masks_to_folder failed: {e}") + + +# ------------------------- +# Main class (PrepareRefs) +# ------------------------- class PrepareRefs: """ Canvas-enabled node for preparing background/ref images with drawn shapes. Exports ref images with lasso shapes as masks. """ + # Node metadata (kept as original) @classmethod def INPUT_TYPES(cls): return { @@ -24,8 +531,6 @@ class PrepareRefs: }, "optional": { "bg_image": ("IMAGE", {"forceInput": True}), - "export_alpha": ("BOOLEAN", {"default": True}), - "to_bounding_box": ("BOOLEAN", {"default": True}), }, "hidden": { "unique_id": "UNIQUE_ID", @@ -42,496 +547,278 @@ class PrepareRefs: Can export ref images with lasso shapes as masks. """ - def prepare(self, mask_width, mask_height, bg_image=None, export_alpha=True, to_bounding_box=True, unique_id=None, prompt=None): - # Determine the final output width and height - # If bg_image is provided, use its dimensions. Otherwise, use mask_width/height as fallback. - final_width = int(mask_width) - final_height = int(mask_height) + def prepare(self, mask_width: int, mask_height: int, bg_image: Optional[torch.Tensor] = None, + unique_id: Optional[str] = None, prompt: Optional[Dict] = None): + """ + Main entry point. High-level orchestration: + - Resolve base dimensions and images + - Parse & validate ref layer data + - Short-circuit if no valid layers + - Create per-layer images and masks + - Create combined mask preview and inpaint background + - Crop to 768x768 bounding boxes and scale images + - Save preview/ref images/masks to disk and return UI metadata + tensors + + Note: export_alpha and to_bounding_box are always True now (hardcoded). + """ + + # Hardcoded settings: always export alpha and use 768x768 bounding boxes + export_alpha = True + to_bounding_box = True + + # Resolve final dims and base images + width, height, output_bg_image, internal_processing_image = self._resolve_base_images(mask_width, mask_height, bg_image) + + # Parse and validate ref_layer_data early + raw_ref_data = parse_ref_layer_data_from_prompt(prompt, unique_id) + valid_ref_layers = filter_layers_with_shapes(raw_ref_data) + + # EARLY GUARD: if no valid layers, warn and return empty outputs (no further processing) + if not valid_ref_layers: + # Build empty outputs matching expected shapes + empty_ref_images = torch.zeros((1, height, width, 4 if export_alpha else 3), dtype=torch.float32) + empty_ref_masks = torch.zeros((1, height, width), dtype=torch.float32) + ui_out = {"bg_image_dims": [{"width": float(width), "height": float(height)}]} + # Persist bg preview only if original bg_image was provided + if bg_image is not None: + try: + save_bg_preview(bg_image, Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg") + ui_out["bg_image_path"] = ["bg/bg_image.png"] + except Exception as exc: + print(f"[PrepareRefs WARNING] Failed to save bg preview: {exc}") + return {"ui": ui_out, "result": (output_bg_image, empty_ref_images, empty_ref_masks, [])} + + # Extract per-layer images and masks (full-size) + ref_images, ref_masks = extract_lasso_shapes_as_images(internal_processing_image, width, height, valid_ref_layers, export_alpha) + + # Ensure we always have tensors for downstream logic + if ref_images is None or ref_images.shape[0] == 0: + ref_images = torch.zeros((1, height, width, 4 if export_alpha else 3), dtype=torch.float32) + ref_masks = torch.zeros((1, height, width), dtype=torch.float32) + + # Optionally save a combined mask preview and create masked BG preview + if ref_masks is not None and ref_masks.shape[0] > 0 and to_bounding_box: + combined_mask = combine_masks_union(ref_masks) + preview_mask_path = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" / "combined_ref_mask.png" + save_combined_mask_preview(combined_mask, preview_mask_path) + + bg_image_path = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" / "bg_image.png" + masked_bg_out = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" / "bg_image_masked.png" + apply_mask_to_bg_with_cv2_preview(combined_mask, bg_image_path, masked_bg_out) + + # Combine all masks at original dims for inpainting BEFORE any cropping + if ref_masks is not None and ref_masks.shape[0] > 0: + combined_original_dims_mask = combine_masks_union(ref_masks) + else: + combined_original_dims_mask = torch.zeros((height, width), dtype=torch.float32, device=output_bg_image.device) + + # Apply inpainting if combined mask has active areas + if torch.any(combined_original_dims_mask > MASK_THRESHOLD): + output_bg_image = inpaint_background_torch(output_bg_image, combined_original_dims_mask) + + # Determine bounding boxes and optionally crop to square images + # Always use 768x768 squares and scale each layer to fill as much space as possible + FIXED_SQUARE_SIZE = 768 + max_dim = max(width, height) + + if ref_images.shape[0] > 0 and to_bounding_box: + all_bboxes = [] + for i in range(ref_images.shape[0]): + bbox = find_bounding_box(ref_images[i], mask_tensor=ref_masks[i] if ref_masks is not None else None) + all_bboxes.append(bbox) + + # Build 768x768 square canvases with scaled and centered images + masks + square_images = [] + square_masks = [] + + for i in range(ref_images.shape[0]): + img = ref_images[i] + mask = ref_masks[i] if ref_masks is not None else torch.zeros((height, width), dtype=torch.float32) + bbox = all_bboxes[i] + + # Scale and center image in 768x768 canvas + sq_img = scale_and_center_in_square(img, bbox, square_size=FIXED_SQUARE_SIZE, has_alpha=export_alpha) + square_images.append(sq_img) + + # Scale and center mask in 768x768 canvas + min_x, min_y, max_x, max_y = bbox + src_w = max_x - min_x + 1 + src_h = max_y - min_y + 1 + visible_mask = mask[min_y:max_y + 1, min_x:max_x + 1] + + # Calculate same scale factor as for image + scale = min(FIXED_SQUARE_SIZE / src_w, FIXED_SQUARE_SIZE / src_h) + new_w = int(src_w * scale) + new_h = int(src_h * scale) + + # Resize mask + mask_chw = visible_mask.unsqueeze(0).unsqueeze(0) # (1, 1, H, W) + resized_mask_chw = F.interpolate(mask_chw, size=(new_h, new_w), mode='bilinear', align_corners=False) + resized_mask = resized_mask_chw.squeeze(0).squeeze(0) # (H, W) + + # Create square mask canvas and center + sq_mask = torch.zeros((FIXED_SQUARE_SIZE, FIXED_SQUARE_SIZE), dtype=torch.float32) + offset_x = (FIXED_SQUARE_SIZE - new_w) // 2 + offset_y = (FIXED_SQUARE_SIZE - new_h) // 2 + sq_mask[offset_y:offset_y + new_h, offset_x:offset_x + new_w] = resized_mask + + square_masks.append(sq_mask) + + # Replace with square versions + ref_images = torch.stack(square_images, dim=0) + ref_masks = torch.stack(square_masks, dim=0) + max_dim = FIXED_SQUARE_SIZE # Update max_dim for subsequent operations + + # Align batch sizes: if base batch size greater than refs, pad refs + ref_batch_size = ref_images.shape[0] + base_batch_size = internal_processing_image.shape[0] + if base_batch_size > ref_batch_size: + extra_count = base_batch_size - ref_batch_size + extra_channels = ref_images.shape[-1] + extra_images = torch.zeros((extra_count, max_dim, max_dim, extra_channels), dtype=ref_images.dtype, device=ref_images.device) + ref_images = torch.cat([ref_images, extra_images], dim=0) + + extra_masks = torch.zeros((extra_count, max_dim, max_dim), dtype=ref_masks.dtype, device=ref_masks.device) + ref_masks = torch.cat([ref_masks, extra_masks], dim=0) + + # Build ui_out + ui_out = {"bg_image_dims": [{"width": float(width), "height": float(height)}]} + + # Save bg preview if bg_image provided + if bg_image is not None: + try: + save_bg_preview(bg_image, Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg") + ui_out["bg_image_path"] = ["bg/bg_image.png"] + except Exception as exc: + print(f"[PrepareRefs WARNING] Failed to save bg preview: {exc}") + + # Save input-like ref images (converted to PIL) and record ui paths + ref_paths = [] + try: + # ensure CPU for PIL conversion + if ref_images.device != torch.device("cpu"): + ref_images_cpu = ref_images.cpu() + else: + ref_images_cpu = ref_images + + transform_to_pil = transforms.ToPILImage() + ref_folder = Path(__file__).parent.parent / "web" / "power_spline_editor" / "ref" + ensure_dir(ref_folder) + + for idx in range(ref_images_cpu.shape[0]): + if idx < len(valid_ref_layers): + layer_name = valid_ref_layers[idx].get("name", f"ref_{idx + 1}") + else: + layer_name = f"ref_{idx + 1}" + + img_tensor = ref_images_cpu[idx] + + # Normalize and permute CHW/HWC as needed for torchvision + if img_tensor.dim() == 3: + # If channels last (HWC) and last dim is channels, torchvision expects CHW + if img_tensor.shape[0] not in (3, 4) and img_tensor.shape[2] in (3, 4): + img_tensor = img_tensor.permute(2, 0, 1) + elif img_tensor.shape[0] in (3, 4) and img_tensor.shape[2] not in (3, 4): + # already CHW + pass + elif img_tensor.dim() == 2: + img_tensor = img_tensor.unsqueeze(0).repeat(3, 1, 1) + + img_tensor = clamp_float_tensor(img_tensor) + pil_img = transform_to_pil(img_tensor) + rel_path = save_image_to_ref_folder(pil_img, layer_name, ref_folder) + if rel_path: + ref_paths.append(rel_path) + + if ref_paths: + ui_out["ref_images_paths"] = ref_paths + except Exception as e: + print(f"[PrepareRefs ERROR] saving input ref images failed: {e}") + + # Save output images (clean bg_image) and ref images/masks per-layer + try: + # Save output bg_image_cl.png + if output_bg_image is not None: + out_ref_folder = Path(__file__).parent.parent / "web" / "power_spline_editor" / "ref" + ensure_dir(out_ref_folder) + # prepare a PIL image + bg_img = output_bg_image + if bg_img.device != torch.device("cpu"): + bg_img = bg_img.cpu() + transform_to_pil = transforms.ToPILImage() + img_tensor = bg_img[0] + if img_tensor.dim() == 3 and img_tensor.shape[2] == 3: + img_tensor = img_tensor.permute(2, 0, 1) + img_tensor = clamp_float_tensor(img_tensor) + pil = transform_to_pil(img_tensor) + if pil.mode not in ("RGB", "RGBA"): + pil = pil.convert("RGB") + pil.save(str(out_ref_folder / "bg_image_cl.png"), format="PNG") + except Exception as e: + print(f"[PrepareRefs ERROR] failed to save bg_image_cl: {e}") + + # Save final per-layer ref images and masks (named by layer) + try: + ref_folder = Path(__file__).parent.parent / "web" / "power_spline_editor" / "ref" + ensure_dir(ref_folder) + + # Save ref_images outputs with layer names + if ref_images.device != torch.device("cpu"): + ref_imgs_cpu = ref_images.cpu() + else: + ref_imgs_cpu = ref_images + + transform_to_pil = transforms.ToPILImage() + for idx in range(ref_imgs_cpu.shape[0]): + if idx < len(valid_ref_layers): + layer_name = valid_ref_layers[idx].get("name", f"ref_{idx + 1}") + else: + layer_name = f"ref_{idx + 1}" + + img_tensor = ref_imgs_cpu[idx] + # ensure CHW for ToPILImage + if img_tensor.dim() == 3 and img_tensor.shape[2] == 3: + img_tensor = img_tensor.permute(2, 0, 1) + + img_tensor = clamp_float_tensor(img_tensor) + pil_img = transform_to_pil(img_tensor) + if pil_img.mode not in ("RGB", "RGBA"): + pil_img = pil_img.convert("RGB") + pil_img.save(str(ref_folder / f"{layer_name}.png"), format="PNG") + except Exception as e: + print(f"[PrepareRefs ERROR] failed to save final ref images: {e}") + + # Save masks + try: + ref_folder = Path(__file__).parent.parent / "web" / "power_spline_editor" / "ref" + save_masks_to_folder(ref_masks, valid_ref_layers, ref_folder) + except Exception as e: + print(f"[PrepareRefs ERROR] failed to save ref masks: {e}") + + return {"ui": ui_out, "result": (output_bg_image, ref_images, ref_masks, valid_ref_layers)} + + # ------------------------- + # Internal helper + # ------------------------- + def _resolve_base_images(self, mask_width: int, mask_height: int, bg_image: Optional[torch.Tensor]): + """ + Determine final width/height and internal images used during processing. + Returns (width, height, output_bg_image, internal_processing_image) + """ + final_w = int(mask_width) + final_h = int(mask_height) output_bg_image = None internal_processing_image = None if bg_image is not None: - final_height = bg_image.shape[1] - final_width = bg_image.shape[2] + # bg_image expected shape: [B,H,W,C] or [1,H,W,C] + final_h = bg_image.shape[1] + final_w = bg_image.shape[2] output_bg_image = bg_image internal_processing_image = bg_image else: - output_bg_image = torch.zeros((1, final_height, final_width, 3), dtype=torch.float32) + output_bg_image = torch.zeros((1, final_h, final_w, DEFAULT_BG_CHANNELS), dtype=torch.float32) internal_processing_image = output_bg_image - # Use final_width and final_height consistently - width = final_width - height = final_height - - ref_layer_data = None - if prompt and unique_id and str(unique_id) in prompt: - node_data = prompt[str(unique_id)] - inputs = node_data.get("inputs", {}) - # The ref_layer_data should be in the inputs from the serialized widget - if "ref_layer_data" in inputs: - raw_data = inputs["ref_layer_data"] - - # Handle the case where ref_layer_data might be a dict or list - if isinstance(raw_data, list): - ref_layer_data = raw_data - elif isinstance(raw_data, dict): - # If it's a dict, it might contain the data in a nested structure - # Try to extract the actual ref layer data from the dict - if '__value__' in raw_data and isinstance(raw_data['__value__'], list): - ref_layer_data = raw_data['__value__'] - elif 'value' in raw_data and isinstance(raw_data['value'], list): - ref_layer_data = raw_data['value'] - elif 'value' in raw_data: - # If value isn't a list, check if it's the data itself - if isinstance(raw_data['value'], (list, tuple)): - ref_layer_data = list(raw_data['value']) - else: - ref_layer_data = [raw_data['value']] if raw_data['value'] else [] - else: - # If not in 'value' or '__value__' key, check if the dict itself contains the list - # This might happen if the widget serialization works differently - ref_layer_data = [raw_data] # Wrap in list if it's a single dict item - - # Create ref images and masks from lasso shapes - ref_images, ref_masks = self._extract_lasso_shapes_as_images( - internal_processing_image, width, height, ref_layer_data, export_alpha, unique_id, prompt) - - # Combined masks - if ref_masks is not None and ref_masks.shape[0] > 0 and to_bounding_box: - self._process_combined_mask_before_crop(ref_masks, output_bg_image) - - # If no ref images from lasso shapes, use empty tensor - if ref_images is None or ref_images.shape[0] == 0: - # Create a single empty image/mask if no shapes exist - if export_alpha: - ref_images = torch.zeros((1, height, width, 4), dtype=torch.float32) - else: - ref_images = torch.zeros((1, height, width, 3), dtype=torch.float32) - ref_masks = torch.zeros((1, height, width), dtype=torch.float32) - - # --- NEW LOGIC START (Mask Combining and Inpainting at Original Dims) --- - # Combine all reference masks (at original dimensions) into a single mask for inpainting the background - combined_original_dims_mask = None - if ref_masks.shape[0] > 0: - combined_original_dims_mask = torch.max(ref_masks, dim=0)[0] - else: - combined_original_dims_mask = torch.zeros((height, width), dtype=torch.float32, device=output_bg_image.device) - - # Store the original combined mask for later use after cropping - original_combined_mask = combined_original_dims_mask.clone() if combined_original_dims_mask is not None else None - - # Apply inpainting to the output_bg_image if a combined mask exists and has active areas - if torch.any(combined_original_dims_mask > 0.01): - output_bg_image = self._inpaint_background_torch(output_bg_image, combined_original_dims_mask) - # --- NEW LOGIC END --- - - # Initialize max_dim for potential use in padding later, ensuring it's always defined. - # This will be overridden if to_bounding_box is true and ref_images are present. - max_dim = max(width, height) - - # Now proceed with to_bounding_box logic for ref_images and ref_masks outputs - if ref_images.shape[0] > 0 and to_bounding_box: # Only apply if there are actual ref_images and to_bounding_box is true - # Find the largest bounding box among all images using their masks - max_dim = 0 # Reset to calculate based on bounding boxes - all_bboxes = [] - - for i in range(ref_images.shape[0]): - image_tensor = ref_images[i] - mask_tensor = ref_masks[i] - bbox = self._find_bounding_box(image_tensor, mask_tensor=mask_tensor) - all_bboxes.append(bbox) - - # Calculate the dimensions of this bounding box - bbox_width = bbox[2] - bbox[0] + 1 - bbox_height = bbox[3] - bbox[1] + 1 - max_dim = max(max_dim, bbox_width, bbox_height) - - - # Create square images for each reference image - square_images = [] - square_masks = [] - - for i in range(ref_images.shape[0]): - image_tensor = ref_images[i] - mask_tensor = ref_masks[i] - bbox = all_bboxes[i] - - # Create a square canvas - channels = 4 if export_alpha else 3 - square_canvas = self._create_square_canvas(max_dim, channels) - - # Place the image content in the square canvas - square_image = self._place_image_in_square(image_tensor, square_canvas, bbox, has_alpha=export_alpha) - square_images.append(square_image) - - # Create a corresponding mask for the square image - square_mask = torch.zeros((max_dim, max_dim), dtype=torch.float32) - - # Calculate the position where the image was placed - bbox_width = bbox[2] - bbox[0] + 1 - bbox_height = bbox[3] - bbox[1] + 1 - offset_x = (max_dim - bbox_width) // 2 - offset_y = (max_dim - bbox_height) // 2 - end_x = offset_x + bbox_width - end_y = offset_y + bbox_height - - # Copy the mask content to the square mask - # Extract only the visible part of the mask based on the bounding box - visible_mask = mask_tensor[bbox[1]:bbox[3]+1, bbox[0]:bbox[2]+1] - square_mask[offset_y:end_y, offset_x:end_x] = visible_mask - square_masks.append(square_mask) - - # Replace the original ref_images and ref_masks with the square versions - ref_images = torch.stack(square_images, dim=0) - ref_masks = torch.stack(square_masks, dim=0) - - # The batch size adjustment logic follows, unchanged. - # This part still ensures ref_images/ref_masks align with internal_processing_image's batch size - # (which is 1) for downstream compatibility. - ref_batch_size = ref_images.shape[0] - base_batch_size = internal_processing_image.shape[0] # Use internal_processing_image - - # We explicitly do NOT pad internal_processing_image based on ref_batch_size, - # as the user wants the output bg_image to be singular. - # If internal_processing_image needs to be batched for other purposes, - # it should be handled internally or by subsequent nodes. - - if base_batch_size > ref_batch_size: - # If there are more internal_processing_image items than ref_images, - # pad the ref_images batch with empty images. - extra_count = base_batch_size - ref_batch_size - extra_channels = ref_images.shape[-1] - extra_images = torch.zeros((extra_count, max_dim, max_dim, extra_channels), - dtype=ref_images.dtype, device=ref_images.device) - ref_images = torch.cat([ref_images, extra_images], dim=0) - - extra_masks = torch.zeros((extra_count, max_dim, max_dim), - dtype=ref_masks.dtype, device=ref_masks.device) - ref_masks = torch.cat([ref_masks, extra_masks], dim=0) - - # Note: We don't need to recombine masks here since we already applied inpainting before cropping - # The masks have been cropped to bounding boxes, so they're no longer at the original dimensions - # We'll use the original_combined_mask that was saved before cropping - - ui_out = { - "bg_image_dims": [{"width": float(width), "height": float(height)}], - } - - # Persist preview to disk so the frontend can pull it without bloating UI payloads - if bg_image is not None: - try: - self._save_bg_preview(bg_image) - ui_out["bg_image_path"] = ["bg/bg_image.png"] - except Exception as exc: # pragma: no cover - UI helper - print(f"PrepareRefs: failed to save bg preview: {exc}") - - # Return the ref layer data as well for the export node - # Use our extracted ref_layer_data which is already properly formatted - - return {"ui": ui_out, "result": (output_bg_image, ref_images, ref_masks, ref_layer_data or [])} - - def _to_image_tensor(self, image, width, height): - if image is not None: - # Ensure image has proper dimensions - if image.shape[1] != height or image.shape[2] != width: - # Resize image to match specified dimensions - img_bchw = image.permute(0, 3, 1, 2) # [B, C, H, W] - resized_bchw = F.interpolate(img_bchw, size=(height, width), mode='bilinear', align_corners=False) - image = resized_bchw.permute(0, 2, 3, 1) # [B, H, W, C] - return image - return torch.zeros((1, height, width, 3), dtype=torch.float32) - - def _process_combined_mask_before_crop(self, ref_masks, bg_image=None): - """ - Called right after masks are created, but BEFORE bounding-box cropping/squaring. - At this point all masks are still full-size (width x height) and perfectly aligned. - - Use this for: - - Previewing the union of all refs - - Running extra logic on combined mask - - Saving a "total reference mask" for debugging - - Applying the combined mask to bg_image using cv2 - """ - if ref_masks.shape[0] == 0: - return - - # Combine all masks (union) - combined_mask = torch.clamp(torch.sum(ref_masks, dim=0), 0, 1) # or torch.max() - - # Example: Save preview of combined mask - mask_pil = transforms.ToPILImage()(combined_mask.cpu()) - mask_pil = mask_pil.convert("RGB") - save_path = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" / "combined_ref_mask.png" - save_path.parent.mkdir(parents=True, exist_ok=True) - mask_pil.save(save_path) - - # Apply the combined mask to bg_image using cv2 - # Get the bg_image path - bg_image_path = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" / "bg_image.png" - - # Check if bg_image exists - if bg_image_path.exists(): - # Load the bg_image - bg_image_cv2 = cv2.imread(str(bg_image_path)) - if bg_image_cv2 is not None: - # Convert mask to numpy for cv2 - mask_np = (combined_mask.cpu().numpy() * 255).astype(np.uint8) - - # Resize mask if it doesn't match bg_image dimensions - if mask_np.shape != (bg_image_cv2.shape[0], bg_image_cv2.shape[1]): - mask_np = cv2.resize(mask_np, (bg_image_cv2.shape[1], bg_image_cv2.shape[0])) - - # Apply the mask to create a masked version of bg_image - # Convert mask to 3-channel for proper masking - mask_3ch = cv2.cvtColor(mask_np, cv2.COLOR_GRAY2BGR) - masked_bg_image = cv2.bitwise_and(bg_image_cv2, mask_3ch) - - # Save the masked bg_image - masked_bg_path = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" / "bg_image_masked.png" - cv2.imwrite(str(masked_bg_path), masked_bg_image) - - # Return it if downstream nodes want it - return combined_mask - - def _extract_lasso_shapes_as_images(self, base_image, width, height, ref_layer_data, export_alpha, unique_id=None, prompt=None): - """ - Extract lasso shapes from the frontend and convert them to image tensors with masks. - This method processes the serialized ref layer data containing the lasso shapes. - """ - # Process the ref layer data to create images with masks - if not ref_layer_data or not isinstance(ref_layer_data, list): - return None, None - - # Filter out layers that don't have shapes (empty additivePaths) - layers_with_shapes = [] - for layer_data in ref_layer_data: - if isinstance(layer_data, dict) and layer_data.get('on', False): - lasso_shape = layer_data.get('lassoShape', {}) - additive_paths = lasso_shape.get('additivePaths', []) - if additive_paths and len(additive_paths) > 0: - # This layer has shapes, add it to our list - layers_with_shapes.append(layer_data) - - if not layers_with_shapes: - # No layers with shapes found - return None, None - - # Create images and masks for each layer with shapes - images = [] - masks = [] - - for layer_data in layers_with_shapes: - # Create an image with the shape filled in - image, mask = self._create_image_from_lasso_shape(base_image, layer_data, width, height, export_alpha) - if image is not None and mask is not None: - images.append(image) - masks.append(mask) - - if not images: - return None, None - - # Stack all images and masks into tensors - images_tensor = torch.stack(images, dim=0) - masks_tensor = torch.stack(masks, dim=0) - - return images_tensor, masks_tensor - - def _create_image_from_lasso_shape(self, base_image, layer_data, width, height, export_alpha): - """ - Create an image and mask from a single lasso shape. - This version extracts the content from the base_image using the mask. - """ - try: - lasso_shape = layer_data.get('lassoShape', {}) - additive_paths = lasso_shape.get('additivePaths', []) - - if not additive_paths: - return None, None - - # Create a PIL image and mask - mask_img = Image.new('L', (width, height), 0) - mask_draw = ImageDraw.Draw(mask_img) - - # Draw all paths on the mask - for path in additive_paths: - if not path or not isinstance(path, list): - continue - actual_points = [] - for point in path: - if isinstance(point, dict) and 'x' in point and 'y' in point: - x = int(point['x'] * width) - y = int(point['y'] * height) - actual_points.append((x, y)) - if len(actual_points) >= 3: - mask_draw.polygon(actual_points, fill=255) - - # Apply 3-pixel feather to the mask using a Gaussian blur - mask_np = np.array(mask_img).astype(np.float32) / 255.0 - mask_np = gaussian_filter(mask_np, sigma=1.5) - - mask_tensor = torch.from_numpy(mask_np).to(base_image.device) - - # Use the first image from the base_image batch as the source - source_image = base_image[0] # Shape: [H, W, C] - - if export_alpha: - # Take the original RGB channels of the source_image - original_rgb = source_image[..., :3] - # Concatenate the original RGB with the feathered mask as the alpha channel - image_tensor = torch.cat((original_rgb, mask_tensor.unsqueeze(-1)), dim=-1) - else: - # If not exporting alpha, just return the original RGB of the source within the mask area. - # The feathering in mask_tensor is still present, but won't be explicitly used as alpha - # in the output image_tensor if it's RGB. - image_tensor = source_image[..., :3] - - return image_tensor, mask_tensor - except Exception as e: - print(f"Error creating image from lasso shape: {e}") - return None, None - - def _save_bg_preview(self, bg_image): - if bg_image.device != torch.device("cpu"): - bg_image = bg_image.cpu() - - transform = transforms.ToPILImage() - img_tensor = bg_image[0] - if img_tensor.dim() == 3 and img_tensor.shape[0] != 3 and img_tensor.shape[2] == 3: - img_tensor = img_tensor.permute(2, 0, 1) - elif img_tensor.dim() == 2: - img_tensor = img_tensor.unsqueeze(0).repeat(3, 1, 1) - - if torch.is_floating_point(img_tensor): - img_tensor = torch.clamp(img_tensor, 0, 1) - - image = transform(img_tensor) - if image.mode != "RGB": - image = image.convert("RGB") - - bg_folder = Path(__file__).parent.parent / "web" / "power_spline_editor" / "bg" - bg_folder.mkdir(parents=True, exist_ok=True) - bg_path = bg_folder / "bg_image.png" - image.save(str(bg_path), format="PNG") - - def _inpaint_background_torch(self, image_tensor: torch.Tensor, mask_tensor: torch.Tensor) -> torch.Tensor: - """ - Applies OpenCV's inpainting to a torch tensor image using a torch tensor mask. - image_tensor: [B, H, W, C] (float32, 0-1) - mask_tensor: [H, W] (float32, 0-1), where 1 indicates masked areas - """ - if image_tensor.shape[0] != 1: - # Handle batching: apply inpaint to each image in batch if necessary. - # For this use case, we expect B=1 for the bg_image - pass # We'll handle a single image for now based on the problem description. - - # Convert image to numpy (H, W, C) for OpenCV, assuming B=1 - img_np = image_tensor[0].mul(255).byte().cpu().numpy() # Convert from 0-1 float to 0-255 uint8 - - # Convert mask to numpy (H, W), boolean to uint8 (0 or 255) - mask_np = (mask_tensor.cpu().numpy() > 0.01).astype(np.uint8) * 255 - - - # Perform inpainting - # cv2.inpaint expects BGR format, so we need to convert if image_tensor is RGB - # ComfyUI typically uses RGB. Let's assume input is RGB and convert to BGR for inpaint. - img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) - inpainted_bgr_np = cv2.inpaint(img_bgr, mask_np, 3, cv2.INPAINT_TELEA) - inpainted_rgb_np = cv2.cvtColor(inpainted_bgr_np, cv2.COLOR_BGR2RGB) - - # Convert back to torch tensor, (B, H, W, C) - inpainted_tensor = torch.from_numpy(inpainted_rgb_np).float().div(255.0).unsqueeze(0).to(image_tensor.device) - - return inpainted_tensor - - def _find_bounding_box(self, image_tensor, mask_tensor=None): - """ - Find the bounding box of non-transparent pixels in an image tensor. - - Args: - image_tensor: Image tensor of shape [H, W, C] - mask_tensor: Optional mask tensor of shape [H, W] indicating visible pixels - - Returns: - Tuple of (min_x, min_y, max_x, max_y) representing the bounding box - """ - if mask_tensor is not None: - # Use the provided mask to find visible pixels - mask = mask_tensor > 0.01 # Threshold for considering a pixel visible - elif image_tensor.shape[-1] == 4: - # Use alpha channel to find non-transparent pixels - alpha = image_tensor[..., 3] - mask = alpha > 0.01 # Threshold for considering a pixel non-transparent - else: - # For RGB images, check if pixels are not black (have some color) - # Sum the RGB channels and check if they have any value - rgb_sum = torch.sum(image_tensor[..., :3], dim=-1) - mask = rgb_sum > 0.01 # Threshold for considering a pixel non-black - - if not mask.any(): - # No non-transparent pixels found - return (0, 0, image_tensor.shape[1], image_tensor.shape[0]) - - # Find coordinates of non-transparent pixels - y_indices, x_indices = torch.where(mask) - - min_x = x_indices.min().item() - max_x = x_indices.max().item() - min_y = y_indices.min().item() - max_y = y_indices.max().item() - - return (min_x, min_y, max_x, max_y) - - def _create_square_canvas(self, max_dim, channels=4): - """ - Create a square canvas with the specified dimensions. - - Args: - max_dim: The size of the square canvas (width and height) - channels: Number of channels (3 for RGB, 4 for RGBA) - - Returns: - A square tensor of shape [max_dim, max_dim, channels] - """ - return torch.zeros((max_dim, max_dim, channels), dtype=torch.float32) - - def _place_image_in_square(self, source_image, target_canvas, bbox, has_alpha=True): - """ - Place the visible content of the source image into the target square canvas. - - Args: - source_image: Source image tensor - target_canvas: Target square canvas tensor - bbox: Bounding box of the source image (min_x, min_y, max_x, max_y) - has_alpha: Whether the source image has an alpha channel - - Returns: - Modified target_canvas with the source image placed in it - """ - min_x, min_y, max_x, max_y = bbox - source_width = max_x - min_x + 1 - source_height = max_y - min_y + 1 - - # Extract the visible content from the source image - if has_alpha: - visible_content = source_image[min_y:max_y+1, min_x:max_x+1, :] - else: - visible_content = source_image[min_y:max_y+1, min_x:max_x+1, :] - - # Calculate position to center the content in the square canvas - canvas_size = target_canvas.shape[0] - offset_x = (canvas_size - source_width) // 2 - offset_y = (canvas_size - source_height) // 2 - - # Place the visible content in the center of the square canvas - end_x = offset_x + source_width - end_y = offset_y + source_height - - if has_alpha and visible_content.shape[-1] == 4: - # Handle alpha channel properly - alpha = visible_content[..., 3:4] - target_canvas[offset_y:end_y, offset_x:end_x, :] = visible_content - else: - target_canvas[offset_y:end_y, offset_x:end_x, :visible_content.shape[-1]] = visible_content - - return target_canvas + return final_w, final_h, output_bg_image, internal_processing_image \ No newline at end of file diff --git a/nodes/trigger_ref_refresh.py b/nodes/trigger_ref_refresh.py new file mode 100644 index 0000000..4deab0f --- /dev/null +++ b/nodes/trigger_ref_refresh.py @@ -0,0 +1,131 @@ +from server import PromptServer +from aiohttp import web +import base64 +from io import BytesIO +from PIL import Image +import torch +import numpy as np + +# Import PrepareRefs class +from .prepare_refs import PrepareRefs + + +def pil_to_tensor(pil_img): + """Convert PIL image to tensor [1, H, W, C] in range [0, 1]""" + img_np = np.array(pil_img).astype(np.float32) / 255.0 + if len(img_np.shape) == 2: # Grayscale + img_np = np.stack([img_np] * 3, axis=-1) + return torch.from_numpy(img_np) + + +@PromptServer.instance.routes.post("/wanvideowrapper_qq/trigger_prepare_refs") +async def trigger_prepare_refs(request): + """ + Backend trigger to run PrepareRefs processing without node execution. + + Input JSON: + { + "bg_image": "data:image/png;base64,...", + "ref_layer_data": [{name, on, lassoShape: {additivePaths}}], + "mask_width": 640, + "mask_height": 480 + } + + Note: export_alpha and to_bounding_box are always True (hardcoded in PrepareRefs). + All ref images are exported as 768x768 PNG with alpha channel. + + Output JSON: + { + "success": true/false, + "paths": { + "bg_image_cl": "ref/bg_image_cl.png", + "bg_preview": "bg/bg_image.png", + "ref_images": ["ref/ref_1.png", ...], + "ref_masks": ["ref/ref_1_mask.png", ...] + }, + "message": "PrepareRefs processing complete: N layers processed", + "error": "error message if failed" + } + """ + try: + data = await request.json() + + # Extract parameters + bg_image_b64 = data.get("bg_image") + ref_layer_data = data.get("ref_layer_data", []) + mask_width = data.get("mask_width", 640) + mask_height = data.get("mask_height", 480) + + print(f"[trigger_prepare_refs] Received request with {len(ref_layer_data)} layers, dims: {mask_width}x{mask_height}") + + # Validate layer data + if not ref_layer_data or len(ref_layer_data) == 0: + return web.json_response({ + "success": False, + "error": "No ref_layer_data provided" + }, status=400) + + # Decode bg_image to tensor + bg_tensor = None + if bg_image_b64: + # Remove data URL prefix if present + if bg_image_b64.startswith('data:image'): + bg_image_b64 = bg_image_b64.split(',', 1)[1] + + image_bytes = base64.b64decode(bg_image_b64) + pil_img = Image.open(BytesIO(image_bytes)) + + # Convert to RGB if needed + if pil_img.mode != 'RGB': + pil_img = pil_img.convert('RGB') + + print(f"[trigger_prepare_refs] Decoded bg_image: {pil_img.size}, mode: {pil_img.mode}") + + # Convert to tensor [1, H, W, C] + bg_tensor = pil_to_tensor(pil_img).unsqueeze(0) + + # Construct prompt structure (PrepareRefs expects this format) + prompt = { + "unique_id": { + "inputs": { + "ref_layer_data": ref_layer_data + } + } + } + + # Create PrepareRefs instance and execute + print(f"[trigger_prepare_refs] Calling PrepareRefs.prepare()...") + prep_refs = PrepareRefs() + result = prep_refs.prepare( + mask_width=mask_width, + mask_height=mask_height, + bg_image=bg_tensor, + unique_id="unique_id", # Key used in prompt dict + prompt=prompt + ) + + # Extract paths from result + ui_data = result.get("ui", {}) + paths = { + "bg_image_cl": ui_data.get("bg_image_path", [None])[0], + "bg_preview": ui_data.get("bg_image_preview_path", [None])[0], + "ref_images": ui_data.get("ref_images_paths", []), + "ref_masks": ui_data.get("ref_masks_paths", []) + } + + print(f"[trigger_prepare_refs] Processing complete. Generated paths: {paths}") + + return web.json_response({ + "success": True, + "paths": paths, + "message": f"PrepareRefs processing complete: {len(ref_layer_data)} layers processed" + }) + + except Exception as e: + print(f"[trigger_prepare_refs] ERROR: {e}") + import traceback + traceback.print_exc() + return web.json_response({ + "success": False, + "error": str(e) + }, status=500) diff --git a/web/power_spline_editor/canvas_lasso.js b/web/power_spline_editor/canvas_lasso.js index 8e41040..6adf225 100644 --- a/web/power_spline_editor/canvas_lasso.js +++ b/web/power_spline_editor/canvas_lasso.js @@ -494,6 +494,9 @@ function enterLassoMode(node, refLayer) { canvas.style.cursor = 'crosshair'; } + // Update button states to show green pen for active layer + node.updateLayerButtonStates?.(); + console.log(`Lasso mode activated for ${refLayer.value.name}`); } @@ -518,6 +521,9 @@ function exitLassoMode(node) { canvas.style.cursor = 'default'; } + // Update button states to reset pen colors + node.updateLayerButtonStates?.(); + console.log('Lasso mode deactivated'); } diff --git a/web/power_spline_editor/canvas_top_row.js b/web/power_spline_editor/canvas_top_row.js index 1bed7a0..8695ca1 100644 --- a/web/power_spline_editor/canvas_top_row.js +++ b/web/power_spline_editor/canvas_top_row.js @@ -4,6 +4,7 @@ import { drawNumberWidgetPart, RgthreeBaseWidget } from './drawing_utils.js'; +import { triggerPrepareRefsBackend } from './trigger_ref_refresh.js'; // === TOP ROW WIDGET (Combined refresh button, bg_image dropdown, and width/height controls) === export class TopRowWidget extends RgthreeBaseWidget { @@ -268,28 +269,65 @@ export class TopRowWidget extends RgthreeBaseWidget { // Set up event handlers for refresh button if (this.visibility.refreshCanvasButton) { this.hitAreas.refreshCanvasButton.onClick = async () => { - // Update all box layer ref images from connected ref_images - if (node.updateAllBoxLayerRefs) { - try { - await node.updateAllBoxLayerRefs(); - } catch (e) { - console.warn('Failed to update box layer refs:', e); + // Check if ref_images input is connected to PrepareRefs + let isPrepareRefsConnected = false; + let prepareRefsNode = null; + try { + // Import the graph query function at the top of the file + const { findConnectedSourceNode } = await import('./graph_query.js'); + const sourceNodeObj = findConnectedSourceNode(node, 'ref_images'); + if (sourceNodeObj && sourceNodeObj.node) { + if (sourceNodeObj.node.type === 'PrepareRefs') { + console.log('[Refresh] PowerSplineEditor connected to PrepareRefs'); + isPrepareRefsConnected = true; + prepareRefsNode = sourceNodeObj.node; + } } + } catch (e) { + console.warn('Failed to check ref_images connection:', e); } - // Wait for reference image update and scale recalculation to complete - if (node.updateReferenceImageFromConnectedNode) { - try { - // This now properly awaits the image loading and scale recalculation - await node.updateReferenceImageFromConnectedNode(); + if (isPrepareRefsConnected && prepareRefsNode) { + // Trigger backend PrepareRefs processing + console.log('[Refresh] Triggering backend PrepareRefs processing...'); + const result = await triggerPrepareRefsBackend(prepareRefsNode); - // Force a final render with the correct scale - // The scale should already be correct from handleImageLoad -> recenterBackgroundImage - if (node.editor && node.editor.layerRenderer) { - node.editor.layerRenderer.render(); + if (result.success) { + console.log('[Refresh] Backend processing complete, loading results...'); + + // Load the newly generated images from ref folder + await this.loadImagesFromRefFolder(node); + + console.log('[Refresh] Canvas refresh complete'); + } else { + // Silent error handling - log only, no user interruption + console.error('[Refresh] Backend processing failed:', result.error); + } + } else { + // Original behavior: fetch from connected nodes + // Update all box layer ref images from connected ref_images + if (node.updateAllBoxLayerRefs) { + try { + await node.updateAllBoxLayerRefs(); + } catch (e) { + console.warn('Failed to update box layer refs:', e); + } + } + + // Wait for reference image update and scale recalculation to complete + if (node.updateReferenceImageFromConnectedNode) { + try { + // This now properly awaits the image loading and scale recalculation + await node.updateReferenceImageFromConnectedNode(); + + // Force a final render with the correct scale + // The scale should already be correct from handleImageLoad -> recenterBackgroundImage + if (node.editor && node.editor.layerRenderer) { + node.editor.layerRenderer.render(); + } + } catch (e) { + console.warn('Failed to update reference image:', e); } - } catch (e) { - console.warn('Failed to update reference image:', e); } } @@ -435,6 +473,155 @@ export class TopRowWidget extends RgthreeBaseWidget { } } + // Load images from ref folder when PrepareRefs is connected + async loadImagesFromRefFolder(node) { + const { loadImageAsBase64 } = await import('./graph_query.js'); + const timestamp = Date.now(); + + // 1. Load bg_image_cl.png from ref folder + try { + const bgImageUrl = new URL(`ref/bg_image_cl.png?t=${timestamp}`, import.meta.url).href; + const bgImageData = await loadImageAsBase64(bgImageUrl); + + if (bgImageData) { + // Store the background image + node.originalRefImageData = { + name: 'bg_image_cl.png', + base64: bgImageData.split(',')[1], + type: 'image/png' + }; + + // Update background image display based on current bg_img selection + const bgImgWidget = node.widgets?.find(w => w.name === "bg_img"); + const bg_img = bgImgWidget ? bgImgWidget.value : "None"; + await node.updateBackgroundImage(bg_img); + + console.log('Successfully loaded bg_image_cl.png from ref folder'); + } + } catch (e) { + console.warn('Failed to load bg_image_cl.png from ref folder:', e); + } + + // 2. Find and load all ref_*.png files from ref folder + const refImages = []; + const maxRefs = 5; // Try to load up to 5 reference images + + for (let i = 1; i <= maxRefs; i++) { + try { + const refImageUrl = new URL(`ref/ref_${i}.png?t=${timestamp}`, import.meta.url).href; + const refImageData = await loadImageAsBase64(refImageUrl); + + if (refImageData) { + refImages.push(refImageData); + console.log(`Successfully loaded ref_${i}.png from ref folder`); + } else { + // Stop trying if we can't load this image + break; + } + } catch (e) { + // Stop trying if we encounter an error (file doesn't exist) + console.log(`No more ref images found after ref_${i - 1}.png`); + break; + } + } + + // 3. Update box layer refs with loaded images + if (refImages.length > 0) { + const activeWidget = node.layerManager?.getActiveWidget?.(); + if (!activeWidget || activeWidget.value?.type !== 'box_layer') { + console.log('No active box layer to attach ref images'); + // Try to update all box layers instead + await this.updateAllBoxLayersFromRefFolder(node, refImages); + } else { + // Update active box layer + await this.updateBoxLayerWithRefImages(node, activeWidget, refImages); + } + } + + // Force a final render + if (node.editor && node.editor.layerRenderer) { + node.editor.layerRenderer.render(); + } + } + + // Update a specific box layer with ref images from folder + async updateBoxLayerWithRefImages(node, boxWidget, refImages) { + const attachments = []; + + for (let i = 0; i < refImages.length; i++) { + const imgData = refImages[i]; + const base64Data = imgData.startsWith('data:') + ? imgData.split(',')[1] + : imgData; + const dataUrl = imgData.startsWith('data:') + ? imgData + : `data:image/png;base64,${base64Data}`; + + // Load dimensions + const dims = await new Promise((resolve) => { + const img = new Image(); + img.onload = () => resolve({ width: img.width, height: img.height }); + img.onerror = () => resolve({ width: 1, height: 1 }); + img.src = dataUrl; + }); + + const filename = `ref_${i + 1}.png`; + attachments.push({ + path: `ref/${filename}`, + type: 'image/png', + width: dims.width, + height: dims.height, + name: filename + }); + } + + boxWidget.value.ref_attachment = { entries: attachments }; + + // Preserve existing ref_selection if valid, otherwise default to first ref + const currentSelection = boxWidget.value.ref_selection; + if (attachments.length > 0) { + // Check if current selection is still valid + if (currentSelection && currentSelection !== 'no_ref') { + const refIndex = parseInt(currentSelection.split('_')[1], 10); + if (refIndex > 0 && refIndex <= attachments.length) { + // Keep the current selection as it's still valid + boxWidget.value.ref_selection = currentSelection; + } else { + // Current selection is out of range, default to ref_1 + boxWidget.value.ref_selection = 'ref_1'; + } + } else { + // No current selection or it was 'no_ref', default to ref_1 + boxWidget.value.ref_selection = 'ref_1'; + } + } else { + boxWidget.value.ref_selection = 'no_ref'; + } + + // Clear ref image cache to force reload + if (node.editor?.layerRenderer?.clearRefImageCache) { + node.editor.layerRenderer.clearRefImageCache(); + } + + node.setDirtyCanvas(true, true); + } + + // Update all box layers with ref images from folder + async updateAllBoxLayersFromRefFolder(node, refImages) { + const widgets = node.layerManager?.getSplineWidgets?.() || []; + const boxWidgets = widgets.filter(w => w?.value?.type === 'box_layer'); + + if (!boxWidgets.length) { + console.log('No box layers found to update'); + return; + } + + for (const boxWidget of boxWidgets) { + await this.updateBoxLayerWithRefImages(node, boxWidget, refImages); + } + + console.log(`Updated ${boxWidgets.length} box layer(s) with ref images from folder`); + } // Mouse event handlers mouse(event, pos, node) { diff --git a/web/power_spline_editor/power_spline_editor.js b/web/power_spline_editor/power_spline_editor.js index d76fcd3..aa0c145 100644 --- a/web/power_spline_editor/power_spline_editor.js +++ b/web/power_spline_editor/power_spline_editor.js @@ -483,21 +483,92 @@ app.registerExtension({ const bgImgWidget = this.widgets && this.widgets.find(w => w.name === "bg_img"); const bg_img = bgImgWidget ? bgImgWidget.value : "None"; - // Try to pull a fresh ref image from connected node silently - const base64Image = await getReferenceImageFromConnectedNode(this); - if (base64Image) { - this.originalRefImageData = { - name: 'ref_image_from_connection.jpg', - base64: base64Image.split(',')[1], - type: 'image/jpeg' - }; - try { await saveRefImageToCache(this.originalRefImageData.base64, 'bg_image.png'); } catch {} - if (this.uuid) { - sessionStorage.removeItem(`spline-editor-img-${this.uuid}`); + // Check if connected to PrepareRefs - if so, load bg_image_cl.png from ref folder instead + const isConnectedToPrepareRefs = this.checkIfConnectedToPrepareRefs(); + + if (isConnectedToPrepareRefs) { + console.log('Connected to PrepareRefs, loading bg_image_cl.png from ref folder'); + try { + const timestamp = Date.now(); + const refImageUrl = new URL(`ref/bg_image_cl.png?t=${timestamp}`, import.meta.url).href; + const response = await fetch(refImageUrl); + if (response.ok) { + const blob = await response.blob(); + const base64Data = await new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(reader.result.split(',')[1]); + reader.onerror = reject; + reader.readAsDataURL(blob); + }); + + this.originalRefImageData = { + name: 'bg_image_cl.png', + base64: base64Data, + type: 'image/png' + }; + + // Process the image immediately based on bg_img selection + if (bg_img === "None") { + // Apply darkening for "None" selection + const img = new Image(); + img.onload = () => { + const canvas = document.createElement('canvas'); + canvas.width = img.width; + canvas.height = img.height; + const ctx = canvas.getContext('2d'); + ctx.drawImage(img, 0, 0); + ctx.fillStyle = 'rgba(0, 0, 0, 0.6)'; + ctx.fillRect(0, 0, canvas.width, canvas.height); + const darkenedDataUrl = canvas.toDataURL('image/png'); + this.imgData = { + name: 'bg_image_cl.png', + base64: darkenedDataUrl.split(',')[1], + type: 'image/png' + }; + if (this.editor) { + this.editor.refreshBackgroundImage(); + } + }; + img.src = `data:image/png;base64,${base64Data}`; + } else { + // For A/B/C selections, set imgData directly for now + this.imgData = { + name: 'bg_image_cl.png', + base64: base64Data, + type: 'image/png' + }; + if (this.editor) { + this.editor.refreshBackgroundImage(); + } + // Also trigger updateBackgroundImage for proper overlay handling + this.updateBackgroundImage(bg_img); + } + + if (this.uuid) { + sessionStorage.removeItem(`spline-editor-img-${this.uuid}`); + } + return; + } + } catch (e) { + console.error('Failed to load bg_image_cl.png from ref folder:', e); + } + } else { + // Try to pull a fresh ref image from connected node silently + const base64Image = await getReferenceImageFromConnectedNode(this); + if (base64Image) { + this.originalRefImageData = { + name: 'ref_image_from_connection.jpg', + base64: base64Image.split(',')[1], + type: 'image/jpeg' + }; + try { await saveRefImageToCache(this.originalRefImageData.base64, 'bg_image.png'); } catch {} + if (this.uuid) { + sessionStorage.removeItem(`spline-editor-img-${this.uuid}`); + } + this.updateBackgroundImage(bg_img); + this.editor?.refreshBackgroundImage?.(); + return; } - this.updateBackgroundImage(bg_img); - this.editor?.refreshBackgroundImage?.(); - return; } } catch {} // Fallback: just apply current selection to force overlay/cached handling @@ -617,6 +688,35 @@ app.registerExtension({ // Set initial node size this.updateNodeHeight(); + // Helper method to load the correct cached ref image based on PrepareRefs connection + this.loadCorrectCachedRefImage = async function() { + const isConnectedToPrepareRefs = this.checkIfConnectedToPrepareRefs(); + + if (isConnectedToPrepareRefs) { + // Load bg_image_cl.png from ref folder + try { + const timestamp = Date.now(); + const refImageUrl = new URL(`ref/bg_image_cl.png?t=${timestamp}`, import.meta.url).href; + const response = await fetch(refImageUrl); + if (response.ok) { + const blob = await response.blob(); + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onloadend = () => resolve(reader.result); + reader.onerror = reject; + reader.readAsDataURL(blob); + }); + } + } catch (e) { + console.error('Failed to load bg_image_cl.png from ref folder:', e); + } + return null; + } else { + // Load bg_image.png from bg folder (standard behavior) + return await loadCachedRefImageAsBase64(); + } + }.bind(this); + // Add method to update background image immediately based on bg_img selection this.updateBackgroundImage = async function(bg_img) { // First check if we have saved dimensions to maintain proportions @@ -661,7 +761,7 @@ app.registerExtension({ }; img.onerror = () => { // Fallback: try to load cached bg_image first, then default A.jpg - loadCachedRefImageAsBase64().then(cachedImageUrl => { + this.loadCorrectCachedRefImage().then(cachedImageUrl => { if (cachedImageUrl) { const fallbackImg = new Image(); fallbackImg.onload = () => { @@ -717,7 +817,7 @@ app.registerExtension({ img.src = `data:image/jpeg;base64,${this.originalRefImageData.base64}`; } else { // If no original ref image data, try to load cached bg_image first, then fallback to default - loadCachedRefImageAsBase64().then(cachedImageUrl => { + this.loadCorrectCachedRefImage().then(cachedImageUrl => { if (cachedImageUrl) { const img = new Image(); img.onload = () => { @@ -781,10 +881,11 @@ app.registerExtension({ if (this.originalRefImageData && this.originalRefImageData.base64) { // Use original ref image if available (from refresh button) - refImageForOverlay = `data:image/jpeg;base64,${this.originalRefImageData.base64}`; + const imgType = this.originalRefImageData.type || 'image/jpeg'; + refImageForOverlay = `data:${imgType};base64,${this.originalRefImageData.base64}`; } else { // Otherwise, load from cached - loadCachedRefImageAsBase64().then(cachedImageUrl => { + this.loadCorrectCachedRefImage().then(cachedImageUrl => { if (cachedImageUrl) { // Create a properly scaled overlay using the original cached ref image this.createScaledImageOverlay(cachedImageUrl, bg_img, imageUrl); @@ -815,6 +916,58 @@ app.registerExtension({ console.log('Attempting to update reference image from connected node...'); try { + // Check if we're connected to PrepareRefs node + const isConnectedToPrepareRefs = this.checkIfConnectedToPrepareRefs(); + + if (isConnectedToPrepareRefs) { + console.log('Connected to PrepareRefs node, loading bg_image_cl.png from ref folder'); + try { + const timestamp = Date.now(); + const refImageUrl = new URL(`ref/bg_image_cl.png?t=${timestamp}`, import.meta.url).href; + const response = await fetch(refImageUrl); + if (!response.ok) { + console.error('Failed to load bg_image_cl.png from ref folder'); + alert('Could not load bg_image_cl.png from ref folder. Make sure the file exists.'); + return; + } + + const blob = await response.blob(); + const base64Data = await new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(reader.result.split(',')[1]); + reader.onerror = reject; + reader.readAsDataURL(blob); + }); + + this.originalRefImageData = { + name: 'bg_image_cl.png', + base64: base64Data, + type: 'image/png' + }; + + // Clear session storage cache for this node's image to force refresh + if (this.uuid) { + sessionStorage.removeItem(`spline-editor-img-${this.uuid}`); + } + + // Get current bg_img selection to update background accordingly + const bgImgWidget = this.widgets.find(w => w.name === "bg_img"); + const bg_img = bgImgWidget ? bgImgWidget.value : "None"; + + // Update background based on current bg_img selection + this.updateBackgroundImage(bg_img); + + // Wait for the image processing to complete + await new Promise(resolve => setTimeout(resolve, 200)); + + console.log('Successfully loaded bg_image_cl.png from ref folder'); + } catch (error) { + console.error('Error loading bg_image_cl.png:', error); + alert('Error loading bg_image_cl.png: ' + error.message); + } + return; + } + // Get reference image from connected bg_image input const base64Image = await getReferenceImageFromConnectedNode(this, 'bg_image'); if (!base64Image) { @@ -864,6 +1017,15 @@ app.registerExtension({ return; } + // Check if we're connected to PrepareRefs node + const isConnectedToPrepareRefs = this.checkIfConnectedToPrepareRefs(); + + if (isConnectedToPrepareRefs) { + // When connected to PrepareRefs, load ref images from ref folder instead + await this.loadRefImagesFromRefFolder(activeWidget, desiredSelection); + return; + } + // Fetch images from ref_images input (first frame used) const images = await getReferenceImagesFromConnectedNode(this); if (!images || images.length === 0) { @@ -980,6 +1142,17 @@ app.registerExtension({ this.layerRenderer.clearRefImageCache(); } + // Check if we're connected to PrepareRefs node + const isConnectedToPrepareRefs = this.checkIfConnectedToPrepareRefs(); + + if (isConnectedToPrepareRefs) { + // When connected to PrepareRefs, load ref images from ref folder for all box layers + for (const boxWidget of boxWidgets) { + await this.loadRefImagesFromRefFolder(boxWidget, boxWidget.value.ref_selection || 'ref_1'); + } + return; + } + // Fetch images from ref_images input (first frame used) const images = await getReferenceImagesFromConnectedNode(this); if (!images || images.length === 0) { @@ -1372,6 +1545,142 @@ app.registerExtension({ } }.bind(this); + // Helper method to check if connected to PrepareRefs node + this.checkIfConnectedToPrepareRefs = function() { + // Check if any of our inputs is connected to a PrepareRefs node + const graph = app.graph; + if (!graph || !graph.links || !this.inputs) { + return false; + } + + for (let i = 0; i < this.inputs.length; i++) { + const input = this.inputs[i]; + if (!input || !input.name) continue; + + // Look for connections to this input + let link = null; + if (graph.links instanceof Map) { + for (const [linkId, linkObj] of graph.links) { + if (linkObj && linkObj.target_id === this.id && linkObj.target_slot === i) { + link = linkObj; + break; + } + } + } else if (Array.isArray(graph.links)) { + link = graph.links.find(linkObj => + linkObj && linkObj.target_id === this.id && linkObj.target_slot === i + ); + } + + if (link) { + const sourceNode = graph._nodes?.find(n => n.id === link.origin_id); + if (sourceNode && sourceNode.type === 'PrepareRefs') { + return true; + } + } + } + return false; + }.bind(this); + + // Helper method to load ref images from ref folder + this.loadRefImagesFromRefFolder = async function(boxWidget, desiredSelection = 'ref_1') { + try { + // Load ALL available ref images (ref_1.png through ref_5.png) from ref folder + const attachments = []; + const maxRefs = 5; + const timestamp = Date.now(); + + for (let i = 1; i <= maxRefs; i++) { + try { + const refImageUrl = new URL(`ref/ref_${i}.png?t=${timestamp}`, import.meta.url).href; + const response = await fetch(refImageUrl); + + if (!response.ok) { + console.log(`ref_${i}.png not found in ref folder, skipping`); + continue; + } + + const blob = await response.blob(); + const base64Data = await new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(reader.result.split(',')[1]); + reader.onerror = reject; + reader.readAsDataURL(blob); + }); + + // Load dimensions to fit box without stretching + const dims = await new Promise((resolve) => { + const img = new Image(); + img.onload = () => resolve({ width: img.width, height: img.height }); + img.onerror = () => resolve({ width: 1, height: 1 }); + img.src = `data:image/png;base64,${base64Data}`; + }); + + // Save image to disk and store path + const filename = `ref_${i}.png`; + await saveRefImageToCache(base64Data, filename); + + // Create attachment + attachments.push({ + path: `ref/${filename}`, + type: 'image/png', + width: dims.width, + height: dims.height, + name: filename + }); + + console.log(`Successfully loaded ref_${i}.png from ref folder`); + } catch (error) { + console.error(`Error loading ref_${i}.png:`, error); + } + } + + if (attachments.length === 0) { + console.error('No ref images found in ref folder'); + return; + } + + // Set all attachments on the box widget + boxWidget.value.ref_attachment = { entries: attachments }; + + // Set the desired selection if valid, otherwise set to first available + const availableOptions = boxWidget._getRefOptions ? boxWidget._getRefOptions() : ['no_ref', 'ref_1', 'ref_2', 'ref_3', 'ref_4', 'ref_5']; + const desiredIdx = desiredSelection && availableOptions.includes(desiredSelection) + ? Math.max(0, (parseInt(desiredSelection.split('_')[1], 10) || 1) - 1) + : 0; + + if (desiredIdx >= attachments.length) { + boxWidget.value.ref_selection = 'no_ref'; + } else { + const clampedIndex = Math.min(attachments.length - 1, desiredIdx); + boxWidget.value.ref_selection = `ref_${clampedIndex + 1}`; + } + + // Persist to sessionStorage + try { + const keyId = this.id ?? this.uuid; + const key = keyId ? `spline-editor-boxref-${keyId}-${boxWidget.value.name || boxWidget.name || 'box'}` : null; + if (key) { + sessionStorage.setItem(key, JSON.stringify({ + attachment: boxWidget.value.ref_attachment, + selection: boxWidget.value.ref_selection, + })); + } + } catch (e) { + console.warn('Failed to persist box ref attachment to session:', e); + } + + // Clear ref image cache to force reload + if (this.layerRenderer && this.layerRenderer.clearRefImageCache) { + this.layerRenderer.clearRefImageCache(); + } + + console.log(`Successfully loaded ${attachments.length} ref images from ref folder for box layer`); + } catch (error) { + console.error('Error loading ref images from ref folder:', error); + } + }.bind(this); + // context menu this.contextMenu = document.createElement("div"); this.contextMenu.className = 'spline-editor-context-menu'; @@ -1488,9 +1797,10 @@ app.registerExtension({ } // Try to load cached bg_image first - loadCachedRefImageAsBase64().then(cachedImageUrl => { + this.loadCorrectCachedRefImage().then(cachedImageUrl => { if (cachedImageUrl) { - this.loadBackgroundImageFromUrl(cachedImageUrl, 'bg_image.png', targetWidth, targetHeight); + const filename = this.checkIfConnectedToPrepareRefs() ? 'bg_image_cl.png' : 'bg_image.png'; + this.loadBackgroundImageFromUrl(cachedImageUrl, filename, targetWidth, targetHeight); } else { // Fallback to default A.jpg const timestamp = Date.now(); @@ -1685,7 +1995,7 @@ app.registerExtension({ // Determine which image to load based on the bg_img selection if (bg_img === "None") { // Try to load cached bg_image first, then fallback to default - loadCachedRefImageAsBase64().then(cachedImageUrl => { + this.loadCorrectCachedRefImage().then(cachedImageUrl => { if (cachedImageUrl) { // Apply darkening effect to match onExecuted behavior const img = new Image(); @@ -1726,9 +2036,9 @@ app.registerExtension({ // Load the selected background image (A, B, or C) from the bg folder const timestamp = Date.now(); const imageUrl = new URL(`bg/${bg_img}.jpg?t=${timestamp}`, import.meta.url).href; - + // Use the overlay function from image_overlay.js - loadCachedRefImageAsBase64().then(cachedImageUrl => { + this.loadCorrectCachedRefImage().then(cachedImageUrl => { if (cachedImageUrl) { createImageOverlayForConfigure( cachedImageUrl, diff --git a/web/power_spline_editor/prepare_refs.js b/web/power_spline_editor/prepare_refs.js index 63adcfc..483773c 100644 --- a/web/power_spline_editor/prepare_refs.js +++ b/web/power_spline_editor/prepare_refs.js @@ -6,598 +6,685 @@ import { RefLayerWidget } from './layer_type_ref.js'; import { attachLassoHelpers } from './canvas_lasso.js'; import RefCanvas from './canvas_main_ref.js'; +// =================================================================== +// Session Storage Helpers +// =================================================================== +function safeSetSessionItem(key, value) { + try { sessionStorage.setItem(key, value); } catch (e) { /* ignore */ } +} + +function safeGetSessionItem(key) { + try { return sessionStorage.getItem(key); } catch (e) { return null; } +} + +// =================================================================== +// Image Loading Utilities +// =================================================================== +async function loadImageFromUrl(url, canvasEditor) { + console.log('[loadImageFromUrl] Loading:', url); + const timestampedUrl = `${url}${url.includes('?') ? '&' : '?'}t=${Date.now()}`; + + return new Promise((resolve, reject) => { + const img = new Image(); + img.onload = () => { + canvasEditor.loadBackgroundImage(img); + resolve(img); + }; + img.onerror = (err) => { + console.error('[loadImageFromUrl] Failed to load image:', url, err); + reject(err); + }; + img.src = timestampedUrl; + }); +} + +async function loadRefImagesFromPaths(paths, node) { + if (!Array.isArray(paths) || paths.length === 0) return []; + + const images = []; + for (let i = 0; i < paths.length; i++) { + const path = paths[i]; + try { + const url = new URL(path, import.meta.url).href + `?t=${Date.now()}`; + const img = await new Promise((resolve, reject) => { + const image = new Image(); + image.onload = () => resolve(image); + image.onerror = reject; + image.src = url; + }); + images.push(img); + } catch (err) { + console.error(`[loadRefImagesFromPaths] Failed to load ref image ${i}:`, path, err); + } + } + return images; +} + +// =================================================================== +// Node Lifecycle: PrepareRefs +// =================================================================== app.registerExtension({ name: 'WanVideoWrapper_QQ.PrepareRefs', async beforeRegisterNodeDef(nodeType, nodeData) { if (nodeData?.name !== 'PrepareRefs') return; + // ================================================================= + // onNodeCreated – Main Setup + // ================================================================= chainCallback(nodeType.prototype, 'onNodeCreated', function () { this.serialize_widgets = true; this.resizable = false; this.properties = this.properties || {}; this.properties.userAdjustedDims = this.properties.userAdjustedDims || false; - const widthWidget = this.widgets.find((w) => w.name === 'mask_width'); - const heightWidget = this.widgets.find((w) => w.name === 'mask_height'); + this.uuid = this.uuid || makeUUID(); - const setUserAdjusted = () => { - this.properties.userAdjustedDims = true; - }; + // Hide original dimension widgets + const widthWidget = this.widgets.find(w => w.name === 'mask_width'); + const heightWidget = this.widgets.find(w => w.name === 'mask_height'); - if (!this.addCustomWidget) { - this.addCustomWidget = function (widget) { - widget.parent = this; - this.widgets = this.widgets || []; - this.widgets.push(widget); - const originalMouse = widget.mouse; - widget.mouse = function (event, pos, node) { - const localPos = [pos[0], pos[1] - (widget.last_y || 0)]; - return originalMouse?.call(this, event, localPos, node); - }; - return widget; - }; - } + const markUserAdjusted = () => { this.properties.userAdjustedDims = true; }; - if (widthWidget) { - const original = widthWidget.callback; - widthWidget.callback = (value) => { - original?.call(widthWidget, value); - setUserAdjusted(); - this.updateCanvasSizeFromWidgets?.(); - return value; - }; - hideWidgetForGood(this, widthWidget); - } - - if (heightWidget) { - const original = heightWidget.callback; - heightWidget.callback = (value) => { - original?.call(heightWidget, value); - setUserAdjusted(); - this.updateCanvasSizeFromWidgets?.(); - return value; - }; - hideWidgetForGood(this, heightWidget); - } + [widthWidget, heightWidget].forEach(w => { + if (w) { + const orig = w.callback; + w.callback = (v) => { orig?.call(w, v); markUserAdjusted(); this.updateCanvasSizeFromWidgets?.(); }; + hideWidgetForGood(this, w); + } + }); + // Size manager for consistent canvas/node sizing this.sizeManager = new NodeSizeManager(this, { spacingAfterCanvas: 36, - canvasWidth: widthWidget ? widthWidget.value : 640, - canvasHeight: heightWidget ? heightWidget.value : 480, + canvasWidth: widthWidget?.value ?? 640, + canvasHeight: heightWidget?.value ?? 480, minNodeWidth: 640, minNodeHeight: 480, }); - this.onResize = function (size) { - const constrained = this.sizeManager.onNodeResized(size); - size[0] = constrained[0]; - size[1] = constrained[1]; - }; + this.onResize = (size) => { + const constrained = this.sizeManager.onNodeResized(size); + size[0] = constrained[0]; + size[1] = constrained[1]; + }; - // Reuse the existing top row widget (refresh + width/height controls) - const topRow = new TopRowWidget('prepare_refs_top_row', { + // Top row controls + this.addCustomWidget(new TopRowWidget('prepare_refs_top_row', { refreshCanvasButton: true, refreshFramesButton: false, bgImgControl: false, animToggleButton: false, + })); + + this.setupUI(); + this.setupCanvas(); + this.setupRefLayers(); + this.restoreSessionImages(); + }); + + // ================================================================= + // UI Construction + // ================================================================= + nodeType.prototype.setupUI = function () { + const container = this.createMainContainer(); + this.domWidget = this.addDOMWidget('PrepareRefs', 'PrepareRefsCanvas', container, { + serialize: false, + hideOnZoom: false, }); - this.addCustomWidget(topRow); + this.domWidget.computeSize = () => { + const h = this.widgets.find(w => w.name === 'mask_height')?.value ?? 480; + return [0, h]; + }; + }; + + nodeType.prototype.createMainContainer = function () { const container = document.createElement('div'); - this.uuid = this.uuid || makeUUID(); container.id = `prepare-refs-${this.uuid}`; - container.style.margin = '0'; - container.style.padding = '0'; - container.style.display = 'flex'; - container.style.flexDirection = 'column'; - container.style.gap = '0px'; + container.style.cssText = ` + display: flex; + flex-direction: column; + margin: 0; padding: 0; gap: 0; + `; - // Two-column layout: sidebar + canvas const row = document.createElement('div'); - row.style.display = 'flex'; - row.style.flexDirection = 'row'; - row.style.flexWrap = 'nowrap'; - row.style.gap = '8px'; - row.style.alignItems = 'stretch'; + row.style.cssText = ` + display: flex; gap: 8px; align-items: stretch; + `; container.appendChild(row); + this.sidebar = this.createSidebar(); + row.appendChild(this.sidebar); + + this.canvasWrap = this.createCanvasWrapper(); + row.appendChild(this.canvasWrap); + + return container; + }; + + nodeType.prototype.createSidebar = function () { const sidebar = document.createElement('div'); - sidebar.style.width = '120px'; - sidebar.style.minWidth = '120px'; - sidebar.style.maxWidth = '120px'; - sidebar.style.flex = '0 0 120px'; - sidebar.style.display = 'flex'; - sidebar.style.flexDirection = 'column'; - sidebar.style.alignItems = 'stretch'; - sidebar.style.justifyContent = 'flex-start'; - sidebar.style.gap = '8px'; - sidebar.style.padding = '8px'; - sidebar.style.color = '#b5b5b5'; - sidebar.style.fontSize = '12px'; - sidebar.style.border = '1px solid #3a3a3a'; - sidebar.style.background = 'transparent'; - sidebar.style.overflow = 'hidden'; - sidebar.style.boxSizing = 'border-box'; + sidebar.style.cssText = ` + width: 120px; flex: 0 0 120px; display: flex; flex-direction: column; + gap: 8px; padding: 8px; background: transparent; border: 1px solid #3a3a3a; + color: #b5b5b5; font-size: 12px; box-sizing: border-box; + `; - const addRefBtn = document.createElement('button'); - addRefBtn.textContent = '+ Add ref'; - addRefBtn.style.cursor = 'pointer'; - addRefBtn.style.padding = '8px 12px'; - addRefBtn.style.border = '0.75px solid #00000044'; - addRefBtn.style.background = '#1e1e1e'; - addRefBtn.style.color = '#ddd'; - addRefBtn.style.borderRadius = '4px'; - addRefBtn.style.fontSize = '12px'; - addRefBtn.style.fontWeight = '500'; - addRefBtn.style.textAlign = 'center'; - addRefBtn.style.position = 'relative'; - addRefBtn.style.boxShadow = '0 1px 0 #000000aa, inset 0 0.75px 0 #ffffff22, inset 0.75px 0 0 #ffffff11, inset -0.75px 0 0 #00000044'; - addRefBtn.style.transition = 'all 0.05s ease'; - addRefBtn.onmousedown = () => { - addRefBtn.style.background = '#444'; - addRefBtn.style.transform = 'translateY(1px)'; - addRefBtn.style.boxShadow = 'inset 0 0.75px 0 #00000088'; - addRefBtn.style.border = '0.75px solid #00000088'; - }; - addRefBtn.onmouseup = () => { - addRefBtn.style.background = '#1e1e1e'; - addRefBtn.style.transform = 'translateY(0)'; - addRefBtn.style.boxShadow = '0 1px 0 #000000aa, inset 0 0.75px 0 #ffffff22, inset 0.75px 0 0 #ffffff11, inset -0.75px 0 0 #00000044'; - addRefBtn.style.border = '0.75px solid #00000044'; - }; - addRefBtn.onmouseleave = () => { - addRefBtn.style.background = '#1e1e1e'; - addRefBtn.style.transform = 'translateY(0)'; - addRefBtn.style.boxShadow = '0 1px 0 #000000aa, inset 0 0.75px 0 #ffffff22, inset 0.75px 0 0 #ffffff11, inset -0.75px 0 0 #00000044'; - addRefBtn.style.border = '0.75px solid #00000044'; - }; - addRefBtn.onclick = () => this.addRefLayer?.(); - sidebar.appendChild(addRefBtn); + const addBtn = this.createAddRefButton(); + sidebar.appendChild(addBtn); - // Container for ref layers - const refLayersContainer = document.createElement('div'); - refLayersContainer.style.display = 'flex'; - refLayersContainer.style.flexDirection = 'column'; - refLayersContainer.style.gap = '4px'; - refLayersContainer.style.marginTop = '4px'; - refLayersContainer.style.overflowY = 'auto'; - refLayersContainer.style.flex = '1'; - sidebar.appendChild(refLayersContainer); + this.refLayersContainer = document.createElement('div'); + this.refLayersContainer.style.cssText = ` + flex: 1; display: flex; flex-direction: column; gap: 4px; + overflow-y: auto; margin-top: 4px; + `; + sidebar.appendChild(this.refLayersContainer); - // Initialize ref layer management - this.refLayers = []; - this.refLayerCount = 0; - this.refLayersContainer = refLayersContainer; - this.sidebar = sidebar; - this.selectedRefLayer = null; + return sidebar; + }; - row.appendChild(sidebar); + nodeType.prototype.createAddRefButton = function () { + const btn = document.createElement('button'); + btn.textContent = '+ Add ref'; + btn.style.cssText = ` + padding: 8px 12px; background: #1e1e1e; color: #ddd; border: 0.75px solid #00000044; + border-radius: 4px; font-size: 12px; font-weight: 500; cursor: pointer; + box-shadow: 0 1px 0 #000000aa, inset 0 0.75px 0 #ffffff22; + transition: all 0.05s ease; + `; + ['mousedown', 'mouseleave', 'mouseup'].forEach(ev => { + btn.addEventListener(ev, () => { + btn.style.background = ev === 'mousedown' ? '#444' : '#1e1e1e'; + btn.style.transform = ev === 'mousedown' ? 'translateY(1px)' : 'translateY(0)'; + }); + }); + btn.onclick = () => this.addRefLayer(); + return btn; + }; - const canvasWrap = document.createElement('div'); - canvasWrap.style.flex = '1 1 0'; - canvasWrap.style.minWidth = '0'; - canvasWrap.style.display = 'flex'; - canvasWrap.style.alignItems = 'center'; - canvasWrap.style.justifyContent = 'center'; - canvasWrap.style.paddingTop = '0'; - canvasWrap.style.overflow = 'hidden'; - canvasWrap.style.boxSizing = 'border-box'; - row.appendChild(canvasWrap); + nodeType.prototype.createCanvasWrapper = function () { + const wrap = document.createElement('div'); + wrap.style.cssText = ` + flex: 1; display: flex; align-items: center; justify-content: center; + overflow: hidden; padding-top: 0; box-sizing: border-box; + `; - const canvas = document.createElement('canvas'); - canvas.style.borderRadius = '0px'; - canvas.style.background = '#222'; - canvas.style.border = '1px solid gray'; - canvas.style.boxShadow = 'none'; - canvas.style.display = 'block'; - canvas.style.maxWidth = '100%'; - canvas.style.maxHeight = '100%'; - canvas.style.width = 'auto'; - canvas.style.height = 'auto'; - canvasWrap.appendChild(canvas); + this.refsCanvas = document.createElement('canvas'); + this.refsCanvas.style.cssText = ` + background: #222; border: 1px solid gray; max-width: 100%; max-height: 100%; + display: block; border-radius: 0; + `; + wrap.appendChild(this.refsCanvas); - this.refsCanvas = canvas; - this.canvasWrap = canvasWrap; + return wrap; + }; - // Initialize RefCanvas (uses PowerSplineEditor coordinate system) - this.refCanvasEditor = new RefCanvas(canvas, this); - - // Attach lasso drawing helpers + // ================================================================= + // Canvas & RefCanvas Setup + // ================================================================= + nodeType.prototype.setupCanvas = function () { + this.refCanvasEditor = new RefCanvas(this.refsCanvas, this); attachLassoHelpers(this); this.refreshCanvas = () => { - if (!this.refCanvasEditor) return; - - // Render the canvas (background + grid) this.refCanvasEditor.render(); - // Add border const ctx = this.refsCanvas.getContext('2d'); const { width, height } = this.refsCanvas; ctx.strokeStyle = 'gray'; ctx.lineWidth = 1; ctx.strokeRect(0, 0, width, height); - // Draw placeholder if no bg_image if (!this.refCanvasEditor.backgroundImage) { ctx.fillStyle = '#787878'; ctx.font = '12px sans-serif'; - ctx.textAlign = 'left'; - ctx.fillText('PrepareRefs canvas', 12 + 8, 18 + 8); + ctx.fillText('PrepareRefs canvas', 20, 26); } - // Render lasso shapes for each layer using refCanvasEditor's coordinate system - if (this.refLayers && this.renderLayerLassoShapes) { - this.refLayers.forEach(refLayer => { - if (refLayer.value.lassoShape) { - // Show green only for selected layer, gray for others - const isSelected = this.selectedRefLayer === refLayer; - this.renderLayerLassoShapes(ctx, refLayer.value.lassoShape, width, height, isSelected, this.refCanvasEditor); - } - }); - } + this.refLayers?.forEach(layer => { + if (layer.value.lassoShape) { + const isSelected = this.selectedRefLayer === layer; + this.renderLayerLassoShapes?.(ctx, layer.value.lassoShape, width, height, isSelected, this.refCanvasEditor); + } + }); - // Render preview if currently drawing - if (this.renderLassoPreview) { - this.renderLassoPreview(); - } - }; - - this.setDimensionValue = (widgetName, value) => { - const widget = this.widgets.find((w) => w.name === widgetName); - if (!widget) return; - const clamp = - widgetName === 'mask_width' - ? this.sizeManager?.constrainCanvasWidth?.bind(this.sizeManager) - : this.sizeManager?.constrainCanvasHeight?.bind(this.sizeManager); - const clamped = clamp ? clamp(Math.round(value)) : Math.round(value); - widget.value = clamped; - widget.callback?.call(widget, clamped); - this.setDirtyCanvas(true, true); - }; - - this.adjustDimension = (widgetName, delta) => { - const widget = this.widgets.find((w) => w.name === widgetName); - if (!widget) return; - const next = Number(widget.value || 0) + delta; - this.setDimensionValue(widgetName, next); - this.updateCanvasSizeFromWidgets(); - }; - - this.promptDimension = (widgetName, label) => { - const widget = this.widgets.find((w) => w.name === widgetName); - if (!widget) return; - const next = prompt(`Set ${label}`, widget.value); - if (next === null) return; - const parsed = parseInt(next, 10); - if (Number.isNaN(parsed)) return; - this.setDimensionValue(widgetName, parsed); - this.updateCanvasSizeFromWidgets(); - this.properties.userAdjustedDims = true; + this.renderLassoPreview?.(); }; this.forceCanvasRefresh = () => { this.refreshCanvas(); this.setDirtyCanvas(true, true); }; + }; - // Wire the TopRow refresh button to pull the connected bg_image and paint it - this.updateReferenceImageFromConnectedNode = async () => { - try { - const base64 = await getReferenceImageFromConnectedNode(this, 'bg_image'); - if (!base64) { - console.warn('PrepareRefs: no bg_image found on refresh'); - this.forceCanvasRefresh(); - return; - } + // ================================================================= + // Ref Layer Management + // ================================================================= + nodeType.prototype.setupRefLayers = function () { + this.refLayers = []; + this.selectedRefLayer = null; - await new Promise((resolve, reject) => { - const img = new Image(); - img.onload = () => { - // Use RefCanvas to load bg_image (sets up coordinate system automatically) - this.refCanvasEditor.loadBackgroundImage(img); - this.loadedBgImage = img; - // Keep user-selected canvas size, only refresh the visual preview - this.forceCanvasRefresh(); - resolve(); - }; - img.onerror = reject; - img.src = base64; - }); - } catch (e) { - console.error('PrepareRefs: failed to refresh bg_image', e); - this.forceCanvasRefresh(); - } - }; - this.handleFramesRefresh = () => { - this.forceCanvasRefresh(); - }; + this.addSerializedRefDataWidget(); + }; - this.updateCanvasSizeFromWidgets = () => { - const width = widthWidget ? widthWidget.value : 640; - const height = heightWidget ? heightWidget.value : 480; - // Use RefCanvas to update size (handles canvas resize and image recentering) - this.refCanvasEditor.setSize(width, height); - // Remove CSS dimensions - let it scale naturally - this.refsCanvas.style.width = ''; - this.refsCanvas.style.height = ''; - this.sizeManager.setCanvasSize(width, height); - this.refreshCanvas(); - this.sizeManager.updateSize(true); - }; - - // Add a method to get ref layer data for serialization - this.getRefLayerData = () => { - if (!this.refLayers || !Array.isArray(this.refLayers)) { - return []; - } - - // Filter and return only layers that have shapes (non-empty additivePaths) - return this.refLayers - .filter(layer => { - // Check if the layer has lasso shapes with actual paths - const lassoShape = layer.value?.lassoShape; - const additivePaths = lassoShape?.additivePaths; - return lassoShape && additivePaths && Array.isArray(additivePaths) && additivePaths.length > 0; - }) - .map(layer => ({ ...layer.value })); - }; - - // Add the ref layer data widget early in the process - if (!this.widgets) this.widgets = []; - const refDataWidget = { + nodeType.prototype.addSerializedRefDataWidget = function () { + const widget = { type: "custom", name: "ref_layer_data", - value: this.getRefLayerData(), - callback: (value) => { - // Update the value when needed - this.widgets.find(w => w.name === "ref_layer_data").value = this.getRefLayerData(); - } + value: [], + serialize: true, + callback: () => { widget.value = this.getRefLayerData(); } }; - // Mark this widget for serialization - refDataWidget.serialize = true; - this.widgets.push(refDataWidget); + this.widgets.push(widget); + }; - this.domWidget = this.addDOMWidget(nodeData.name, 'PrepareRefsCanvas', container, { - serialize: false, - hideOnZoom: false, + nodeType.prototype.getRefLayerData = function () { + return (this.refLayers || []) + .filter(l => l.value?.lassoShape?.additivePaths?.length > 0) + .map(l => ({ ...l.value })); + }; + + nodeType.prototype.restoreRefLayers = function (layerDataArray) { + if (!Array.isArray(layerDataArray) || layerDataArray.length === 0) return; + + console.log('[PrepareRefs] Restoring', layerDataArray.length, 'layers from saved data'); + + // Clear existing layers + this.refLayers.forEach(layer => { + if (layer.domElement?.parentNode) { + layer.domElement.parentNode.removeChild(layer.domElement); + } + }); + this.refLayers = []; + this.selectedRefLayer = null; + + // Recreate layers from saved data + layerDataArray.forEach((layerData, index) => { + const layer = new RefLayerWidget(layerData.name || `ref_${index + 1}`); + layer.value = { ...layerData }; + layer.node = this; + this.refLayers.push(layer); + + const el = this.createRefLayerElement(layer); + this.refLayersContainer.appendChild(el); + layer.domElement = el; }); - this.domWidget.computeSize = function (width) { - const heightWidget = this.node?.widgets?.find((w) => w.name === 'mask_height'); - const canvasHeight = heightWidget ? heightWidget.value : 480; - return [width, canvasHeight]; - }.bind(this.domWidget); + // Save to session storage for quick restore + try { + const sessionKey = `prepare-refs-layers-${this.uuid}`; + safeSetSessionItem(sessionKey, JSON.stringify(layerDataArray)); + } catch (e) { + console.warn('[PrepareRefs] Failed to save layers to session:', e); + } + + this.forceCanvasRefresh(); + }; + + nodeType.prototype.addRefLayer = function () { + const nextNum = (this.refLayers.reduce((m, l) => { + const n = parseInt(l.value.name.match(/ref_(\d+)/)?.[1] || 0); + return Math.max(m, n); + }, 0)) + 1; + + const layer = new RefLayerWidget(`ref_${nextNum}`); + layer.value = { on: true, name: `ref_${nextNum}` }; + layer.node = this; + this.refLayers.push(layer); + + const el = this.createRefLayerElement(layer); + this.refLayersContainer.appendChild(el); + layer.domElement = el; + + this.selectRefLayer(layer); + this.enterLassoMode?.(layer); + this.updateRefDataWidget(); + this.forceCanvasRefresh(); + }; + + nodeType.prototype.createRefLayerElement = function (layer) { + const el = document.createElement('div'); + el.style.cssText = ` + display: flex; align-items: center; padding: 4px 6px; background: #262626; + border: 1px solid #3a3a3a; border-radius: 3px; font-size: 11px; color: #e0e0e0; + cursor: pointer; user-select: none; transition: background 0.1s ease; + `; + + const name = document.createElement('span'); + name.textContent = layer.value.name; + name.style.flex = '1'; + el.appendChild(name); + + const editBtn = this.createIconButton('Edit', '✎', () => { + this.selectRefLayer(layer); + if (this._lassoDrawingActive && this._lassoActiveLayer === layer) { + this.exitLassoMode?.(); + } else { + this.enterLassoMode?.(layer); + } + }); + el.appendChild(editBtn); + + // Store reference to edit button for color updates + layer.editButton = editBtn; + + const delBtn = this.createIconButton('Delete', '✕', () => this.removeRefLayer(layer)); + el.appendChild(delBtn); + + el.onclick = (e) => { + if (e.target.tagName === 'BUTTON') return; + if (this.selectedRefLayer === layer) { + this.deselectRefLayer(); + } else { + this.selectRefLayer(layer); + } + }; + + return el; + }; + + nodeType.prototype.createIconButton = function (title, symbol, callback) { + const btn = document.createElement('button'); + btn.innerHTML = symbol; + btn.title = title; + btn.style.cssText = ` + width: 18px; height: 18px; padding: 0; background: transparent; + border: none; color: #888; cursor: pointer; font-size: 12px; + display: flex; align-items: center; justify-content: center; + `; + btn.onmousedown = (e) => { e.stopPropagation(); btn.style.color = '#fff'; }; + btn.onmouseleave = () => { + // Only reset to gray if not in edit mode + const isInEditMode = this._lassoDrawingActive && this._lassoActiveLayer && + this._lassoActiveLayer.editButton === btn; + if (!isInEditMode) { + btn.style.color = '#888'; + } + }; + btn.onmouseup = () => { + // Only reset to gray if not in edit mode + const isInEditMode = this._lassoDrawingActive && this._lassoActiveLayer && + this._lassoActiveLayer.editButton === btn; + if (!isInEditMode) { + btn.style.color = '#888'; + } + }; + btn.onclick = (e) => { e.stopPropagation(); callback(); }; + return btn; + }; + + nodeType.prototype.selectRefLayer = function (layer) { + if (this.selectedRefLayer?.domElement) { + this.selectedRefLayer.domElement.style.background = '#262626'; + } + this.selectedRefLayer = layer; + if (layer.domElement) layer.domElement.style.background = '#0c0c0c'; + this.exitLassoMode?.(); + this.updateLayerButtonStates(); + this.forceCanvasRefresh(); + }; + + nodeType.prototype.deselectRefLayer = function () { + if (this.selectedRefLayer?.domElement) { + this.selectedRefLayer.domElement.style.background = '#262626'; + } + this.selectedRefLayer = null; + this.exitLassoMode?.(); + this.updateLayerButtonStates(); + this.forceCanvasRefresh(); + }; + + nodeType.prototype.removeRefLayer = function (layer) { + const idx = this.refLayers.indexOf(layer); + if (idx === -1) return; + + if (this.selectedRefLayer === layer) this.selectedRefLayer = null; + this.refLayers.splice(idx, 1); + if (layer.domElement?.parentNode) layer.domElement.parentNode.removeChild(layer.domElement); + + this.renumberRefLayers(); + this.updateRefDataWidget(); + this.forceCanvasRefresh(); + }; + + nodeType.prototype.renumberRefLayers = function () { + this.refLayers.forEach((l, i) => { + l.value.name = `ref_${i + 1}`; + if (l.domElement) l.domElement.querySelector('span').textContent = l.value.name; + }); + }; + + nodeType.prototype.updateRefDataWidget = function () { + const w = this.widgets.find(w => w.name === "ref_layer_data"); + if (w) { + w.value = this.getRefLayerData(); + + // Also save to sessionStorage for quick restore on page refresh + try { + const sessionKey = `prepare-refs-layers-${this.uuid}`; + safeSetSessionItem(sessionKey, JSON.stringify(w.value)); + } catch (e) { + console.warn('[PrepareRefs] Failed to save layers to session:', e); + } + } + }; + + nodeType.prototype.updateLayerButtonStates = function () { + this.refLayers?.forEach(layer => { + if (layer.editButton) { + // Make pen icon green if this layer is in edit mode + const isInEditMode = this._lassoDrawingActive && this._lassoActiveLayer === layer; + layer.editButton.style.color = isInEditMode ? '#4CAF50' : '#888'; + } + }); + }; + + // ================================================================= + // Dimension & Canvas Size Helpers + // ================================================================= + nodeType.prototype.updateCanvasSizeFromWidgets = function () { + const w = this.widgets.find(w => w.name === 'mask_width')?.value ?? 640; + const h = this.widgets.find(w => w.name === 'mask_height')?.value ?? 480; + + this.refCanvasEditor.setSize(w, h); + this.refsCanvas.style.width = this.refsCanvas.style.height = ''; + this.sizeManager.setCanvasSize(w, h); + this.refreshCanvas(); + this.sizeManager.updateSize(true); + }; + + // ================================================================= + // Image Refresh & Persistence + // ================================================================= + nodeType.prototype.updateReferenceImageFromConnectedNode = async function () { + try { + const base64 = await getReferenceImageFromConnectedNode(this, 'bg_image'); + if (!base64) return this.forceCanvasRefresh(); + + const img = await new Promise((resolve, reject) => { + const i = new Image(); + i.onload = () => { this.refCanvasEditor.loadBackgroundImage(i); resolve(i); }; + i.onerror = reject; + i.src = base64; + }); + + this.loadedBgImage = img; + + // Save via backend + await fetch('/wanvideowrapper_qq/save_prepare_refs_images', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ bg_image: base64, ref_images: [] }) + }) + .then(r => r.json()) + .then(res => { + if (res.success && res.paths?.bg_image_path) { + this.properties.bg_image_path = res.paths.bg_image_path; + safeSetSessionItem(`prepare-refs-bg-${this.uuid}`, res.paths.bg_image_path); + } + }) + .catch(console.error); + + this.forceCanvasRefresh(); + } catch (err) { + console.error('[PrepareRefs] Failed to refresh background:', err); + this.forceCanvasRefresh(); + } + }; + + // ================================================================= + // Session Restore on Load + // ================================================================= + nodeType.prototype.restoreSessionImages = function () { + // Restore layers from sessionStorage + const sessionLayers = safeGetSessionItem(`prepare-refs-layers-${this.uuid}`); + if (sessionLayers) { + try { + const layerData = JSON.parse(sessionLayers); + if (Array.isArray(layerData) && layerData.length > 0) { + this.restoreRefLayers(layerData); + console.log('[PrepareRefs] Restored', layerData.length, 'layers from session'); + } + } catch (e) { + console.warn('[PrepareRefs] Failed to restore layers from session:', e); + } + } + + const bgPath = safeGetSessionItem(`prepare-refs-bg-${this.uuid}`); + const refPathsJson = safeGetSessionItem(`prepare-refs-images-${this.uuid}`); + + if (bgPath) { + try { + // CRITICAL: Resolve relative path correctly using import.meta.url + const url = new URL(bgPath, import.meta.url).href; + loadImageFromUrl(url, this.refCanvasEditor) + .then(() => { + console.log('[PrepareRefs] Restored bg_image from session:', bgPath); + this.forceCanvasRefresh(); + }) + .catch(err => console.error('[PrepareRefs] Failed to restore bg_image:', err)); + } catch (e) { + console.error('[PrepareRefs] Invalid bg_image path in session:', bgPath, e); + } + } + + if (refPathsJson) { + try { + const paths = JSON.parse(refPathsJson); + if (Array.isArray(paths) && paths.length > 0) { + const resolvedPaths = paths.map(p => new URL(p, import.meta.url).href); + loadRefImagesFromPaths(resolvedPaths, this).then(imgs => { + this.loadedRefImages = imgs; + this.forceCanvasRefresh(); + }); + } + } catch (e) { + console.error('[PrepareRefs] Failed to parse/restore ref images from session:', e); + } + } this.updateCanvasSizeFromWidgets(); + }; - // Add ref layer methods - this.addRefLayer = () => { - // Find the highest existing ref number - let maxNum = 0; - this.refLayers.forEach(layer => { - const match = layer.value.name.match(/ref_(\d+)/); - if (match) { - const num = parseInt(match[1], 10); - if (num > maxNum) maxNum = num; - } - }); - const nextNum = maxNum + 1; + // ================================================================= + // Lifecycle Callbacks + // ================================================================= + chainCallback(nodeType.prototype, 'onExecuted', async function (message) { + // Auto-size from incoming image if user hasn't manually adjusted + if (!this.properties.userAdjustedDims && message?.ui?.bg_image_dims?.[0]) { + const { width, height } = message.ui.bg_image_dims[0]; + this.setDimensionValue?.('mask_width', Math.round(width)); + this.setDimensionValue?.('mask_height', Math.round(height)); + this.updateCanvasSizeFromWidgets(); + } - const refWidget = new RefLayerWidget(`ref_${nextNum}`); - refWidget.value = { on: true, name: `ref_${nextNum}` }; - refWidget.node = this; - this.refLayers.push(refWidget); + // Refresh canvas with connected bg_image + await this.updateReferenceImageFromConnectedNode(); - // Create DOM element for the ref layer - const refElement = document.createElement('div'); - refElement.style.display = 'flex'; - refElement.style.alignItems = 'center'; - refElement.style.padding = '4px 6px'; - refElement.style.background = '#262626'; - refElement.style.border = '1px solid #3a3a3a'; - refElement.style.borderRadius = '3px'; - refElement.style.fontSize = '11px'; - refElement.style.color = '#e0e0e0'; - refElement.style.cursor = 'pointer'; - refElement.style.userSelect = 'none'; - refElement.style.transition = 'background 0.1s ease'; + // Restore bg/ref images from execution message + if (message?.ui?.bg_image_path?.[0]) { + const path = message.ui.bg_image_path[0]; + this.properties.bg_image_path = path; + safeSetSessionItem(`prepare-refs-bg-${this.uuid}`, path); + const url = new URL(path, import.meta.url).href; + await loadImageFromUrl(url, this.refCanvasEditor); + } - // Layer name - const nameSpan = document.createElement('span'); - nameSpan.textContent = refWidget.value.name; - nameSpan.style.flex = '1'; - refElement.appendChild(nameSpan); + if (message?.ui?.ref_images_paths?.length) { + const paths = message.ui.ref_images_paths; + this.properties.ref_images_paths = paths; + safeSetSessionItem(`prepare-refs-images-${this.uuid}`, JSON.stringify(paths)); + const resolvedPaths = paths.map(p => new URL(p, import.meta.url).href); + this.loadedRefImages = await loadRefImagesFromPaths(resolvedPaths, this); + } - // Edit button - const editBtn = document.createElement('button'); - editBtn.textContent = '✎'; - editBtn.style.width = '18px'; - editBtn.style.height = '18px'; - editBtn.style.padding = '0'; - editBtn.style.border = 'none'; - editBtn.style.background = 'transparent'; - editBtn.style.color = '#888888'; - editBtn.style.cursor = 'pointer'; - editBtn.style.fontSize = '12px'; - editBtn.style.transition = 'color 0.1s ease'; - editBtn.onmousedown = (e) => { - e.stopPropagation(); - editBtn.style.color = '#ffffff'; - }; - editBtn.onmouseup = () => { - editBtn.style.color = '#888888'; - }; - editBtn.onmouseleave = () => { - editBtn.style.color = '#888888'; - }; - editBtn.onclick = (e) => { - e.stopPropagation(); - // Select this layer when edit is clicked - if (this.selectedRefLayer && this.selectedRefLayer.domElement) { - this.selectedRefLayer.domElement.style.background = '#262626'; - } - this.selectedRefLayer = refWidget; - refElement.style.background = '#0c0c0c'; - - // Toggle lasso mode - if (this._lassoDrawingActive && this._lassoActiveLayer === refWidget) { - this.exitLassoMode?.(); - } else { - this.enterLassoMode?.(refWidget); - } - - this.setDirtyCanvas(true, true); - }; - refElement.appendChild(editBtn); - - // Delete button (X) - const deleteBtn = document.createElement('button'); - deleteBtn.textContent = '✕'; - deleteBtn.style.width = '18px'; - deleteBtn.style.height = '18px'; - deleteBtn.style.padding = '0'; - deleteBtn.style.border = 'none'; - deleteBtn.style.background = 'transparent'; - deleteBtn.style.color = '#888888'; - deleteBtn.style.cursor = 'pointer'; - deleteBtn.style.fontSize = '12px'; - deleteBtn.style.transition = 'color 0.1s ease'; - deleteBtn.onmousedown = (e) => { - e.stopPropagation(); - deleteBtn.style.color = '#ffffff'; - }; - deleteBtn.onmouseup = () => { - deleteBtn.style.color = '#888888'; - }; - deleteBtn.onmouseleave = () => { - deleteBtn.style.color = '#888888'; - }; - deleteBtn.onclick = (e) => { - e.stopPropagation(); - this.removeRefLayer(refWidget, refElement); - }; - refElement.appendChild(deleteBtn); - - // Store reference to element on widget - refWidget.domElement = refElement; - - // Click handler to select/deselect layer - refElement.onclick = () => { - // If clicking on already selected layer, deselect it - if (this.selectedRefLayer === refWidget) { - this.selectedRefLayer = null; - refElement.style.background = '#262626'; - // Exit lasso mode when deselecting - if (this._lassoDrawingActive) { - this.exitLassoMode?.(); - } - this.refreshCanvas(); - this.setDirtyCanvas(true, true); - return; - } - // Deselect previous layer - if (this.selectedRefLayer && this.selectedRefLayer.domElement) { - this.selectedRefLayer.domElement.style.background = '#262626'; - } - // Exit lasso mode when switching layers - if (this._lassoDrawingActive) { - this.exitLassoMode?.(); - } - // Select this layer - this.selectedRefLayer = refWidget; - refElement.style.background = '#0c0c0c'; - this.refreshCanvas(); - this.setDirtyCanvas(true, true); - }; - - this.refLayersContainer.appendChild(refElement); - - // Auto-select the newly added layer - if (this.selectedRefLayer && this.selectedRefLayer.domElement) { - this.selectedRefLayer.domElement.style.background = '#262626'; - } - this.selectedRefLayer = refWidget; - refElement.style.background = '#0c0c0c'; - - // Auto-enter lasso mode for the new layer - if (this.enterLassoMode) { - this.enterLassoMode(refWidget); - } - - // Update the serialized ref layer data widget - const refDataWidget = this.widgets.find(w => w.name === "ref_layer_data"); - if (refDataWidget) { - refDataWidget.value = this.getRefLayerData(); - } - - this.refreshCanvas(); - this.setDirtyCanvas(true, true); - }; - - this.removeRefLayer = (refWidget, refElement) => { - const idx = this.refLayers.indexOf(refWidget); - if (idx > -1) { - // Clear selection if this was the selected layer - if (this.selectedRefLayer === refWidget) { - this.selectedRefLayer = null; - } - - this.refLayers.splice(idx, 1); - // Remove from widgets - const widgetIdx = this.widgets.indexOf(refWidget); - if (widgetIdx > -1) { - this.widgets.splice(widgetIdx, 1); - } - // Remove from DOM - if (refElement && refElement.parentNode) { - refElement.parentNode.removeChild(refElement); - } else if (refWidget.domElement && refWidget.domElement.parentNode) { - refWidget.domElement.parentNode.removeChild(refWidget.domElement); - } - // Renumber remaining layers - this.refLayers.forEach((layer, i) => { - layer.value.name = `ref_${i + 1}`; - if (layer.domElement) { - layer.domElement.querySelector('span').textContent = `ref_${i + 1}`; - } - }); - - // Update the serialized ref layer data widget - const refDataWidget = this.widgets.find(w => w.name === "ref_layer_data"); - if (refDataWidget) { - refDataWidget.value = this.getRefLayerData(); - } - - this.refreshCanvas(); - this.setDirtyCanvas(true, true); - } - }; + this.forceCanvasRefresh(); }); - chainCallback(nodeType.prototype, 'onExecuted', function (message) { - const dims = message?.ui?.bg_image_dims; - if (!this.properties.userAdjustedDims && Array.isArray(dims) && dims.length > 0) { - const first = dims[0]; - if (first?.width && first?.height) { - this.setDimensionValue('mask_width', Math.round(first.width)); - this.setDimensionValue('mask_height', Math.round(first.height)); - this.updateCanvasSizeFromWidgets(); + chainCallback(nodeType.prototype, 'onConfigure', async function (info) { + // Restore ref layers from widget data + const refDataWidget = this.widgets?.find(w => w.name === 'ref_layer_data'); + if (refDataWidget && Array.isArray(refDataWidget.value) && refDataWidget.value.length > 0) { + this.restoreRefLayers(refDataWidget.value); + } else { + // Try to restore from session storage as fallback + const sessionLayers = safeGetSessionItem(`prepare-refs-layers-${this.uuid}`); + if (sessionLayers) { + try { + const layerData = JSON.parse(sessionLayers); + if (Array.isArray(layerData) && layerData.length > 0) { + this.restoreRefLayers(layerData); + } + } catch (e) { + console.warn('[onConfigure] Failed to restore layers from session:', e); + } } } - this.setDirtyCanvas(true, true); + + const savedBgPath = this.properties.bg_image_path; + const savedRefPaths = this.properties.ref_images_paths; + + const sessionBgPath = safeGetSessionItem(`prepare-refs-bg-${this.uuid}`); + const sessionRefPaths = safeGetSessionItem(`prepare-refs-images-${this.uuid}`); + + const finalBgPath = savedBgPath || sessionBgPath; + const finalRefPaths = savedRefPaths || (sessionRefPaths ? JSON.parse(sessionRefPaths) : null); + + if (finalBgPath) { + try { + const url = new URL(finalBgPath, import.meta.url).href; + await loadImageFromUrl(url, this.refCanvasEditor); + console.log('[onConfigure] Restored background image:', finalBgPath); + } catch (e) { + console.error('[onConfigure] Failed to load bg_image:', finalBgPath, e); + } + } + + if (finalRefPaths && Array.isArray(finalRefPaths) && finalRefPaths.length > 0) { + try { + const urls = finalRefPaths.map(p => new URL(p, import.meta.url).href); + this.loadedRefImages = await loadRefImagesFromPaths(urls, this); + console.log('[onConfigure] Restored', this.loadedRefImages.length, 'ref images'); + } catch (e) { + console.error('[onConfigure] Failed to restore ref images:', e); + } + } + + this.forceCanvasRefresh(); }); - // Update the ref layer data widget when node is serialized - const originalOnSerialize = nodeType.prototype.onSerialize; + // Ensure ref data and paths are serialized + const origSerialize = nodeType.prototype.onSerialize; nodeType.prototype.onSerialize = function (o) { - // Update the ref layer data widget before serialization - const refDataWidget = this.widgets?.find(w => w.name === "ref_layer_data"); - if (refDataWidget) { - refDataWidget.value = this.getRefLayerData?.() || []; - } - originalOnSerialize?.apply(this, arguments); + const refWidget = this.widgets?.find(w => w.name === "ref_layer_data"); + if (refWidget) refWidget.value = this.getRefLayerData(); + + o.properties = o.properties || {}; + if (this.properties.bg_image_path) o.properties.bg_image_path = this.properties.bg_image_path; + if (this.properties.ref_images_paths) o.properties.ref_images_paths = this.properties.ref_images_paths; + + origSerialize?.call(this, o); }; }, -}); +}); \ No newline at end of file diff --git a/web/power_spline_editor/trigger_ref_refresh.js b/web/power_spline_editor/trigger_ref_refresh.js new file mode 100644 index 0000000..f80362d --- /dev/null +++ b/web/power_spline_editor/trigger_ref_refresh.js @@ -0,0 +1,99 @@ +/** + * Frontend trigger for PrepareRefs backend processing + * + * This module provides functionality to trigger PrepareRefs processing + * from the frontend without running the full ComfyUI workflow. + */ + +import { getReferenceImageFromConnectedNode } from './graph_query.js'; + +/** + * Trigger PrepareRefs processing in the backend + * + * Sends lasso layer data and bg_image to backend for processing. + * All ref images are exported as 768x768 PNG with alpha channel. + * + * @param {Object} node - The PrepareRefs node instance + * @returns {Promise} Response with success status and file paths + */ +export async function triggerPrepareRefsBackend(node) { + try { + console.log('[triggerPrepareRefsBackend] Starting backend trigger...'); + + // 1. Get current bg_image from connected node or canvas + let bgImageBase64 = null; + + // Try to get from connected node first + bgImageBase64 = await getReferenceImageFromConnectedNode(node, 'bg_image'); + + // Fallback: get from canvas if available + if (!bgImageBase64 && node.refCanvasEditor?.backgroundImage) { + const canvas = document.createElement('canvas'); + const img = node.refCanvasEditor.backgroundImage; + canvas.width = img.width; + canvas.height = img.height; + const ctx = canvas.getContext('2d'); + ctx.drawImage(img, 0, 0); + bgImageBase64 = canvas.toDataURL('image/png'); + console.log('[triggerPrepareRefsBackend] Using bg_image from canvas'); + } + + if (!bgImageBase64) { + console.warn('[triggerPrepareRefsBackend] No bg_image found'); + return { success: false, error: 'No background image available' }; + } + + // 2. Get ref_layer_data from node + const refLayerData = node.getRefLayerData?.() || []; + + if (refLayerData.length === 0) { + console.warn('[triggerPrepareRefsBackend] No ref layers with shapes found'); + return { success: false, error: 'No ref layers with shapes to process' }; + } + + console.log('[triggerPrepareRefsBackend] Found', refLayerData.length, 'layers with shapes'); + + // 3. Get dimensions + const widthWidget = node.widgets?.find(w => w.name === 'mask_width'); + const heightWidget = node.widgets?.find(w => w.name === 'mask_height'); + const maskWidth = widthWidget?.value || 640; + const maskHeight = heightWidget?.value || 480; + + console.log('[triggerPrepareRefsBackend] Dimensions:', maskWidth, 'x', maskHeight); + + // 4. Prepare payload + const payload = { + bg_image: bgImageBase64, + ref_layer_data: refLayerData, + mask_width: maskWidth, + mask_height: maskHeight + }; + + console.log('[triggerPrepareRefsBackend] Sending request with', refLayerData.length, 'layers'); + + // 5. Call backend endpoint + const response = await fetch('/wanvideowrapper_qq/trigger_prepare_refs', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(payload) + }); + + const result = await response.json(); + + if (!result.success) { + throw new Error(result.error || 'Backend processing failed'); + } + + console.log('[triggerPrepareRefsBackend] Success:', result.message); + console.log('[triggerPrepareRefsBackend] Generated files:', result.paths); + + return result; + + } catch (error) { + console.error('[triggerPrepareRefsBackend] Error:', error); + return { + success: false, + error: error.message || 'Unknown error during backend trigger' + }; + } +}