jbilcke-hf HF staff commited on
Commit
ad68cdf
β€’
1 Parent(s): d69879c

cleaning up

Browse files
Files changed (1) hide show
  1. app.py +1 -77
app.py CHANGED
@@ -7,19 +7,16 @@ Date: September 30, 2024
7
 
8
  import sys
9
  import asyncio
10
- import hashlib
11
  from aiohttp import web, WSMsgType
12
  import json
13
  import uuid
14
  import logging
15
  import os
16
- import zipfile
17
  import signal
18
  from typing import Dict, Any, List, Optional
19
  import base64
20
  import io
21
  from PIL import Image
22
- import numpy as np
23
 
24
  # Configure logging
25
  logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -58,8 +55,6 @@ async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
58
  """
59
  ws = web.WebSocketResponse()
60
  await ws.prepare(request)
61
-
62
- session: Optional[FacePokeSession] = None
63
  try:
64
  logger.info("New WebSocket connection established")
65
 
@@ -87,11 +82,6 @@ async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
87
  except Exception as e:
88
  logger.error(f"Error in websocket_handler: {str(e)}")
89
  logger.exception("Full traceback:")
90
- finally:
91
- if session:
92
- await session.stop()
93
- del active_sessions[session.session_id]
94
- logger.info("WebSocket connection closed")
95
  return ws
96
 
97
  async def handle_modify_image(request: web.Request, ws: web.WebSocketResponse, msg: Dict[str, Any], uuid: str):
@@ -152,10 +142,6 @@ async def hf_logo(request: web.Request) -> web.Response:
152
  async def on_shutdown(app: web.Application):
153
  """Cleanup function to be called on server shutdown."""
154
  logger.info("Server shutdown initiated, cleaning up resources...")
155
- for session in list(active_sessions.values()):
156
- await session.stop()
157
- active_sessions.clear()
158
- logger.info("All active sessions have been closed")
159
 
160
  if 'engine' in app:
161
  await app['engine'].cleanup()
@@ -191,74 +177,12 @@ async def initialize_app() -> web.Application:
191
  logger.exception("Full traceback:")
192
  raise
193
 
194
- async def start_background_tasks(app: web.Application):
195
- """
196
- Start background tasks for the application.
197
-
198
- Args:
199
- app (web.Application): The web application instance.
200
- """
201
- app['cleanup_task'] = asyncio.create_task(periodic_cleanup(app))
202
-
203
- async def cleanup_background_tasks(app: web.Application):
204
- """
205
- Clean up background tasks when the application is shutting down.
206
-
207
- Args:
208
- app (web.Application): The web application instance.
209
- """
210
- app['cleanup_task'].cancel()
211
- await app['cleanup_task']
212
-
213
- async def periodic_cleanup(app: web.Application):
214
- """
215
- Perform periodic cleanup tasks for the application.
216
-
217
- Args:
218
- app (web.Application): The web application instance.
219
- """
220
- while True:
221
- try:
222
- await asyncio.sleep(3600) # Run cleanup every hour
223
- await cleanup_inactive_sessions(app)
224
- except asyncio.CancelledError:
225
- break
226
- except Exception as e:
227
- logger.error(f"Error in periodic cleanup: {str(e)}")
228
- logger.exception("Full traceback:")
229
-
230
- async def cleanup_inactive_sessions(app: web.Application):
231
- """
232
- Clean up inactive sessions.
233
-
234
- Args:
235
- app (web.Application): The web application instance.
236
- """
237
- logger.info("Starting cleanup of inactive sessions")
238
- inactive_sessions = [
239
- session_id for session_id, session in active_sessions.items()
240
- if not session.is_running.is_set()
241
- ]
242
- for session_id in inactive_sessions:
243
- session = active_sessions.pop(session_id)
244
- await session.stop()
245
- logger.info(f"Cleaned up inactive session: {session_id}")
246
- logger.info(f"Cleaned up {len(inactive_sessions)} inactive sessions")
247
-
248
- def main():
249
- """
250
- Main function to start the FacePoke application.
251
- """
252
  try:
253
  logger.info("Starting FacePoke application")
254
  app = asyncio.run(initialize_app())
255
- app.on_startup.append(start_background_tasks)
256
- app.on_cleanup.append(cleanup_background_tasks)
257
  logger.info("Application initialized, starting web server")
258
  web.run_app(app, host="0.0.0.0", port=8080)
259
  except Exception as e:
260
  logger.critical(f"🚨 FATAL: Failed to start the app: {str(e)}")
261
  logger.exception("Full traceback:")
262
-
263
- if __name__ == "__main__":
264
- main()
 
7
 
8
  import sys
9
  import asyncio
 
10
  from aiohttp import web, WSMsgType
11
  import json
12
  import uuid
13
  import logging
14
  import os
 
15
  import signal
16
  from typing import Dict, Any, List, Optional
17
  import base64
18
  import io
19
  from PIL import Image
 
20
 
21
  # Configure logging
22
  logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
55
  """
56
  ws = web.WebSocketResponse()
57
  await ws.prepare(request)
 
 
58
  try:
59
  logger.info("New WebSocket connection established")
60
 
 
82
  except Exception as e:
83
  logger.error(f"Error in websocket_handler: {str(e)}")
84
  logger.exception("Full traceback:")
 
 
 
 
 
85
  return ws
86
 
87
  async def handle_modify_image(request: web.Request, ws: web.WebSocketResponse, msg: Dict[str, Any], uuid: str):
 
142
  async def on_shutdown(app: web.Application):
143
  """Cleanup function to be called on server shutdown."""
144
  logger.info("Server shutdown initiated, cleaning up resources...")
 
 
 
 
145
 
146
  if 'engine' in app:
147
  await app['engine'].cleanup()
 
177
  logger.exception("Full traceback:")
178
  raise
179
 
180
+ if __name__ == "__main__":
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
181
  try:
182
  logger.info("Starting FacePoke application")
183
  app = asyncio.run(initialize_app())
 
 
184
  logger.info("Application initialized, starting web server")
185
  web.run_app(app, host="0.0.0.0", port=8080)
186
  except Exception as e:
187
  logger.critical(f"🚨 FATAL: Failed to start the app: {str(e)}")
188
  logger.exception("Full traceback:")