308 lines
12 KiB
Python
308 lines
12 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 os
|
|
import folder_paths
|
|
from datetime import datetime
|
|
|
|
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)
|
|
|
|
@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)
|
|
|
|
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
|
|
|