File size: 6,535 Bytes
4450790
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
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