ComfyUI Proxy

auth_db.py

#!/usr/bin/env python3
import sqlite3
import uuid
from datetime import datetime, timedelta

class AuthDB:
    def __init__(self, db_path='auth.db'):
        self.db_path = db_path
        self.init_db()

    def init_db(self):
        conn = sqlite3.connect(self.db_path)
        conn.execute('''
            CREATE TABLE IF NOT EXISTS image_ownership (
                image_id TEXT PRIMARY KEY,
                client_token TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        conn.commit()
        conn.close()

    def register_image(self, image_id, client_token):
        conn = sqlite3.connect(self.db_path)
        conn.execute('INSERT INTO image_ownership (image_id, client_token) VALUES (?, ?)',
                    (image_id, client_token))
        conn.commit()
        conn.close()

    def verify_access(self, image_id, client_token):
        conn = sqlite3.connect(self.db_path)
        result = conn.execute('SELECT client_token FROM image_ownership WHERE image_id = ?',
                             (image_id,)).fetchone()
        conn.close()
        return result and result[0] == client_token

    def cleanup_old_images(self, days=30):
        conn = sqlite3.connect(self.db_path)
        cutoff = datetime.now() - timedelta(days=days)
        conn.execute('DELETE FROM image_ownership WHERE created_at < ?', (cutoff,))
        conn.commit()
        conn.close()

comfyui_proxy.py

#!/usr/bin/env python3
from flask import Flask, request, Response, jsonify
import requests
import json
import datetime
import logging
import re
from jsonschema import validate, ValidationError
from auth_db import AuthDB

app = Flask(__name__)
COMFYUI_HOST = "127.0.0.1:8188"
COMFYUI_URL = f"http://{COMFYUI_HOST}"
auth_db = AuthDB()

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger(__name__)

with open('schemas/workflow.json', 'r') as f:
    WORKFLOW_SCHEMA = json.load(f)
with open('schemas/image_request.json', 'r') as f:
    IMAGE_REQUEST_SCHEMA = json.load(f)


def validate_request(instance, schema):
    try:
        validate(instance=instance, schema=schema)
        return None
    except ValidationError as e:
        return jsonify({"Validation error": str(e)}), 400
    except Exception as e:
        return jsonify({"Unexpected error": str(e)}), 400


@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE'])
def proxy(path):

    client_token = request.headers.get('X-Client-Token')
    if not client_token:
        return jsonify({"error": "Missing X-Client-Token header"}), 401

    if request.method == 'POST' and path == 'prompt':
        payload = request.get_json(force=True)
        validation_result = validate_request(instance=payload, schema=WORKFLOW_SCHEMA)
        if validation_result:
            return validation_result

        # Extract and register image ID for this client
        image_id = payload['prompt']['9']['inputs']['filename_prefix']
        auth_db.register_image(image_id, client_token)

    elif request.method == 'GET' and path == 'view':
        validation_result = validate_request(instance=dict(request.args), schema=IMAGE_REQUEST_SCHEMA)
        if validation_result:
            return validation_result

        # Extract image ID from filename and verify ownership
        filename = request.args.get('filename', '')
        image_id = re.match(r'^([0-9a-f-]+)_', filename)
        if not image_id or not auth_db.verify_access(image_id.group(1), client_token):
            return jsonify({"error": "Access denied"}), 403

    else:
        return jsonify({"error": "Not found"}), 404

    log_data = {
        'timestamp': datetime.datetime.now().isoformat(),
        'method': request.method,
        'path': path,
        'query_params': dict(request.args),
        'headers': dict(request.headers)
    }
    logger.info(f"Proxying request: {json.dumps(log_data)}")

    comfyui_response = requests.request(
        method=request.method,
        url=f"{COMFYUI_URL}/{path}",
        headers={},
        data=request.get_data(),
        params=request.args,
        allow_redirects=False
    )

    logger.info(f"Proxied response: {comfyui_response.status_code}")
    return Response(comfyui_response.content, comfyui_response.status_code, comfyui_response.headers.items())


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8189, debug=False)