custom_nodes/rgthree-comfy/py/server/routes_model_info.py

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":  # Force the user to supply file=ALL to trigger all clearing.
      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