From 61c8200443e826efc74e6a416b76be6e941eb558 Mon Sep 17 00:00:00 2001 From: Bryce Date: Fri, 27 Mar 2026 17:17:19 -0700 Subject: [PATCH] Fix ORA editor mask preview: server-side compositing with simple toggle - Add /api/image/masked endpoint that applies mask as alpha channel - Simplify mask preview modal: just toggle between 'With Background' and 'Masked Only' - Remove complex CSS mask/blend mode approach - Server returns pre-composited masked image (transparent where mask is black) --- tools/ora_editor/app.py | 512 ++++++++++---- tools/ora_editor/templates/editor.html | 932 +++++++++++++------------ 2 files changed, 861 insertions(+), 583 deletions(-) diff --git a/tools/ora_editor/app.py b/tools/ora_editor/app.py index 99ddd32..df3682c 100644 --- a/tools/ora_editor/app.py +++ b/tools/ora_editor/app.py @@ -1,16 +1,19 @@ #!/usr/bin/env python3 """Flask web application for ORA editing.""" +import base64 import io +import json import os +import shutil import sys +import threading import time import zipfile from datetime import datetime from pathlib import Path from xml.etree import ElementTree as ET -try: from flask import ( Flask, request, jsonify, send_file, Response, render_template, make_response @@ -19,15 +22,23 @@ from flask import ( from PIL import Image, ImageDraw import logging import urllib.parse +import urllib.request -# Configure logging +# Configure logging first (before imports that might need logger) logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) -# Configure paths +# Configuration from environment variables APP_DIR = Path(__file__).parent PROJECT_ROOT = Path("/home/noti/dev/ai-game-2") TEMP_DIR = APP_DIR / "temp" + +# ComfyUI configuration via environment variables +COMFYUI_BASE_URL = os.environ.get('COMFYUI_BASE_URL', '127.0.0.1:8188') + +logger.info(f"[CONFIG] COMFYUI_BASE_URL={COMFYUI_BASE_URL}") + +# Ensure temp directory exists TEMP_DIR.mkdir(exist_ok=True) # Import ora operations from sibling directory @@ -41,6 +52,10 @@ from ora_editor.ora_ops import ( # In-memory storage for polygon points (per-request basis) _polygon_storage = {} +# Webhook response holder (single request at a time per user spec) +_webhook_response = None +_webhook_ready = threading.Event() + app = Flask(__name__, template_folder='templates') @@ -343,9 +358,98 @@ def api_polygon_clear(): return jsonify({'success': True}) - @app.route('/api/mask/extract', methods=['POST']) +@app.route('/api/webhook/comfyui', methods=['POST']) +def api_webhook_comfyui(): + """Webhook endpoint for ComfyUI to post completed mask images.""" + logger.info("[WEBHOOK] Received webhook from ComfyUI") + + global _webhook_response, _webhook_ready + + # Reset the event in case we're reusing + _webhook_response = None + _webhook_ready.clear() + + try: + # Get JSON metadata if present (from form field or request body) + metadata = request.form.get('metadata', '') if request.form else '' + external_uid = request.form.get('external_uid', '') if request.form else '' + + logger.info(f"[WEBHOOK] Metadata: {metadata}") + logger.info(f"[WEBHOOK] External UID: {external_uid}") + + # Webhook sends image as multipart form file with key "file" + if 'file' in request.files: + img_file = request.files['file'] + timestamp = str(int(time.time())) + final_mask_path = TEMP_DIR / f"mask_{timestamp}.png" + + # Re-save as PNG (webhook sends JPEG) + from PIL import Image as PILImage + img = PILImage.open(img_file).convert('RGBA') + img.save(str(final_mask_path), format='PNG') + + logger.info(f"[WEBHOOK] Image file saved to {final_mask_path}") + _webhook_response = {'success': True, 'path': final_mask_path} + elif 'image' in request.files: + # Fallback for "image" key + img_file = request.files['image'] + timestamp = str(int(time.time())) + final_mask_path = TEMP_DIR / f"mask_{timestamp}.png" + + with open(final_mask_path, 'wb') as f: + img_file.save(str(final_mask_path)) + + logger.info(f"[WEBHOOK] Image file (fallback) saved to {final_mask_path}") + _webhook_response = {'success': True, 'path': final_mask_path} + elif request.files: + # Try first available file if 'image' key not present + img_file = list(request.files.values())[0] + timestamp = str(int(time.time())) + final_mask_path = TEMP_DIR / f"mask_{timestamp}.png" + + img_file.save(str(final_mask_path)) + + logger.info(f"[WEBHOOK] First file saved to {final_mask_path}") + _webhook_response = {'success': True, 'path': final_mask_path} + elif request.data: + # Fallback to raw data + timestamp = str(int(time.time())) + final_mask_path = TEMP_DIR / f"mask_{timestamp}.png" + + with open(final_mask_path, 'wb') as f: + f.write(request.data) + + logger.info(f"[WEBHOOK] Raw image data saved to {final_mask_path}") + _webhook_response = {'success': True, 'path': final_mask_path} + else: + logger.error("[WEBHOOK] No image file or data in request") + logger.debug(f"[WEBHOOK] Request form keys: {request.form.keys() if request.form else 'None'}") + logger.debug(f"[WEBHOOK] Request files keys: {request.files.keys() if request.files else 'None'}") + _webhook_response = {'success': False, 'error': 'No image data received'} + + # Signal that webhook is ready + _webhook_ready.set() + + # Return success to ComfyUI + response = make_response(jsonify({'status': 'ok', 'message': 'Image received'})) + response.status_code = 200 + return response + + except Exception as e: + logger.error(f"[WEBHOOK] Error processing webhook: {e}") + import traceback + traceback.print_exc() + _webhook_response = {'success': False, 'error': str(e)} + _webhook_ready.set() + + response = make_response(jsonify({'status': 'error', 'message': str(e)})) + response.status_code = 500 + return response + + +@app.route('/api/mask/extract', methods=['POST']) def api_mask_extract(): - """Extract mask using ComfyUI.""" + """Extract a mask from the current base image using ComfyUI.""" data = request.get_json() required = ['subject', 'ora_path'] @@ -356,24 +460,13 @@ def api_mask_extract(): subject = data['subject'] use_polygon = data.get('use_polygon', False) ora_path = data['ora_path'] - comfy_url = data.get('comfy_url', '127.0.0.1:8188') + comfy_url = data.get('comfy_url', COMFYUI_BASE_URL) logger.info(f"[MASK EXTRACT] Subject: {subject}") logger.info(f"[MASK EXTRACT] Use polygon: {use_polygon}") logger.info(f"[MASK EXTRACT] ORA path: {ora_path}") logger.info(f"[MASK EXTRACT] ComfyUI URL: {comfy_url}") - # Check if we need to apply polygon to image - if use_polygon: - polygon_data = _polygon_storage.get(ora_path, {}) - polygon_points = polygon_data.get('points', []) - - if not polygon_points or len(polygon_points) < 3: - logger.warning(f"[MASK EXTRACT] Use polygon requested but no valid polygon points stored") - return jsonify({'success': False, 'error': 'No valid polygon points. Please draw polygon first.'}), 400 - - logger.info(f"[MASK EXTRACT] Applying polygon overlay: {len(polygon_points)} points") - # Load workflow template workflow_path = APP_DIR.parent / "image_mask_extraction.json" if not workflow_path.exists(): @@ -384,9 +477,7 @@ def api_mask_extract(): workflow = json.load(f) logger.info(f"[MASK EXTRACT] Loaded workflow") - - # Update prompt text in workflow - + # Load base image from ORA base_img = None try: @@ -394,181 +485,264 @@ def api_mask_extract(): img_data = zf.read('mergedimage.png') base_img = Image.open(io.BytesIO(img_data)).convert('RGBA') logger.info(f"[MASK EXTRACT] Loaded base image: {base_img.size}") - + except Exception as e: logger.error(f"[MASK EXTRACT] Error loading base image: {e}") return jsonify({'success': False, 'error': f'Error loading image: {str(e)}'}), 500 - + # Apply polygon overlay if requested if use_polygon: polygon_data = _polygon_storage.get(ora_path, {}) polygon_points = polygon_data.get('points', []) - + if not polygon_points or len(polygon_points) < 3: logger.warning(f"[MASK EXTRACT] Use polygon requested but no valid polygon points stored") return jsonify({'success': False, 'error': 'No valid polygon points. Please draw polygon first.'}), 400 - + logger.info(f"[MASK EXTRACT] Applying polygon overlay: {len(polygon_points)} points") - + # Convert percentage points to pixel coordinates w, h = base_img.size pixel_points = [(int(p['x'] * w), int(p['y'] * h)) for p in polygon_points] - + # Draw polygon on image polygon_color = polygon_data.get('color', '#FF0000') polygon_width = polygon_data.get('width', 2) - + draw = ImageDraw.Draw(base_img) - + # Draw the polygon hex_color = polygon_color if len(polygon_color) == 7 else polygon_color + 'FF' draw.polygon(pixel_points, outline=hex_color, width=polygon_width) - + logger.info(f"[MASK EXTRACT] Drew polygon on image") - + # Encode (possibly modified) image img_io = io.BytesIO() base_img.save(img_io, format='PNG') img_io.seek(0) - + base64_image = base64.b64encode(img_io.read()).decode('utf-8') logger.info(f"[MASK EXTRACT] Encoded image as base64: {len(base64_image)} chars") - - # Set image input in workflow - workflow["168"]["inputs"]["image"] = base64_image - - # Update prompt text - for node_id, node in workflow.items(): - if 'inputs' in node and 'text' in node['inputs']: - node['inputs']['text'] = f"Create a black and white alpha mask of {subject}, leaving everything else black" - break - - # Poll for completion (up to 4 minutes) - start_time = time.time() - timeout = 240 - - while time.time() - start_time < timeout: - try: - req = urllib.request.Request( - f'http://{comfy_url}/history/{prompt_id}', - headers=headers, - method='GET' - ) - - with urllib.request.urlopen(req, timeout=30) as response: - history = json.loads(response.read().decode()) - - if prompt_id in history: - outputs = history[prompt_id].get('outputs', {}) - - logger.info(f"[MASK EXTRACT] Status check - prompt_id found in history") - - for node_id, node_output in outputs.items(): - if 'images' in node_output: - images = node_output['images'] - - logger.info(f"[MASK EXTRACT] Found {len(images)} images in node {node_id}") - - for img_info in images: - filename = img_info.get('filename', '') - subfolder = img_info.get('subfolder', '') - - logger.info(f"[MASK EXTRACT] Downloading image: {filename} from subfolder: {subfolder}") - - # Download the image - params = { - 'filename': filename, - 'subfolder': subfolder, - 'type': 'output' - } - - download_req = urllib.request.Request( - f'http://{comfy_url}/view?{urllib.parse.urlencode(params)}', - headers=headers, - method='GET' - ) - - with urllib.request.urlopen(download_req, timeout=30) as img_response: - img_data = img_response.read() - - # Save to temp directory - timestamp = str(int(time.time())) - mask_path = TEMP_DIR / f"mask_{timestamp}.png" - - with open(mask_path, 'wb') as f: - f.write(img_data) - - logger.info(f"[MASK EXTRACT] Saved mask to: {mask_path} ({len(img_data)} bytes)") - - return jsonify({ - 'success': True, - 'mask_path': str(mask_path), - 'mask_url': f'/api/file/mask?path={mask_path}' - }) - else: - logger.info(f"[MASK EXTRACT] Prompt_id not in history yet, waiting...") - except urllib.error.HTTPError as e: - if e.code != 404: - logger.error(f"[MASK EXTRACT] HTTP Error polling: {e.code}") - break - time.sleep(2) - - # Timeout - logger.error(f"[MASK EXTRACT] Timeout waiting for ComfyUI to complete") - return jsonify({'success': False, 'error': 'Mask extraction timed out'}), 500 - - except Exception as e: - logger.exception(f"[MASK EXTRACT] Error: {e}") - return jsonify({'success': False, 'error': str(e)}), 500 + # Set image input in workflow (node 87 is ETN_LoadImageBase64) + workflow["87"]["inputs"]["image"] = base64_image + + # Update prompt text in workflow (node 1:68 is TextEncodeQwenImageEditPlus) + if "1:68" in workflow and 'inputs' in workflow["1:68"]: + workflow["1:68"]["inputs"]["prompt"] = f"Create a black and white alpha mask of {subject}, leaving everything else black" + + # Set webhook URL for node 96 (Webhook Image Saver) + webhook_url = f"http://localhost:5001/api/webhook/comfyui" + if "96" in workflow and 'inputs' in workflow["96"]: + workflow["96"]["inputs"]["webhook_url"] = webhook_url + logger.info(f"[MASK EXTRACT] Webhook URL set to: {webhook_url}") + + # Set random seed to prevent cache hits (node 50 is Seed (rgthree)) + import random + random_seed = random.randint(0, 2**31-1) + if "50" in workflow and 'inputs' in workflow["50"]: + workflow["50"]["inputs"]["seed"] = random_seed + logger.info(f"[MASK EXTRACT] Random seed set to: {random_seed}") + + logger.info(f"[MASK EXTRACT] Workflow prepared, sending to ComfyUI at http://{comfy_url}") + + # Prepare headers for ComfyUI requests + headers = {'Content-Type': 'application/json'} + + # Submit workflow to ComfyUI + try: + payload = json.dumps({"prompt": workflow}).encode('utf-8') + submit_req = urllib.request.Request( + f'http://{comfy_url}/prompt', + data=payload, + headers=headers, + method='POST' + ) + + with urllib.request.urlopen(submit_req, timeout=30) as response: + result = json.loads(response.read().decode()) + prompt_id = result.get('prompt_id') + + if not prompt_id: + logger.error("[MASK EXTRACT] No prompt_id returned from ComfyUI") + return jsonify({'success': False, 'error': 'Failed to submit workflow to ComfyUI'}), 500 + + logger.info(f"[MASK EXTRACT] Prompt submitted with ID: {prompt_id}") + + except urllib.error.HTTPError as e: + logger.error(f"[MASK EXTRACT] HTTP Error submitting prompt: {e.code} - {e.read().decode()}") + return jsonify({'success': False, 'error': f'ComfyUI error: {str(e)}'}), 500 + except Exception as e: + logger.error(f"[MASK EXTRACT] Error submitting to ComfyUI: {e}") + return jsonify({'success': False, 'error': f'Failed to connect to ComfyUI: {str(e)}'}), 500 + + # Poll for completion (up to 4 minutes) + start_time = time.time() + timeout = 240 + + + while time.time() - start_time < timeout: + try: + req = urllib.request.Request( + f'http://{comfy_url}/history/{prompt_id}', + headers=headers, + method='GET' + ) + + with urllib.request.urlopen(req, timeout=30) as response: + history = json.loads(response.read().decode()) + + if prompt_id in history: + outputs = history[prompt_id].get('outputs', {}) + status = history[prompt_id].get('status', {}) + + if status.get('status_str') == 'success': + logger.info("[MASK EXTRACT] Workflow completed successfully!") + break # Exit polling loop, workflow is done + + logger.info(f"[MASK EXTRACT] Prompt_id in history but no outputs yet, waiting...") + else: + logger.info("[MASK EXTRACT] Prompt_id not in history yet, waiting...") + + time.sleep(2) + + except urllib.error.HTTPError as e: + if e.code == 404: + logger.debug("[MASK EXTRACT] Prompt not ready yet (404)") + time.sleep(2) + else: + logger.error(f"[MASK EXTRACT] HTTP Error polling: {e.code}") + break + except Exception as e: + logger.error(f"[MASK EXTRACT] Error polling history: {e}") + time.sleep(2) + + # Check if we timed out + if time.time() - start_time >= timeout: + logger.error("[MASK EXTRACT] Timeout waiting for ComfyUI to complete") + return jsonify({'success': False, 'error': 'Mask extraction timed out'}), 500 + + # Workflow completed - wait for webhook callback + logger.info(f"[MASK EXTRACT] Checking/waiting for webhook callback from ComfyUI...") + + global _webhook_response, _webhook_ready + + # Check if webhook already arrived (can happen before history check completes) + if _webhook_ready.is_set() and _webhook_response is not None: + logger.info("[MASK EXTRACT] Webhook already received!") + webhook_received = True + else: + # Clear event for this wait, but preserve response if it arrives now + _webhook_ready.clear() + _webhook_response = None + + # Wait for webhook with timeout (up to 60 seconds) + webhook_timeout = 60.0 + webhook_received = _webhook_ready.wait(timeout=webhook_timeout) + + if not webhook_received: + logger.error(f"[MASK EXTRACT] Timeout waiting for webhook callback") + return jsonify({'success': False, 'error': 'Webhook timeout - mask extraction may have failed'}), 500 + + # Check if webhook was successful + logger.info(f"[MASK EXTRACT] Webhook received: {_webhook_response}") + + if not _webhook_response or not _webhook_response.get('success'): + error_msg = _webhook_response.get('error', 'Unknown error') if _webhook_response else 'Empty response' + logger.error(f"[MASK EXTRACT] Webhook failed: {error_msg}") + return jsonify({'success': False, 'error': f'Webhook error: {error_msg}'}), 500 + + # Get the mask path from webhook response + final_mask_path = _webhook_response.get('path') + + if not final_mask_path: + logger.error("[MASK EXTRACT] No mask path in webhook response") + return jsonify({'success': False, 'error': 'No mask path returned from webhook'}), 500 + + try: + # Verify file exists + if not Path(final_mask_path).exists(): + logger.error(f"[MASK EXTRACT] Mask file not found at {final_mask_path}") + return jsonify({'success': False, 'error': f'Mask file not found: {final_mask_path}'}), 500 + + logger.info(f"[MASK EXTRACT] Mask received via webhook at {final_mask_path}") + + return jsonify({ + 'success': True, + 'mask_path': str(final_mask_path), + 'mask_url': f'/api/file/mask?path={final_mask_path}' + }) + except Exception as e: + logger.error(f"[MASK EXTRACT] Error processing mask file: {e}") + return jsonify({'success': False, 'error': f'Error accessing mask: {str(e)}'}), 500 @app.route('/api/krita/open', methods=['POST']) def api_krita_open(): - """Copy a layer to temp and return file URL for Krita.""" + """Open in Krita - behavior depends on mode. + + Review mode: Open the full ORA file for layer editing. + Add mode: Export base image with polygon overlay for manual annotation. + """ data = request.get_json() - required = ['ora_path', 'layer_name'] - for field in required: - if field not in data: - return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400 + if not data or 'ora_path' not in data: + return jsonify({'success': False, 'error': 'Missing ora_path parameter'}), 400 ora_path = data['ora_path'] - layer_name = data['layer_name'] try: - # Extract the layer image from ORA - root = parse_stack_xml(ora_path) - groups = load_ora(ora_path)['groups'] + # Review mode: Open full ORA file + if data.get('open_full_ora'): + timestamp = str(int(time.time())) + temp_file = TEMP_DIR / f"edit_{timestamp}.ora" + shutil.copy2(ora_path, temp_file) + + return jsonify({ + 'success': True, + 'file_url': f'file://{temp_file}', + 'temp_path': str(temp_file), + 'base_mtime': os.path.getmtime(temp_file) + }) - found = None - for group in groups: - for layer in group['layers']: - if layer['name'] == layer_name: - found = (group, layer) - break - if found: - break - - if not found: - return jsonify({'success': False, 'error': f'Layer {layer_name} not found'}), 404 - - _, layer_info = found + # Add mode: Open base with polygon overlay for annotation + elif data.get('open_base_with_polygon'): + points = data.get('points', []) + color = data.get('color', '#FF0000') + width = data.get('width', 2) + + # Load base image from ORA + with zipfile.ZipFile(ora_path, 'r') as zf: + img = Image.open(zf.open('mergedimage.png')).convert('RGBA') + + # Draw polygon overlay if points exist and we have 3+ points + if points and len(points) >= 3: + w, h = img.size + pixel_points = [(int(p['x'] * w), int(p['y'] * h)) for p in points] + + draw = ImageDraw.Draw(img) + hex_color = color if len(color) == 7 else color + 'FF' + draw.polygon(pixel_points, outline=hex_color, width=width) + + # Save to temp directory + timestamp = str(int(time.time())) + temp_file = TEMP_DIR / f"annotation_{timestamp}.png" + img.save(str(temp_file), format='PNG') + + return jsonify({ + 'success': True, + 'file_url': f'file://{temp_file}', + 'temp_path': str(temp_file), + 'base_mtime': os.path.getmtime(temp_file) + }) - with zipfile.ZipFile(ora_path, 'r') as zf: - img = Image.open(zf.open(layer_info['src'])).convert('RGBA') + else: + return jsonify({'success': False, 'error': 'Invalid request - specify open_full_ora or open_base_with_polygon'}), 400 - # Save to temp directory - timestamp = str(int(time.time())) - temp_file = TEMP_DIR / f"{layer_name}_{timestamp}.png" - img.save(temp_file) - - return jsonify({ - 'success': True, - 'file_url': f'file://{temp_file}', - 'temp_path': str(temp_file), - 'base_mtime': os.path.getmtime(temp_file) - }) except Exception as e: + import traceback + logger.error(f"[KRITA OPEN] Error: {e}") + traceback.print_exc() return jsonify({'success': False, 'error': str(e)}), 500 @@ -600,6 +774,36 @@ def api_krita_status(layer_name): }) +@app.route('/api/image/masked') +def api_image_masked(): + """Serve the base image with mask applied as alpha channel.""" + ora_path = request.args.get('ora_path') + mask_path = request.args.get('mask_path') + + if not ora_path or not mask_path: + return jsonify({'error': 'Missing ora_path or mask_path'}), 400 + + try: + with zipfile.ZipFile(ora_path, 'r') as zf: + base_img = Image.open(zf.open('mergedimage.png')).convert('RGBA') + + mask_img = Image.open(mask_path).convert('L') + + if mask_img.size != base_img.size: + mask_img = mask_img.resize(base_img.size, Image.Resampling.BILINEAR) + + r, g, b, _ = base_img.split() + result = Image.merge('RGBA', (r, g, b, mask_img)) + + img_io = io.BytesIO() + result.save(img_io, format='PNG') + img_io.seek(0) + + return send_file(img_io, mimetype='image/png') + except Exception as e: + return Response(f"Error creating masked image: {e}", status=500) + + @app.route('/api/file/mask') def api_file_mask(): """Serve a mask file from temp directory.""" @@ -652,4 +856,4 @@ def api_image_layer(layer_name): if __name__ == '__main__': - app.run(debug=True, port=5000, host='127.0.0.1') + app.run(debug=False, port=5001, host='127.0.0.1') diff --git a/tools/ora_editor/templates/editor.html b/tools/ora_editor/templates/editor.html index c078e9f..7b5f91a 100644 --- a/tools/ora_editor/templates/editor.html +++ b/tools/ora_editor/templates/editor.html @@ -10,10 +10,354 @@ ::-webkit-scrollbar { width: 8px; height: 8px; } ::-webkit-scrollbar-track { background: #374151; } ::-webkit-scrollbar-thumb { background: #6B7280; border-radius: 4px; } - .polygon-point { position: absolute; width: 20px; height: 20px; background: #FFFFFF; border: 3px solid #FF0000; border-radius: 50%; transform: translate(-50%, -50%); pointer-events: none; z-index: 100; } - - + + +
+ + +
+
+ + +
+ + + + + + + +
+ Scale: + + +
+ + +
+ + + + + +
+

Open a PNG or ORA file to start editing

+
+ + +
+
+
+

Extracted Mask

+ +
+ + +
+
+ +
+
+ + + + +
+
+ +
+ + + +
+
+
+ + +
+
+

Settings

+ +
+ + +
+ +
+ + +
+
+
+ +
+ + - - -
- -
-
- - -
- - -
- Scale: - - -
- -
- -