brantje
00558f031d
Apparently, success_resp didn’t need a message—it was doing just fine without it. Cleaned it up so it stops oversharing like that one friend at brunch. You're welcome, future me.
209 lines
7.0 KiB
Python
209 lines
7.0 KiB
Python
import logging
|
|
import os
|
|
import traceback
|
|
from urllib.parse import urlparse
|
|
from aiohttp.web import Request, json_response
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
from .model_utils.refresh import refresh_folder
|
|
from .metrics.prometheus import get_metrics
|
|
from .model_utils.install import install_model_url
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
|
|
def success_resp(**kwargs):
|
|
return json_response({"code": 200, "message": "success", **kwargs})
|
|
|
|
|
|
def error_resp(code, message, **kwargs):
|
|
return json_response({"code": code, "message": message, **kwargs})
|
|
|
|
@routes.get("/api-tools/v1/models")
|
|
async def get_model_folders(request: Request):
|
|
""" List all model folders """
|
|
folders = list(folder_paths.folder_names_and_paths.keys())
|
|
|
|
return success_resp(result={"folders": folders })
|
|
|
|
@routes.post("/api-tools/v1/models/install")
|
|
async def install_model(request: Request):
|
|
"""Install a model from url, filtered with a allow list """
|
|
json_data = await request.json()
|
|
url = json_data.get('url')
|
|
allowed_domains = ["civitai.com", "github.com", "raw.githubusercontent.com", "huggingface.co"]
|
|
if url:
|
|
try:
|
|
parsed_url = urlparse(url)
|
|
domain = parsed_url.netloc
|
|
if domain.startswith('www.'):
|
|
domain = domain[4:]
|
|
|
|
# Check if the domain is in the allowed list
|
|
if domain not in allowed_domains:
|
|
return error_resp(403, f"Domain '{domain}' is not in the allowed domains list")
|
|
except Exception as e:
|
|
return error_resp(400, f"Invalid URL format: {str(e)}")
|
|
else:
|
|
return error_resp(400, "URL is required")
|
|
|
|
success, message = install_model_url(json_data)
|
|
|
|
if not success:
|
|
return error_resp(409, message)
|
|
|
|
return success_resp()
|
|
|
|
|
|
|
|
@routes.get("/api-tools/v1/models/{folder}")
|
|
async def get_model_folder(request: Request):
|
|
""" List all models in a folder """
|
|
folder = request.match_info['folder']
|
|
try:
|
|
checkpoints = folder_paths.get_filename_list(folder)
|
|
result = [
|
|
{
|
|
"name": os.path.basename(folder),
|
|
"path": ckpt,
|
|
"full_path": folder_paths.get_full_path(folder, ckpt) or ckpt,
|
|
}
|
|
for ckpt in checkpoints
|
|
]
|
|
return success_resp(result=result)
|
|
except Exception as e:
|
|
return error_resp(500, str('Path not found'))
|
|
|
|
|
|
@routes.post("/api-tools/v1/models/{folder}/refresh")
|
|
async def refresh_checkpoints(request: Request):
|
|
"""Refresh the checkpoints list and return the updated list(maybe useful for models that stored in a network storage)"""
|
|
folder = request.match_info['folder']
|
|
try:
|
|
data = refresh_folder(folder)
|
|
return success_resp(data=data)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.delete("/api-tools/v1/models/{folder}/{model:.+}")
|
|
async def remove_model(request: Request):
|
|
"""Removes a model from the file system """
|
|
input_folder = request.match_info['folder']
|
|
model_input = request.match_info['model']
|
|
try:
|
|
models = folder_paths.get_filename_list(input_folder)
|
|
full_path = None
|
|
for model in models:
|
|
if model == model_input:
|
|
model_full_path = folder_paths.get_full_path(input_folder, model)
|
|
full_path = model_full_path
|
|
break
|
|
|
|
if full_path:
|
|
os.remove(full_path)
|
|
return success_resp(result=full_path)
|
|
else:
|
|
return error_resp(404, 'Not Found')
|
|
|
|
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.get("/api-tools/v1/images/output")
|
|
async def get_output_images(request: Request):
|
|
try:
|
|
is_temp = request.rel_url.query.get("temp", "false") == "true"
|
|
folder = (
|
|
folder_paths.get_temp_directory()
|
|
if is_temp
|
|
else folder_paths.get_output_directory()
|
|
)
|
|
# iterate through the folder and get the list of images
|
|
images = []
|
|
for root, dirs, files in os.walk(folder):
|
|
for file in files:
|
|
if file.endswith(".png") or file.endswith(".jpg"):
|
|
image = {"name": file, "full_path": os.path.join(root, file)}
|
|
images.append(image)
|
|
return success_resp(images=images)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
|
|
@routes.delete("/api-tools/v1/images/output/{filename}")
|
|
async def delete_output_images(request: Request):
|
|
try:
|
|
filename = request.match_info.get("filename")
|
|
if filename is None:
|
|
return error_resp(400, "filename is required")
|
|
|
|
if filename[0] == "/" or ".." in filename:
|
|
return error_resp(400, "invalid filename")
|
|
|
|
is_temp = request.rel_url.query.get("temp", "false") == "true"
|
|
annotated_file = f"{filename} [{'temp' if is_temp else 'output'}]"
|
|
if not folder_paths.exists_annotated_filepath(annotated_file):
|
|
return error_resp(404, f"file {filename} not found")
|
|
|
|
filepath = folder_paths.get_annotated_filepath(annotated_file)
|
|
os.remove(filepath)
|
|
return success_resp()
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.get("/api-tools/v1/images/input")
|
|
async def get_input_images(request: Request):
|
|
try:
|
|
folder = folder_paths.get_input_directory()
|
|
|
|
# iterate through the folder and get the list of images
|
|
images = []
|
|
for root, dirs, files in os.walk(folder):
|
|
for file in files:
|
|
if file.endswith(".png") or file.endswith(".jpg"):
|
|
image = {"name": file, "full_path": os.path.join(root, file)}
|
|
images.append(image)
|
|
return success_resp(images=images)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.delete("/api-tools/v1/images/input/{filename}")
|
|
async def delete_input_images(request: Request):
|
|
try:
|
|
filename = request.match_info.get("filename")
|
|
if filename is None:
|
|
return error_resp(400, "filename is required")
|
|
|
|
if filename[0] == '/' or '..' in filename:
|
|
return error_resp(400, "invalid filename")
|
|
|
|
is_temp = request.rel_url.query.get("temp", "false") == "true"
|
|
annotated_file = f"{filename} [{'temp' if is_temp else 'input'}]"
|
|
if not folder_paths.exists_annotated_filepath(annotated_file):
|
|
return error_resp(404, f"file {filename} not found")
|
|
|
|
filepath = folder_paths.get_annotated_filepath(annotated_file)
|
|
os.remove(filepath)
|
|
return success_resp()
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
@routes.get("/api-tools/v1/metrics")
|
|
async def get_prometheus_metrics(request: Request):
|
|
"""Return Prometheus metrics"""
|
|
try:
|
|
metrics_text = get_metrics()
|
|
return json_response(text=metrics_text, content_type="text/plain")
|
|
except Exception as e:
|
|
logging.error(f"Error generating metrics: {str(e)}")
|
|
return error_resp(500, f"Error generating metrics: {str(e)}")
|
|
|
|
|
|
def run_comfyui_api_tools():
|
|
print("extra API server started") |