535 lines
21 KiB
Python
535 lines
21 KiB
Python
import json
|
|
import base64
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
import os
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import folder_paths
|
|
import subprocess
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
import os
|
|
import folder_paths
|
|
from datetime import datetime
|
|
|
|
# Import SAM2 masker for point-click masking
|
|
from .nodes.sam2_masker import sam2_masker
|
|
from .nodes.prepare_refs import convert_mask_to_contour
|
|
|
|
def get_bg_folder_path():
|
|
"""Get the path to the bg folder where A.jpg and bg_image.png are stored"""
|
|
# Look for the bg folder relative to this file
|
|
import pathlib
|
|
current_dir = pathlib.Path(__file__).parent
|
|
bg_path = current_dir / "web" / "power_spline_editor" / "bg"
|
|
|
|
# Create the directory if it doesn't exist
|
|
bg_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
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"""
|
|
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 base64 image data
|
|
if image_data.startswith('data:image'):
|
|
# Remove data URL prefix (e.g., "data:image/png;base64,")
|
|
header, encoded = image_data.split(',', 1)
|
|
image_format = header.split('/')[1].split(';')[0] # Extract format like 'png' or 'jpeg'
|
|
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))
|
|
img.verify() # Verify it's a valid image
|
|
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 = {}
|
|
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(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:
|
|
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("/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:
|
|
# Get the LoRA folder paths from ComfyUI
|
|
lora_folders = folder_paths.get_folder_paths("loras")
|
|
|
|
# Search for the LoRA file in all LoRA folders
|
|
for folder in lora_folders:
|
|
lora_path = os.path.join(folder, lora_name)
|
|
if os.path.exists(lora_path):
|
|
return os.path.abspath(lora_path)
|
|
|
|
# Also check for common extensions if not provided
|
|
if not lora_name.lower().endswith(('.safetensors', '.ckpt', '.pt')):
|
|
for ext in ['.safetensors', '.ckpt', '.pt']:
|
|
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
|
|
|
|
@PromptServer.instance.routes.post("/wanvideo_wrapper/open_explorer")
|
|
async def open_explorer(request):
|
|
"""API endpoint to open Windows Explorer at the LoRA file location"""
|
|
try:
|
|
if not sys.platform.startswith('win'):
|
|
return web.json_response({"error": "This feature is only available on Windows"}, status=400)
|
|
|
|
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)
|
|
|
|
# Get the full path to the LoRA file
|
|
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)
|
|
|
|
# Open Windows Explorer and select the file
|
|
try:
|
|
# Use start command to open Explorer with the file selected
|
|
subprocess.run(['explorer', '/select,', lora_path], check=True, shell=False)
|
|
return web.json_response({"success": True, "path": lora_path})
|
|
except subprocess.CalledProcessError as e:
|
|
# Fallback: open the folder if selecting the file fails
|
|
try:
|
|
folder_path = os.path.dirname(lora_path)
|
|
subprocess.run(['explorer', folder_path], check=True, shell=False)
|
|
return web.json_response({"success": True, "path": folder_path, "message": "Opened folder instead of selecting file"})
|
|
except subprocess.CalledProcessError as e2:
|
|
return web.json_response({"error": f"Failed to open Explorer: {str(e2)}"}, status=500)
|
|
|
|
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("/wanvideo_wrapper/get_lora_path")
|
|
async def get_lora_path(request):
|
|
"""API endpoint to get the full path of a LoRA file"""
|
|
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)
|
|
|
|
# Get the full path to the LoRA file
|
|
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):
|
|
"""API endpoint to rename a LoRA file and its associated files"""
|
|
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)
|
|
|
|
# Get the full path to the LoRA file
|
|
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)
|
|
|
|
# Get the directory and file extension
|
|
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]
|
|
|
|
# Create the new file paths
|
|
new_lora_path = os.path.join(lora_dir, f"{new_name}{lora_ext}")
|
|
|
|
# Check if new file already exists
|
|
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)
|
|
|
|
# Rename the main LoRA file
|
|
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)
|
|
|
|
# Rename associated files (JSON and preview images)
|
|
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"))
|
|
|
|
if is_model:
|
|
# Handle model files - check checkpoints directory
|
|
model_folders = folder_paths.get_folder_paths("checkpoints")
|
|
for folder in model_folders:
|
|
# Look for _power_preview subfolder
|
|
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))
|
|
else:
|
|
# Handle LoRA files - check loras directory
|
|
lora_folders = folder_paths.get_folder_paths("loras")
|
|
for folder in lora_folders:
|
|
# Look for _power_preview subfolder
|
|
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)
|
|
|
|
def _rename_preview_files(old_base_name, new_name, preview_folder):
|
|
"""Helper function to rename preview files (JSON and images) in a preview folder."""
|
|
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 various suffixes)
|
|
for suffix in ["_01", "_02", "_03", "_thumb"]:
|
|
for ext in [".jpg", ".jpeg", ".png", ".webp"]:
|
|
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}")
|
|
|
|
# Also check for image without suffix (backward compatibility)
|
|
for ext in [".jpg", ".jpeg", ".png", ".webp"]:
|
|
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
|
|
|
|
@PromptServer.instance.routes.post("/wanvideowrapper_qq/extract_first_frame")
|
|
async def extract_first_frame(request):
|
|
"""API endpoint to extract the first frame from a video file"""
|
|
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}")
|
|
|
|
# Try to find the video file
|
|
video_file_path = None
|
|
|
|
# Check if it's already an absolute path
|
|
if os.path.exists(video_path):
|
|
video_file_path = video_path
|
|
else:
|
|
# Search in input directories
|
|
input_folders = folder_paths.get_folder_paths("input")
|
|
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
|
|
cap = cv2.VideoCapture(video_file_path)
|
|
|
|
if not cap.isOpened():
|
|
return web.json_response({"success": False, "error": "Failed to open video file"}, status=500)
|
|
|
|
# Read first frame
|
|
ret, frame = cap.read()
|
|
cap.release()
|
|
|
|
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
|
|
})
|
|
|
|
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("/wanvideowrapper_qq/sam2_predict")
|
|
async def sam2_predict(request):
|
|
"""
|
|
API endpoint for SAM2 point-click prediction.
|
|
|
|
Expects JSON with:
|
|
image: base64 encoded image (data URL format)
|
|
points: list of {'x': float, 'y': float, 'label': int} points
|
|
x, y are in pixel coordinates, label is 1 (foreground) or 0 (background)
|
|
|
|
Returns JSON with:
|
|
success: bool
|
|
path: list of {'x': float, 'y': float} normalized coordinates (0-1)
|
|
score: float (confidence score)
|
|
mask_shape: [height, width] of the predicted mask
|
|
"""
|
|
try:
|
|
data = await request.json()
|
|
|
|
# Validate required parameters
|
|
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
|
|
image_data = data["image"]
|
|
if image_data.startswith('data:image'):
|
|
header, encoded = image_data.split(',', 1)
|
|
else:
|
|
encoded = image_data
|
|
|
|
image_bytes = base64.b64decode(encoded)
|
|
img = Image.open(BytesIO(image_bytes))
|
|
|
|
# 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": mask.shape
|
|
})
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.json_response({"success": False, "error": str(e)}, status=500)
|
|
|