Restructure ORA editor into modular blueprints with browse dialog
- Split app.py into route blueprints (files, layers, images, polygon, mask, krita) - Create services layer (polygon_storage, comfyui, file_browser) - Extract config constants to config.py - Split templates into Jinja partials (base, components, modals) - Add browse dialog for visual file navigation - Add /api/browse endpoint for directory listing
This commit is contained in:
@@ -1,858 +1,39 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Flask web application for ORA editing."""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import logging
|
||||
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 flask import Flask, render_template
|
||||
|
||||
# Ensure the package can be imported
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from ora_editor.config import TEMP_DIR
|
||||
from ora_editor.routes import (
|
||||
files_bp, layers_bp, images_bp, polygon_bp, mask_bp, krita_bp
|
||||
)
|
||||
|
||||
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)
|
||||
# Register blueprints
|
||||
app.register_blueprint(files_bp)
|
||||
app.register_blueprint(layers_bp)
|
||||
app.register_blueprint(images_bp)
|
||||
app.register_blueprint(polygon_bp)
|
||||
app.register_blueprint(mask_bp)
|
||||
app.register_blueprint(krita_bp)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
Reference in New Issue
Block a user