Bryce
61c8200443
- 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)
860 lines
30 KiB
Python
860 lines
30 KiB
Python
#!/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
|
|
|
|
from flask import (
|
|
Flask, request, jsonify, send_file, Response, render_template,
|
|
make_response
|
|
)
|
|
|
|
from PIL import Image, ImageDraw
|
|
import logging
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
# Configure logging first (before imports that might need logger)
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 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
|
|
sys.path.insert(0, str(APP_DIR.parent))
|
|
from ora_editor.ora_ops import (
|
|
load_ora, create_ora_from_png, add_masked_layer, rename_layer,
|
|
delete_layer, reorder_layer, get_layer_visibility, set_layer_visibility,
|
|
save_ora, parse_stack_xml
|
|
)
|
|
|
|
# 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')
|
|
|
|
|
|
def ensure_ora_exists(input_path: str) -> str | None:
|
|
"""Ensure an ORA file exists, creating from PNG if necessary."""
|
|
full_path = PROJECT_ROOT / input_path
|
|
|
|
if not full_path.exists():
|
|
return None
|
|
|
|
# If it's a PNG, create ORA
|
|
if full_path.suffix.lower() == '.png':
|
|
ora_path = str(full_path.with_suffix('.ora'))
|
|
result = create_ora_from_png(str(full_path), ora_path)
|
|
if not result.get('success'):
|
|
return None
|
|
return ora_path
|
|
|
|
# It's an ORA, verify it's valid
|
|
try:
|
|
parse_stack_xml(str(full_path))
|
|
return str(full_path)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
@app.route('/')
|
|
def index():
|
|
"""Serve the main editor UI."""
|
|
return render_template('editor.html')
|
|
|
|
|
|
@app.route('/api/open', methods=['POST'])
|
|
def api_open():
|
|
"""Open a file (PNG or ORA)."""
|
|
data = request.get_json()
|
|
|
|
if not data or 'path' not in data:
|
|
return jsonify({'success': False, 'error': 'Missing path parameter'}), 400
|
|
|
|
input_path = data['path']
|
|
ora_path = ensure_ora_exists(input_path)
|
|
|
|
if not ora_path:
|
|
return jsonify({'success': False, 'error': f'File not found or invalid: {input_path}'}), 404
|
|
|
|
try:
|
|
loaded = load_ora(ora_path)
|
|
|
|
# Build layers list with visibility
|
|
layers = []
|
|
for layer in loaded['layers']:
|
|
layers.append({
|
|
'name': layer['name'],
|
|
'src': layer['src'],
|
|
'group': layer.get('group'),
|
|
'visible': layer.get('visible', True)
|
|
})
|
|
|
|
return jsonify({
|
|
'success': True,
|
|
'ora_path': ora_path,
|
|
'width': loaded['width'],
|
|
'height': loaded['height'],
|
|
'layers': layers
|
|
})
|
|
except Exception as e:
|
|
return jsonify({'success': False, 'error': f'Error loading ORA: {str(e)}'}), 500
|
|
|
|
|
|
@app.route('/api/save', methods=['POST'])
|
|
def api_save():
|
|
"""Save the current ORA state."""
|
|
data = request.get_json()
|
|
|
|
if not data or 'ora_path' not in data:
|
|
return jsonify({'success': False, 'error': 'Missing ora_path parameter'}), 400
|
|
|
|
ora_path = data['ora_path']
|
|
|
|
# Save by re-reading and re-writing (preserves current state)
|
|
result = save_ora(ora_path)
|
|
|
|
if result:
|
|
return jsonify({'success': True, 'ora_path': ora_path})
|
|
else:
|
|
return jsonify({'success': False, 'error': 'Failed to save ORA'}), 500
|
|
|
|
|
|
@app.route('/api/layers', methods=['GET'])
|
|
def api_layers():
|
|
"""Get the current layer structure."""
|
|
ora_path = request.args.get('ora_path')
|
|
|
|
if not ora_path:
|
|
return jsonify({'success': False, 'error': 'Missing ora_path parameter'}), 400
|
|
|
|
try:
|
|
loaded = load_ora(ora_path)
|
|
|
|
layers = []
|
|
for layer in loaded['layers']:
|
|
layers.append({
|
|
'name': layer['name'],
|
|
'src': layer['src'],
|
|
'group': layer.get('group'),
|
|
'visible': layer.get('visible', True)
|
|
})
|
|
|
|
return jsonify({'success': True, 'layers': layers})
|
|
except Exception as e:
|
|
return jsonify({'success': False, 'error': str(e)}), 500
|
|
|
|
|
|
@app.route('/api/layer/add', methods=['POST'])
|
|
def api_layer_add():
|
|
"""Add a new masked layer."""
|
|
data = request.get_json()
|
|
|
|
required = ['ora_path', 'entity_name', 'mask_path']
|
|
for field in required:
|
|
if field not in data:
|
|
return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400
|
|
|
|
ora_path = data['ora_path']
|
|
entity_name = data['entity_name']
|
|
mask_path = data['mask_path']
|
|
source_layer = data.get('source_layer', 'base')
|
|
|
|
result = add_masked_layer(ora_path, entity_name, mask_path, source_layer)
|
|
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route('/api/layer/rename', methods=['POST'])
|
|
def api_layer_rename():
|
|
"""Rename a layer."""
|
|
data = request.get_json()
|
|
|
|
required = ['ora_path', 'old_name', 'new_name']
|
|
for field in required:
|
|
if field not in data:
|
|
return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400
|
|
|
|
result = rename_layer(data['ora_path'], data['old_name'], data['new_name'])
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route('/api/layer/delete', methods=['POST'])
|
|
def api_layer_delete():
|
|
"""Delete a layer."""
|
|
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
|
|
|
|
result = delete_layer(data['ora_path'], data['layer_name'])
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route('/api/layer/reorder', methods=['POST'])
|
|
def api_layer_reorder():
|
|
"""Reorder a layer."""
|
|
data = request.get_json()
|
|
|
|
required = ['ora_path', 'layer_name', 'direction']
|
|
for field in required:
|
|
if field not in data:
|
|
return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400
|
|
|
|
result = reorder_layer(data['ora_path'], data['layer_name'], data['direction'])
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route('/api/layer/visibility', methods=['POST'])
|
|
def api_layer_visibility():
|
|
"""Set layer visibility."""
|
|
data = request.get_json()
|
|
|
|
required = ['ora_path', 'layer_name', 'visible']
|
|
for field in required:
|
|
if field not in data:
|
|
return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400
|
|
|
|
result = set_layer_visibility(data['ora_path'], data['layer_name'], data['visible'])
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route('/api/image/<layer_name>')
|
|
def api_image(layer_name):
|
|
"""Serve a specific layer image."""
|
|
ora_path = request.args.get('ora_path')
|
|
|
|
if not ora_path:
|
|
return jsonify({'error': 'Missing ora_path'}), 400
|
|
|
|
try:
|
|
root = parse_stack_xml(ora_path)
|
|
|
|
with zipfile.ZipFile(ora_path, 'r') as zf:
|
|
# Get the image from mergedimage.png
|
|
img_data = zf.read('mergedimage.png')
|
|
return send_file(io.BytesIO(img_data), mimetype='image/png')
|
|
except Exception as e:
|
|
return Response(f"Error loading image: {e}", status=500)
|
|
|
|
|
|
@app.route('/api/image/base')
|
|
def api_image_base():
|
|
"""Serve the base/merged image."""
|
|
ora_path = request.args.get('ora_path')
|
|
|
|
if not ora_path:
|
|
return jsonify({'error': 'Missing ora_path'}), 400
|
|
|
|
try:
|
|
with zipfile.ZipFile(ora_path, 'r') as zf:
|
|
img_data = zf.read('mergedimage.png')
|
|
return send_file(io.BytesIO(img_data), mimetype='image/png')
|
|
except Exception as e:
|
|
return Response(f"Error loading image: {e}", status=500)
|
|
|
|
|
|
@app.route('/api/image/polygon')
|
|
def api_image_polygon():
|
|
"""Serve polygon overlay if stored."""
|
|
ora_path = request.args.get('ora_path')
|
|
|
|
if not ora_path or ora_path not in _polygon_storage:
|
|
return Response("No polygon data", status=404)
|
|
|
|
try:
|
|
with zipfile.ZipFile(ora_path, 'r') as zf:
|
|
base_img = Image.open(zf.open('mergedimage.png')).convert('RGBA')
|
|
|
|
points = _polygon_storage[ora_path].get('points', [])
|
|
color = _polygon_storage[ora_path].get('color', '#FF0000')
|
|
width = _polygon_storage[ora_path].get('width', 2)
|
|
|
|
if len(points) < 3:
|
|
return Response("Not enough points", status=404)
|
|
|
|
# Convert percentage points to pixel coordinates
|
|
w, h = base_img.size
|
|
pixel_points = [(int(p['x'] * w), int(p['y'] * h)) for p in points]
|
|
|
|
draw = ImageDraw.Draw(base_img)
|
|
# Draw as hex color with alpha
|
|
hex_color = color if len(color) == 7 else color + 'FF'
|
|
draw.polygon(pixel_points, outline=hex_color, width=width)
|
|
|
|
img_io = io.BytesIO()
|
|
base_img.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 drawing polygon: {e}", status=500)
|
|
|
|
|
|
@app.route('/api/polygon', methods=['POST'])
|
|
def api_polygon():
|
|
"""Store polygon points for overlay."""
|
|
data = request.get_json()
|
|
|
|
required = ['ora_path', 'points']
|
|
for field in required:
|
|
if field not in data:
|
|
return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400
|
|
|
|
ora_path = data['ora_path']
|
|
points = data['points']
|
|
color = data.get('color', '#FF0000')
|
|
width = data.get('width', 2)
|
|
|
|
logger.info(f"[POLYGON] Storing polygon: {len(points)} points, color: {color}, width: {width}")
|
|
|
|
_polygon_storage[ora_path] = {
|
|
'points': points,
|
|
'color': color,
|
|
'width': width
|
|
}
|
|
|
|
return jsonify({
|
|
'success': True,
|
|
'overlay_url': f'/api/image/polygon?ora_path={ora_path}'
|
|
})
|
|
|
|
|
|
@app.route('/api/polygon/clear', methods=['POST'])
|
|
def api_polygon_clear():
|
|
"""Clear stored polygon points."""
|
|
ora_path = request.json.get('ora_path') if request.data else None
|
|
|
|
if ora_path and ora_path in _polygon_storage:
|
|
del _polygon_storage[ora_path]
|
|
|
|
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 a mask from the current base image using ComfyUI."""
|
|
data = request.get_json()
|
|
|
|
required = ['subject', 'ora_path']
|
|
for field in required:
|
|
if field not in data:
|
|
return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400
|
|
|
|
subject = data['subject']
|
|
use_polygon = data.get('use_polygon', False)
|
|
ora_path = data['ora_path']
|
|
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}")
|
|
|
|
# Load workflow template
|
|
workflow_path = APP_DIR.parent / "image_mask_extraction.json"
|
|
if not workflow_path.exists():
|
|
logger.error(f"[MASK EXTRACT] Workflow file not found: {workflow_path}")
|
|
return jsonify({'success': False, 'error': f'Workflow file not found: {workflow_path}'}), 500
|
|
|
|
with open(workflow_path) as f:
|
|
workflow = json.load(f)
|
|
|
|
logger.info(f"[MASK EXTRACT] Loaded workflow")
|
|
|
|
# Load base image from ORA
|
|
base_img = None
|
|
try:
|
|
with zipfile.ZipFile(ora_path, 'r') as zf:
|
|
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 (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():
|
|
"""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()
|
|
|
|
if not data or 'ora_path' not in data:
|
|
return jsonify({'success': False, 'error': 'Missing ora_path parameter'}), 400
|
|
|
|
ora_path = data['ora_path']
|
|
|
|
try:
|
|
# 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)
|
|
})
|
|
|
|
# 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
|
|
|
|
|
|
@app.route('/api/krita/status/<layer_name>')
|
|
def api_krita_status(layer_name):
|
|
"""Check if a Krita temp file has been modified."""
|
|
# Find the temp file for this layer
|
|
temp_files = list(TEMP_DIR.glob(f"{layer_name}_*.png"))
|
|
|
|
if not temp_files:
|
|
return jsonify({'success': False, 'error': 'No temp file found'}), 404
|
|
|
|
# Get most recent file
|
|
temp_file = max(temp_files, key=lambda f: f.stat().st_mtime)
|
|
|
|
mtime = temp_file.stat().st_mtime
|
|
|
|
stored_mtime = request.args.get('stored_mtime')
|
|
if stored_mtime:
|
|
modified = float(stored_mtime) != mtime
|
|
else:
|
|
modified = True
|
|
|
|
return jsonify({
|
|
'success': True,
|
|
'modified': modified,
|
|
'mtime': mtime,
|
|
'file_url': f'file://{temp_file}'
|
|
})
|
|
|
|
|
|
@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."""
|
|
path = request.args.get('path')
|
|
|
|
if not path:
|
|
return jsonify({'error': 'Missing path'}), 400
|
|
|
|
full_path = Path(path)
|
|
|
|
if not full_path.exists():
|
|
return jsonify({'error': 'File not found'}), 404
|
|
|
|
return send_file(full_path, mimetype='image/png')
|
|
|
|
|
|
@app.route('/api/image/layer/<layer_name>')
|
|
def api_image_layer(layer_name):
|
|
"""Serve a specific layer as image."""
|
|
ora_path = request.args.get('ora_path')
|
|
|
|
if not ora_path:
|
|
return jsonify({'error': 'Missing ora_path'}), 400
|
|
|
|
try:
|
|
# Find the layer source path
|
|
root = parse_stack_xml(ora_path)
|
|
|
|
with zipfile.ZipFile(ora_path, 'r') as zf:
|
|
stack = ET.fromstring(zf.read('stack.xml'))
|
|
|
|
for child in stack.find('stack'):
|
|
if child.tag == 'layer' and child.get('name') == layer_name:
|
|
src = child.get('src')
|
|
img_data = zf.read(src)
|
|
return send_file(io.BytesIO(img_data), mimetype='image/png')
|
|
|
|
elif child.tag == 'stack': # Group
|
|
for layer in child:
|
|
if layer.tag == 'layer' and layer.get('name') == layer_name:
|
|
src = layer.get('src')
|
|
img_data = zf.read(src)
|
|
return send_file(io.BytesIO(img_data), mimetype='image/png')
|
|
|
|
return jsonify({'error': 'Layer not found'}), 404
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
return Response(f"Error loading layer: {e}", status=500)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app.run(debug=False, port=5001, host='127.0.0.1')
|