ai-game-2/tools/ora_editor/routes/mask.py

"""Mask extraction routes for ORA Editor."""

import io
import json
import time
import zipfile
from pathlib import Path
from flask import Blueprint, request, jsonify, make_response

from ora_editor.config import APP_DIR, COMFYUI_BASE_URL, TEMP_DIR
from ora_editor.services import polygon_storage
from ora_editor.services.comfyui import ComfyUIService
from ora_editor.ora_ops import parse_stack_xml

import logging
logger = logging.getLogger(__name__)

mask_bp = Blueprint('mask', __name__)
comfy_service = ComfyUIService(COMFYUI_BASE_URL)


@mask_bp.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")
    
    result = comfy_service.handle_webhook(
        request.files,
        request.form if request.form else {},
        request.data,
        TEMP_DIR
    )
    
    if result.get('success'):
        response = make_response(jsonify({'status': 'ok', 'message': 'Image received'}))
        response.status_code = 200
    else:
        response = make_response(jsonify({'status': 'error', 'message': result.get('error', 'Unknown error')}))
        response.status_code = 500
    
    return response


@mask_bp.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}")
    
    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_template = json.load(f)
    
    logger.info(f"[MASK EXTRACT] Loaded workflow")
    
    base_img = None
    try:
        with zipfile.ZipFile(ora_path, 'r') as zf:
            img_data = zf.read('mergedimage.png')
            base_img = __import__('PIL').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
    
    polygon_points = None
    polygon_color = '#FF0000'
    polygon_width = 2
    
    if use_polygon:
        poly_data = polygon_storage.get(ora_path)
        if poly_data:
            polygon_points = poly_data.get('points', [])
            polygon_color = poly_data.get('color', '#FF0000')
            polygon_width = poly_data.get('width', 2)
        
        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
    
    webhook_url = f"http://localhost:5001/api/webhook/comfyui"
    
    workflow = comfy_service.prepare_mask_workflow(
        base_image=base_img,
        subject=subject,
        webhook_url=webhook_url,
        polygon_points=polygon_points,
        polygon_color=polygon_color,
        polygon_width=polygon_width,
        workflow_template=workflow_template
    )
    
    logger.info(f"[MASK EXTRACT] Workflow prepared, sending to ComfyUI at http://{comfy_url}")
    
    try:
        prompt_id = comfy_service.submit_workflow(workflow, comfy_url)
        logger.info(f"[MASK EXTRACT] Prompt submitted with ID: {prompt_id}")
    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
    
    completed = comfy_service.poll_for_completion(prompt_id, comfy_url, timeout=240)
    
    if not completed:
        logger.error("[MASK EXTRACT] Timeout waiting for ComfyUI to complete")
        return jsonify({'success': False, 'error': 'Mask extraction timed out'}), 500
    
    logger.info(f"[MASK EXTRACT] Checking/waiting for webhook callback from ComfyUI...")
    
    webhook_result = comfy_service.wait_for_webhook(timeout=60.0)
    
    if not webhook_result:
        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] Webhook received: {webhook_result}")
    
    if not webhook_result.get('success'):
        error_msg = webhook_result.get('error', 'Unknown error')
        logger.error(f"[MASK EXTRACT] Webhook failed: {error_msg}")
        return jsonify({'success': False, 'error': f'Webhook error: {error_msg}'}), 500
    
    final_mask_path = webhook_result.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:
        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