import base64 import json import os import sys from datetime import datetime from io import BytesIO from pathlib import Path from typing import Optional, Tuple, List, Dict, Any import cv2 import folder_paths import subprocess from aiohttp import web from PIL import Image from server import PromptServer # Import SAM2 masker for point-click masking from .nodes.sam2_masker import sam2_masker from .nodes.prepare_refs import convert_mask_to_contour # Constants IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp'} LORA_EXTENSIONS = {'.safetensors', '.ckpt', '.pt'} PREVIEW_SUFFIXES = ["_01", "_02", "_03", "_thumb"] # Endpoint prefix constant for consistency API_PREFIX = "/wanvideowrapper_qq" def _get_bg_folder_path() -> str: """Get the path to the bg folder where A.jpg and bg_image.png are stored.""" current_dir = Path(__file__).parent bg_path = current_dir / "web" / "power_spline_editor" / "bg" bg_path.mkdir(parents=True, exist_ok=True) return str(bg_path) def _get_ref_folder_path() -> str: """Get the path to the ref folder where reference images are stored.""" current_dir = Path(__file__).parent ref_path = current_dir / "web" / "power_spline_editor" / "ref" ref_path.mkdir(parents=True, exist_ok=True) return str(ref_path) def _decode_base64_image(image_data: str) -> Tuple[Image.Image, str]: if image_data.startswith('data:image'): header, encoded = image_data.split(',', 1) image_format = header.split('/')[1].split(';')[0] else: encoded = image_data # Try to infer format from data or default to png image_format = 'png' try: image_bytes = base64.b64decode(encoded) img = Image.open(BytesIO(image_bytes)) img.verify() # Reopen after verification since verify() closes the file img = Image.open(BytesIO(image_bytes)) return img, image_format except Exception as e: raise ValueError(f"Invalid image data: {e}") def _save_image_with_format(img: Image.Image, filepath: str, image_format: Optional[str] = None) -> None: fmt = (image_format or "png").upper() save_kwargs = {} if fmt in ("JPG", "JPEG"): img = img.convert('RGB') fmt = "JPEG" save_kwargs["quality"] = 95 else: # Preserve alpha if present if img.mode not in ("RGBA", "LA"): img = img.convert("RGBA") fmt = "PNG" img.save(filepath, format=fmt, **save_kwargs) def _get_lora_full_path(lora_name: str) -> Optional[str]: try: lora_folders = folder_paths.get_folder_paths("loras") for folder in lora_folders: # Check exact name first lora_path = os.path.join(folder, lora_name) if os.path.exists(lora_path): return os.path.abspath(lora_path) # Check with common extensions if not provided if not any(lora_name.lower().endswith(ext) for ext in LORA_EXTENSIONS): for ext in LORA_EXTENSIONS: lora_path_with_ext = os.path.join(folder, lora_name + ext) if os.path.exists(lora_path_with_ext): return os.path.abspath(lora_path_with_ext) return None except Exception as e: print(f"Error finding LoRA path: {e}") return None def _rename_preview_files(old_base_name: str, new_name: str, preview_folder: str) -> List[str]: renamed_files = [] try: # Rename JSON file old_json_path = os.path.join(preview_folder, f"{old_base_name}.json") new_json_path = os.path.join(preview_folder, f"{new_name}.json") if os.path.exists(old_json_path): try: os.rename(old_json_path, new_json_path) renamed_files.append(new_json_path) print(f"Renamed JSON: {old_json_path} -> {new_json_path}") except Exception as e: print(f"Warning: Failed to rename JSON file: {e}") # Rename preview images with suffixes for suffix in PREVIEW_SUFFIXES: for ext in IMAGE_EXTENSIONS: old_img_path = os.path.join(preview_folder, f"{old_base_name}{suffix}{ext}") new_img_path = os.path.join(preview_folder, f"{new_name}{suffix}{ext}") if os.path.exists(old_img_path): try: os.rename(old_img_path, new_img_path) renamed_files.append(new_img_path) print(f"Renamed image: {old_img_path} -> {new_img_path}") except Exception as e: print(f"Warning: Failed to rename image file: {e}") # Check for images without suffix (backward compatibility) for ext in IMAGE_EXTENSIONS: old_img_path = os.path.join(preview_folder, f"{old_base_name}{ext}") new_img_path = os.path.join(preview_folder, f"{new_name}{ext}") if os.path.exists(old_img_path): try: os.rename(old_img_path, new_img_path) renamed_files.append(new_img_path) print(f"Renamed image: {old_img_path} -> {new_img_path}") except Exception as e: print(f"Warning: Failed to rename image file: {e}") except Exception as e: print(f"Error in _rename_preview_files: {e}") return renamed_files def _open_explorer(path: str, select_file: bool = True) -> Dict[str, Any]: if not sys.platform.startswith('win'): return {"error": "This feature is only available on Windows", "status": 400} try: if select_file and os.path.isfile(path): subprocess.run(['explorer', '/select,', path], check=True, shell=False) return {"success": True, "path": path} else: folder_path = path if os.path.isdir(path) else os.path.dirname(path) subprocess.run(['explorer', folder_path], check=True, shell=False) return {"success": True, "path": folder_path, "message": "Opened folder instead of selecting file"} except subprocess.CalledProcessError as e: # Fallback: try opening just the directory try: if os.path.isfile(path): path = os.path.dirname(path) subprocess.run(['explorer', path], check=True, shell=False) return {"success": True, "path": path} except subprocess.CalledProcessError as e2: return {"error": f"Failed to open Explorer: {str(e2)}", "status": 500} @PromptServer.instance.routes.post(f"{API_PREFIX}/save_ref_image") async def save_ref_image(request) -> web.Response: try: 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 and validate image try: img, image_format = _decode_base64_image(image_data) except ValueError as e: return web.json_response({"error": str(e)}, status=400) bg_folder = _get_bg_folder_path() ref_image_path = os.path.join(bg_folder, safe_name) # Save with appropriate format handling _save_image_with_format(img, ref_image_path, image_format) return web.json_response({"success": True, "path": ref_image_path, "message": "Reference image saved successfully"}) except Exception as e: print(f"Error saving ref image: {e}") import traceback traceback.print_exc() return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post(f"{API_PREFIX}/save_prepare_refs_images") async def save_prepare_refs_images(request) -> web.Response: 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: " f"{len(ref_images_data) if isinstance(ref_images_data, list) else 0}") saved_paths = { "bg_image_path": None, "ref_images_paths": [] } ref_folder = _get_ref_folder_path() print(f"[API save_prepare_refs_images] Using ref folder: {ref_folder}") # Save background image to REF folder if bg_image_data: try: print(f"[API save_prepare_refs_images] Saving bg_image to REF folder...") img, _ = _decode_base64_image(bg_image_data) 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...") img, _ = _decode_base64_image(ref_image_data) 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) @PromptServer.instance.routes.post("/wanvideo_wrapper/open_explorer") async def open_explorer_handler(request) -> web.Response: try: post = await request.post() data = await request.json() if request.content_type == 'application/json' else {} lora_name = data.get("lora_name") or post.get("lora_name") if not lora_name: return web.json_response({"error": "No LoRA name provided"}, status=400) lora_path = _get_lora_full_path(lora_name) if not lora_path: return web.json_response({"error": f"LoRA file not found: {lora_name}"}, status=404) result = _open_explorer(lora_path, select_file=True) if "error" in result: return web.json_response(result, status=result.get("status", 500)) return web.json_response({"success": True, **result}) except Exception as e: print(f"Error opening Explorer: {e}") import traceback traceback.print_exc() return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/power_load_video/open_input_dir") async def open_input_dir(request) -> web.Response: try: if not sys.platform.startswith('win'): return web.json_response({"error": "This feature is only available on Windows"}, status=400) input_dir = folder_paths.get_input_directory() if not os.path.isdir(input_dir): return web.json_response({"error": f"Input directory not found: {input_dir}"}, status=404) subprocess.run(['explorer', input_dir], check=True, shell=False) return web.json_response({"success": True, "path": input_dir}) except subprocess.CalledProcessError as e: return web.json_response({"error": f"Failed to open Explorer: {str(e)}"}, status=500) except Exception as e: return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/power_load_video/check_workflow") async def check_video_workflow(request) -> web.Response: try: data = await request.json() filename = data.get("filename", "") subfolder = data.get("subfolder", "") input_dir = folder_paths.get_input_directory() file_path = os.path.join(input_dir, subfolder, filename) if subfolder else \ os.path.join(input_dir, filename) if not os.path.isfile(file_path): return web.json_response({"has_workflow": False}) # Extract metadata via ffprobe/ffmpeg try: result = subprocess.run( ["ffprobe", "-show_format", "-print_format", "json", file_path], capture_output=True, text=True, timeout=5 ) if result.returncode == 0: probe = json.loads(result.stdout) tags = probe.get("format", {}).get("tags", {}) workflow = None # Check for 'workflow' tag (VHS VideoCombine) if "workflow" in tags: try: workflow = json.loads(tags["workflow"]) except (json.JSONDecodeError, TypeError): pass # Check for 'prompt' tag and reconstruct workflow if workflow is None and "prompt" in tags: try: prompt = json.loads(tags["prompt"]) workflow = {"prompt": prompt} except (json.JSONDecodeError, TypeError): pass if workflow is not None: return web.json_response({"has_workflow": True, "workflow": workflow}) except Exception: pass return web.json_response({"has_workflow": False}) except Exception as e: return web.json_response({"has_workflow": False, "error": str(e)}) @PromptServer.instance.routes.post("/wanvideo_wrapper/get_lora_path") async def get_lora_path(request) -> web.Response: try: post = await request.post() data = await request.json() if request.content_type == 'application/json' else {} lora_name = data.get("lora_name") or post.get("lora_name") if not lora_name: return web.json_response({"error": "No LoRA name provided"}, status=400) lora_path = _get_lora_full_path(lora_name) if not lora_path: return web.json_response({"error": f"LoRA file not found: {lora_name}"}, status=404) return web.json_response({"success": True, "path": lora_path}) except Exception as e: print(f"Error getting LoRA path: {e}") import traceback traceback.print_exc() return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/wanvideo_wrapper/rename_lora") async def rename_lora(request) -> web.Response: try: post = await request.post() data = await request.json() if request.content_type == 'application/json' else {} old_name = data.get("old_name") or post.get("old_name") new_name = data.get("new_name") or post.get("new_name") variant = data.get("variant") or post.get("variant", "high") if not old_name: return web.json_response({"error": "No old LoRA name provided"}, status=400) if not new_name: return web.json_response({"error": "No new LoRA name provided"}, status=400) old_lora_path = _get_lora_full_path(old_name) if not old_lora_path: return web.json_response({"error": f"LoRA file not found: {old_name}"}, status=404) lora_dir = os.path.dirname(old_lora_path) old_base_name = os.path.splitext(os.path.basename(old_lora_path))[0] lora_ext = os.path.splitext(old_lora_path)[1] new_lora_path = os.path.join(lora_dir, f"{new_name}{lora_ext}") if os.path.exists(new_lora_path): return web.json_response({"error": f"A file with the name '{new_name}{lora_ext}' already exists"}, status=409) try: os.rename(old_lora_path, new_lora_path) print(f"Renamed LoRA: {old_lora_path} -> {new_lora_path}") except Exception as e: return web.json_response({"error": f"Failed to rename LoRA file: {str(e)}"}, status=500) renamed_files = [new_lora_path] # Determine if this is a model or LoRA based on file path is_model = any(old_lora_path.startswith(folder) for folder in folder_paths.get_folder_paths("checkpoints")) folder_type = "checkpoints" if is_model else "loras" folders_to_check = folder_paths.get_folder_paths(folder_type) for folder in folders_to_check: preview_folder = os.path.join(folder, "_power_preview") if os.path.exists(preview_folder): renamed_files.extend(_rename_preview_files(old_base_name, new_name, preview_folder)) return web.json_response({ "success": True, "message": f"Successfully renamed LoRA and {len(renamed_files)-1} associated files", "new_high_path": new_lora_path, "new_low_path": new_lora_path, # For high variant, low path is the same "renamed_files": renamed_files }) except Exception as e: print(f"Error renaming LoRA: {e}") import traceback traceback.print_exc() return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post(f"{API_PREFIX}/extract_first_frame") async def extract_first_frame(request) -> web.Response: try: data = await request.json() video_path = data.get("video_path") if not video_path: return web.json_response({"success": False, "error": "No video path provided"}, status=400) print(f"[API extract_first_frame] Extracting first frame from: {video_path}") # Find the video file if os.path.exists(video_path): video_file_path = video_path else: input_folders = folder_paths.get_folder_paths("input") video_file_path = None for folder in input_folders: potential_path = os.path.join(folder, video_path) if os.path.exists(potential_path): video_file_path = potential_path break if not video_file_path: return web.json_response({"success": False, "error": f"Video file not found: {video_path}"}, status=404) print(f"[API extract_first_frame] Found video at: {video_file_path}") # Extract first frame using OpenCV with proper resource management cap = None try: cap = cv2.VideoCapture(video_file_path) if not cap.isOpened(): return web.json_response({"success": False, "error": "Failed to open video file"}, status=500) ret, frame = cap.read() if not ret: return web.json_response({"success": False, "error": "Failed to read first frame from video"}, status=500) # Convert BGR to RGB frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert to PIL Image img = Image.fromarray(frame_rgb) # Convert to base64 buffered = BytesIO() img.save(buffered, format="PNG") img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8') img_data_url = f"data:image/png;base64,{img_base64}" print(f"[API extract_first_frame] Successfully extracted first frame") return web.json_response({ "success": True, "image": img_data_url }) finally: if cap is not None and cap.isOpened(): cap.release() except Exception as e: print(f"[API extract_first_frame] ERROR: {e}") import traceback traceback.print_exc() return web.json_response({"success": False, "error": str(e)}, status=500) @PromptServer.instance.routes.post(f"{API_PREFIX}/sam2_predict") async def sam2_predict(request) -> web.Response: try: data = await request.json() if not data.get("image"): return web.json_response({"success": False, "error": "Missing image parameter"}, status=400) if not data.get("points") or len(data["points"]) == 0: return web.json_response({"success": False, "error": "Missing points parameter"}, status=400) # Decode base64 image img, _ = _decode_base64_image(data["image"]) # Extract points points = data["points"] point_coords = [[p["x"], p["y"]] for p in points] point_labels = [p.get("label", 1) for p in points] # Default to foreground # Predict mask using SAM2 mask, score = sam2_masker.predict_from_points(img, point_coords, point_labels) # Convert mask to contour path path = convert_mask_to_contour(mask) return web.json_response({ "success": True, "path": path, "score": score, "mask_shape": list(mask.shape) if hasattr(mask, 'shape') else None }) except Exception as e: import traceback traceback.print_exc() return web.json_response({"success": False, "error": str(e)}, status=500) @PromptServer.instance.routes.get(f"{API_PREFIX}/video_metadata") async def get_video_metadata(request) -> web.Response: filename = request.query.get("filename") if not filename: return web.json_response({"success": False, "error": "Missing filename"}, status=400) try: filepath = folder_paths.get_annotated_filepath(filename) except Exception: input_dir = folder_paths.get_input_directory() filepath = os.path.join(input_dir, filename) if not os.path.exists(filepath): return web.json_response({"success": False, "error": "File not found"}, status=404) cap = None try: cap = cv2.VideoCapture(filepath) if not cap.isOpened(): return web.json_response({"success": False, "error": "Failed to open video file"}, status=500) fps = cap.get(cv2.CAP_PROP_FPS) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) if fps <= 0: fps = 24.0 return web.json_response({ "success": True, "fps": float(fps), "frame_count": frame_count, "width": width, "height": height, }) except Exception as e: return web.json_response({"success": False, "error": str(e)}, status=500) finally: if cap is not None and cap.isOpened(): cap.release()