Create app.py

app.py

@@ -0,0 +1,127 @@
+import sys
+import asyncio
+from aiohttp import web, WSMsgType
+import json
+from json import JSONEncoder
+import numpy as np
+import uuid
+import logging
+import os
+import signal
+from typing import Dict, Any, List, Optional
+import base64
+import io
+
+from PIL import Image
+
+import pillow_avif
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+def SIGSEGV_signal_arises(signalNum, stack):
+    logger.critical(f"{signalNum} : SIGSEGV arises")
+    logger.critical(f"Stack trace: {stack}")
+
+signal.signal(signal.SIGSEGV, SIGSEGV_signal_arises)
+
+from loader import initialize_models
+from engine import Engine, base64_data_uri_to_PIL_Image
+
+# Global constants
+DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
+MODELS_DIR = os.path.join(DATA_ROOT, "models")
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.ndarray):
+            return obj.tolist()
+        else:
+            return super(NumpyEncoder, self).default(obj)
+
+async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
+    ws = web.WebSocketResponse()
+    await ws.prepare(request)
+    engine = request.app['engine']
+    try:
+        #logger.info("New WebSocket connection established")
+        while True:
+            msg = await ws.receive()
+
+            if msg.type in (WSMsgType.CLOSE, WSMsgType.ERROR):
+                #logger.warning(f"WebSocket connection closed: {msg.type}")
+                break
+
+            try:
+                if msg.type == WSMsgType.BINARY:
+                    res = await engine.load_image(msg.data)
+                    json_res = json.dumps(res, cls=NumpyEncoder)
+                    await ws.send_str(json_res)
+
+                elif msg.type == WSMsgType.TEXT:
+                    data = json.loads(msg.data)
+                    webp_bytes = await engine.transform_image(data.get('uuid'), data.get('params'))
+                    await ws.send_bytes(webp_bytes)
+
+            except Exception as e:
+                logger.error(f"Error in engine: {str(e)}")
+                logger.exception("Full traceback:")
+                await ws.send_json({"error": str(e)})
+
+    except Exception as e:
+        logger.error(f"Error in websocket_handler: {str(e)}")
+        logger.exception("Full traceback:")
+    return ws
+
+async def handle_upload(request: web.Request) -> web.Response:
+    """Recebe uma imagem e retorna informações sobre ela."""
+    engine = request.app['engine']
+    data = await request.content.read()
+    res = await engine.load_image(data)
+    return web.json_response(res)
+
+async def handle_modify(request: web.Request) -> web.Response:
+    """Recebe uma imagem e retorna informações sobre ela."""
+    engine = request.app['engine']
+    data = await request.json()
+    webp_bytes = await engine.transform_image(data.get('uuid'), data.get('params'))
+    return web.Response(body=webp_bytes, content_type="image/webp")
+    
+async def initialize_app() -> web.Application:
+    """Initialize and configure the web application."""
+    try:
+        logger.info("Initializing application...")
+        live_portrait = await initialize_models()
+
+        logger.info("🚀 Creating Engine instance...")
+        engine = Engine(live_portrait=live_portrait)
+        logger.info("✅ Engine instance created.")
+
+        app = web.Application()
+        app['engine'] = engine
+
+        # Configure routes
+        app.router.add_post("/upload", handle_upload)
+        app.router.add_post("/modify", handle_modify)
+
+        logger.info("Application routes configured")
+
+        return app
+    except Exception as e:
+        logger.error(f"🚨 Error during application initialization: {str(e)}")
+        logger.exception("Full traceback:")
+        raise
+
+if __name__ == "__main__":
+    try:
+        logger.info("Starting FacePoke application")
+        app = asyncio.run(initialize_app())
+        logger.info("Application initialized, starting web server")
+        web.run_app(app, host="0.0.0.0", port=8080)
+    except Exception as e:
+        logger.critical(f"🚨 FATAL: Failed to start the app: {str(e)}")
+        logger.exception("Full traceback:")