mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
import logging
|
|
from aiohttp import web
|
|
from ..utils.usage_stats import UsageStats
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class UsageStatsRoutes:
|
|
"""Routes for handling usage statistics updates"""
|
|
|
|
@staticmethod
|
|
def setup_routes(app):
|
|
"""Register usage stats routes"""
|
|
app.router.add_post('/loras/api/update-usage-stats', UsageStatsRoutes.update_usage_stats)
|
|
app.router.add_get('/loras/api/get-usage-stats', UsageStatsRoutes.get_usage_stats)
|
|
|
|
@staticmethod
|
|
async def update_usage_stats(request):
|
|
"""
|
|
Update usage statistics based on a prompt_id
|
|
|
|
Expects a JSON body with:
|
|
{
|
|
"prompt_id": "string"
|
|
}
|
|
"""
|
|
try:
|
|
# Parse the request body
|
|
data = await request.json()
|
|
prompt_id = data.get('prompt_id')
|
|
|
|
if not prompt_id:
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': 'Missing prompt_id'
|
|
}, status=400)
|
|
|
|
# Call the UsageStats to process this prompt_id synchronously
|
|
usage_stats = UsageStats()
|
|
await usage_stats.process_execution(prompt_id)
|
|
|
|
return web.json_response({
|
|
'success': True
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to update usage stats: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|
|
|
|
@staticmethod
|
|
async def get_usage_stats(request):
|
|
"""Get current usage statistics"""
|
|
try:
|
|
usage_stats = UsageStats()
|
|
stats = await usage_stats.get_stats()
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'data': stats
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to get usage stats: {e}", exc_info=True)
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|