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)
This commit is contained in:
2026-03-27 17:17:19 -07:00
parent 087ac4ca0e
commit 61c8200443
2 changed files with 861 additions and 583 deletions

View File

@@ -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/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():
@@ -385,8 +478,6 @@ def api_mask_extract():
logger.info(f"[MASK EXTRACT] Loaded workflow")
# Update prompt text in workflow
# Load base image from ORA
base_img = None
try:
@@ -434,19 +525,63 @@ def api_mask_extract():
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
# Set image input in workflow (node 87 is ETN_LoadImageBase64)
workflow["87"]["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
# 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(
@@ -460,107 +595,108 @@ def api_mask_extract():
if prompt_id in history:
outputs = history[prompt_id].get('outputs', {})
status = history[prompt_id].get('status', {})
logger.info(f"[MASK EXTRACT] Status check - prompt_id found in history")
if status.get('status_str') == 'success':
logger.info("[MASK EXTRACT] Workflow completed successfully!")
break # Exit polling loop, workflow is done
for node_id, node_output in outputs.items():
if 'images' in node_output:
images = node_output['images']
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...")
logger.info(f"[MASK EXTRACT] Found {len(images)} images in node {node_id}")
time.sleep(2)
for img_info in images:
filename = img_info.get('filename', '')
subfolder = img_info.get('subfolder', '')
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)
logger.info(f"[MASK EXTRACT] Downloading image: {filename} from subfolder: {subfolder}")
# 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
# Download the image
params = {
'filename': filename,
'subfolder': subfolder,
'type': 'output'
}
# Workflow completed - wait for webhook callback
logger.info(f"[MASK EXTRACT] Checking/waiting for webhook callback from ComfyUI...")
download_req = urllib.request.Request(
f'http://{comfy_url}/view?{urllib.parse.urlencode(params)}',
headers=headers,
method='GET'
)
global _webhook_response, _webhook_ready
with urllib.request.urlopen(download_req, timeout=30) as img_response:
img_data = img_response.read()
# 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
# Save to temp directory
timestamp = str(int(time.time()))
mask_path = TEMP_DIR / f"mask_{timestamp}.png"
# Wait for webhook with timeout (up to 60 seconds)
webhook_timeout = 60.0
webhook_received = _webhook_ready.wait(timeout=webhook_timeout)
with open(mask_path, 'wb') as f:
f.write(img_data)
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
logger.info(f"[MASK EXTRACT] Saved mask to: {mask_path} ({len(img_data)} bytes)")
# 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(mask_path),
'mask_url': f'/api/file/mask?path={mask_path}'
'mask_path': str(final_mask_path),
'mask_url': f'/api/file/mask?path={final_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
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']
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
with zipfile.ZipFile(ora_path, 'r') as zf:
img = Image.open(zf.open(layer_info['src'])).convert('RGBA')
# Save to temp directory
# Review mode: Open full ORA file
if data.get('open_full_ora'):
timestamp = str(int(time.time()))
temp_file = TEMP_DIR / f"{layer_name}_{timestamp}.png"
img.save(temp_file)
temp_file = TEMP_DIR / f"edit_{timestamp}.ora"
shutil.copy2(ora_path, temp_file)
return jsonify({
'success': True,
@@ -568,7 +704,45 @@ def api_krita_open():
'temp_path': str(temp_file),
'base_mtime': os.path.getmtime(temp_file)
})
# 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)
})
else:
return jsonify({'success': False, 'error': 'Invalid request - specify open_full_ora or open_base_with_polygon'}), 400
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')

File diff suppressed because it is too large Load Diff