Spaces:
Paused
Paused
File size: 4,514 Bytes
f13f9ba 65b8171 2d74807 65b8171 2d74807 65b8171 f13f9ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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_hello(request: web.Request) -> web.Response:
return web.json_response({
"message": "hello folks!"
})
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)
app.router.add_get("/", handle_hello)
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:") |