|
import os
|
|
import json
|
|
from aiohttp import web
|
|
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
from ..utils import path_exists
|
|
from .utils_server import get_param, is_param_falsy
|
|
from .utils_info import delete_model_info, get_model_info, set_model_info_partial
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
|
|
def _check_valid_model_type(request):
|
|
model_type = request.match_info['type']
|
|
if model_type not in ['loras']:
|
|
return web.json_response({'status': 404, 'error': f'Invalid model type: {model_type}'})
|
|
return None
|
|
|
|
|
|
@routes.get('/rgthree/api/{type}')
|
|
async def api_get_models_list(request):
|
|
""" Returns a list of model types from user configuration. """
|
|
if _check_valid_model_type(request):
|
|
return _check_valid_model_type(request)
|
|
|
|
model_type = request.match_info['type']
|
|
data = folder_paths.get_filename_list(model_type)
|
|
return web.json_response(list(data))
|
|
|
|
|
|
@routes.get('/rgthree/api/{type}/info')
|
|
async def api_get_models_info(request):
|
|
""" Returns a list model info; either all or a single if provided a 'file' param. """
|
|
if _check_valid_model_type(request):
|
|
return _check_valid_model_type(request)
|
|
|
|
model_type = request.match_info['type']
|
|
file_param = get_param(request, 'file')
|
|
maybe_fetch_metadata = file_param is not None
|
|
if not is_param_falsy(request, 'light'):
|
|
maybe_fetch_metadata = False
|
|
api_response = await models_info_response(request,
|
|
model_type,
|
|
maybe_fetch_metadata=maybe_fetch_metadata)
|
|
return web.json_response(api_response)
|
|
|
|
|
|
@routes.get('/rgthree/api/{type}/info/refresh')
|
|
async def api_get_refresh_get_models_info(request):
|
|
""" Refreshes model info; either all or a single if provided a 'file' param. """
|
|
if _check_valid_model_type(request):
|
|
return _check_valid_model_type(request)
|
|
|
|
model_type = request.match_info['type']
|
|
api_response = await models_info_response(request,
|
|
model_type,
|
|
maybe_fetch_civitai=True,
|
|
maybe_fetch_metadata=True)
|
|
return web.json_response(api_response)
|
|
|
|
|
|
@routes.get('/rgthree/api/{type}/info/clear')
|
|
async def api_get_delete_model_info(request):
|
|
"""Clears model info from the filesystem for the provided file."""
|
|
if _check_valid_model_type(request):
|
|
return _check_valid_model_type(request)
|
|
|
|
model_type = request.match_info['type']
|
|
api_response = {'status': 200}
|
|
model_type = request.match_info['type']
|
|
file_param = get_param(request, 'file')
|
|
del_info = not is_param_falsy(request, 'del_info')
|
|
del_metadata = not is_param_falsy(request, 'del_metadata')
|
|
del_civitai = not is_param_falsy(request, 'del_civitai')
|
|
if file_param is None:
|
|
api_response['status'] = '404'
|
|
api_response['error'] = 'No file provided'
|
|
else:
|
|
file_params = [file_param]
|
|
if file_param == "ALL":
|
|
file_params = folder_paths.get_filename_list(model_type)
|
|
for file_param in file_params:
|
|
await delete_model_info(file_param,
|
|
model_type,
|
|
del_info=del_info,
|
|
del_metadata=del_metadata,
|
|
del_civitai=del_civitai)
|
|
return web.json_response(api_response)
|
|
|
|
|
|
@routes.post('/rgthree/api/{type}/info')
|
|
async def api_post_save_model_data(request):
|
|
"""Saves data to a model by name. """
|
|
if _check_valid_model_type(request):
|
|
return _check_valid_model_type(request)
|
|
|
|
model_type = request.match_info['type']
|
|
api_response = {'status': 200}
|
|
file_param = get_param(request, 'file')
|
|
if file_param is None:
|
|
api_response['status'] = '404'
|
|
api_response['error'] = 'No model found at path'
|
|
else:
|
|
post = await request.post()
|
|
await set_model_info_partial(file_param, model_type, json.loads(post.get("json")))
|
|
info_data = await get_model_info(file_param, model_type)
|
|
api_response['data'] = info_data
|
|
return web.json_response(api_response)
|
|
|
|
|
|
@routes.get('/rgthree/api/{type}/img')
|
|
async def api_get_models_info_img(request):
|
|
""" Returns an image response if one exists for the model. """
|
|
if _check_valid_model_type(request):
|
|
return _check_valid_model_type(request)
|
|
|
|
model_type = request.match_info['type']
|
|
file_param = get_param(request, 'file')
|
|
file_path = folder_paths.get_full_path(model_type, file_param)
|
|
if not path_exists(file_path):
|
|
file_path = os.path.abspath(file_path)
|
|
|
|
img_path = None
|
|
for ext in ['jpg', 'png', 'jpeg']:
|
|
try_path = f'{os.path.splitext(file_path)[0]}.{ext}'
|
|
if path_exists(try_path):
|
|
img_path = try_path
|
|
break
|
|
|
|
if not path_exists(img_path):
|
|
api_response = {}
|
|
api_response['status'] = '404'
|
|
api_response['error'] = 'No model found at path'
|
|
return web.json_response(api_response)
|
|
|
|
return web.FileResponse(img_path)
|
|
|
|
|
|
async def models_info_response(request,
|
|
model_type,
|
|
maybe_fetch_civitai=False,
|
|
maybe_fetch_metadata=False):
|
|
"""Gets model info for all or a single model type."""
|
|
api_response = {'status': 200}
|
|
file_param = get_param(request, 'file')
|
|
light = not is_param_falsy(request, 'light')
|
|
if file_param is not None:
|
|
info_data = await get_model_info(file_param,
|
|
model_type,
|
|
maybe_fetch_civitai=maybe_fetch_civitai,
|
|
maybe_fetch_metadata=maybe_fetch_metadata,
|
|
light=light)
|
|
if info_data is None:
|
|
api_response['status'] = '404'
|
|
api_response['error'] = f'No {model_type} found at path'
|
|
else:
|
|
api_response['data'] = info_data
|
|
else:
|
|
api_response['data'] = []
|
|
file_params = folder_paths.get_filename_list(model_type)
|
|
for file_param in file_params:
|
|
info_data = await get_model_info(file_param,
|
|
model_type,
|
|
maybe_fetch_civitai=maybe_fetch_civitai,
|
|
maybe_fetch_metadata=maybe_fetch_metadata,
|
|
light=light)
|
|
api_response['data'].append(info_data)
|
|
return api_response
|
|
|
