196 lines
8.2 KiB
Python
196 lines
8.2 KiB
Python
"""LayerForge HTTP and WebSocket route registration."""
|
|
|
|
import os
|
|
import time
|
|
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
from server import PromptServer
|
|
|
|
from .image_serialization import file_to_data_url, pil_to_data_url
|
|
from .matting.api import register_matting_routes
|
|
from .node import log
|
|
|
|
|
|
def register_routes(node_class):
|
|
"""Register the legacy LayerForge routes and the matting routes."""
|
|
|
|
@PromptServer.instance.routes.get("/layerforge/canvas_ws")
|
|
async def handle_canvas_websocket(request):
|
|
ws = web.WebSocketResponse(max_msg_size=33554432)
|
|
await ws.prepare(request)
|
|
|
|
async for msg in ws:
|
|
if msg.type == web.WSMsgType.TEXT:
|
|
try:
|
|
data = msg.json()
|
|
node_id = data.get("nodeId")
|
|
if not node_id:
|
|
await ws.send_json({"status": "error", "message": "nodeId is required"})
|
|
continue
|
|
|
|
with node_class._storage_lock:
|
|
node_class._canvas_data_storage[node_id] = {
|
|
"image": data.get("image"),
|
|
"mask": data.get("mask"),
|
|
"timestamp": time.time(),
|
|
}
|
|
|
|
log.info(f"Received canvas data for node {node_id} via WebSocket")
|
|
await ws.send_json({"type": "ack", "nodeId": node_id, "status": "success"})
|
|
log.debug(f"Sent ACK for node {node_id}")
|
|
except Exception as error:
|
|
log.error(f"Error processing WebSocket message: {error}")
|
|
await ws.send_json({"status": "error", "message": str(error)})
|
|
elif msg.type == web.WSMsgType.ERROR:
|
|
log.error(f"WebSocket connection closed with exception {ws.exception()}")
|
|
|
|
log.info("WebSocket connection closed")
|
|
return ws
|
|
|
|
@PromptServer.instance.routes.get("/layerforge/get_input_data/{node_id}")
|
|
async def get_input_data(request):
|
|
try:
|
|
node_id = request.match_info["node_id"]
|
|
log.debug(f"Checking for input data for node: {node_id}")
|
|
with node_class._storage_lock:
|
|
input_data = node_class._canvas_data_storage.get(f"{node_id}_input")
|
|
|
|
if input_data:
|
|
log.info(f"Input data found for node {node_id}, sending to frontend")
|
|
return web.json_response({"success": True, "has_input": True, "data": input_data})
|
|
|
|
log.debug(f"No input data found for node {node_id}")
|
|
return web.json_response({"success": True, "has_input": False})
|
|
except Exception as error:
|
|
log.error(f"Error in get_input_data: {error}")
|
|
return web.json_response({"success": False, "error": str(error)}, status=500)
|
|
|
|
@PromptServer.instance.routes.post("/layerforge/clear_input_data/{node_id}")
|
|
async def clear_input_data(request):
|
|
try:
|
|
node_id = request.match_info["node_id"]
|
|
log.info(f"Clearing input data for node: {node_id}")
|
|
with node_class._storage_lock:
|
|
input_key = f"{node_id}_input"
|
|
if input_key in node_class._canvas_data_storage:
|
|
del node_class._canvas_data_storage[input_key]
|
|
log.info(f"Input data cleared for node {node_id}")
|
|
else:
|
|
log.debug(f"No input data to clear for node {node_id}")
|
|
|
|
return web.json_response({
|
|
"success": True,
|
|
"message": f"Input data cleared for node {node_id}",
|
|
})
|
|
except Exception as error:
|
|
log.error(f"Error in clear_input_data: {error}")
|
|
return web.json_response({"success": False, "error": str(error)}, status=500)
|
|
|
|
@PromptServer.instance.routes.get("/ycnode/get_canvas_data/{node_id}")
|
|
async def get_canvas_data(request):
|
|
del request
|
|
try:
|
|
cache_data = node_class._canvas_cache
|
|
response_data = {
|
|
"success": True,
|
|
"data": {"image": None, "mask": None},
|
|
}
|
|
|
|
if cache_data["image"] is not None:
|
|
response_data["data"]["image"] = pil_to_data_url(cache_data["image"])
|
|
|
|
if cache_data["mask"] is not None:
|
|
response_data["data"]["mask"] = pil_to_data_url(cache_data["mask"])
|
|
|
|
return web.json_response(response_data)
|
|
except Exception as error:
|
|
log.error(f"Error in get_canvas_data: {error}")
|
|
return web.json_response({"success": False, "error": str(error)})
|
|
|
|
@PromptServer.instance.routes.get("/layerforge/get-latest-images/{since}")
|
|
async def get_latest_images_route(request):
|
|
try:
|
|
since_timestamp = float(request.match_info.get("since", 0))
|
|
latest_image_paths = node_class.get_latest_images(since_timestamp / 1000.0)
|
|
images_data = []
|
|
for image_path in latest_image_paths:
|
|
images_data.append(file_to_data_url(image_path))
|
|
|
|
return web.json_response({"success": True, "images": images_data})
|
|
except Exception as error:
|
|
log.error(f"Error in get_latest_images_route: {error}")
|
|
return web.json_response({"success": False, "error": str(error)}, status=500)
|
|
|
|
@PromptServer.instance.routes.get("/ycnode/get_latest_image")
|
|
async def get_latest_image_route(request):
|
|
del request
|
|
try:
|
|
latest_image_path = node_class.get_latest_image()
|
|
if latest_image_path:
|
|
return web.json_response({
|
|
"success": True,
|
|
"image_data": file_to_data_url(latest_image_path),
|
|
})
|
|
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": "No images found in output directory.",
|
|
}, status=404)
|
|
except Exception as error:
|
|
return web.json_response({"success": False, "error": str(error)}, status=500)
|
|
|
|
@PromptServer.instance.routes.post("/ycnode/load_image_from_path")
|
|
async def load_image_from_path_route(request):
|
|
try:
|
|
data = await request.json()
|
|
file_path = data.get("file_path")
|
|
if not file_path:
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": "file_path is required",
|
|
}, status=400)
|
|
|
|
log.info(f"Attempting to load image from path: {file_path}")
|
|
if not os.path.exists(file_path):
|
|
log.warning(f"File not found: {file_path}")
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": f"File not found: {file_path}",
|
|
}, status=404)
|
|
|
|
valid_extensions = (
|
|
".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff", ".tif", ".ico", ".avif"
|
|
)
|
|
if not file_path.lower().endswith(valid_extensions):
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": f"Invalid image file extension. Supported: {valid_extensions}",
|
|
}, status=400)
|
|
|
|
try:
|
|
with Image.open(file_path) as image:
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
log.info(f"Successfully loaded image from path: {file_path}")
|
|
return web.json_response({
|
|
"success": True,
|
|
"image_data": pil_to_data_url(image),
|
|
"width": image.width,
|
|
"height": image.height,
|
|
})
|
|
except Exception as image_error:
|
|
log.error(f"Error processing image file {file_path}: {image_error}")
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": f"Error processing image file: {image_error}",
|
|
}, status=500)
|
|
except Exception as error:
|
|
log.error(f"Error in load_image_from_path_route: {error}")
|
|
return web.json_response({"success": False, "error": str(error)}, status=500)
|
|
|
|
register_matting_routes()
|
|
|
|
|
|
__all__ = ["register_routes"]
|