160 lines
5.1 KiB
Python
160 lines
5.1 KiB
Python
import logging
|
|
import os
|
|
import traceback
|
|
|
|
from aiohttp.web import Request, json_response
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
from .model_utils.refresh import refresh_folder
|
|
from .model_utils.lora import (
|
|
list_available_networks,
|
|
available_networks,
|
|
create_lora_json,
|
|
)
|
|
from .utils.images import base64_decode_to_pil, extract_img_metadata
|
|
|
|
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("/comfyapi/v1/checkpoints")
|
|
async def get_checkpoints(request: Request):
|
|
try:
|
|
checkpoints = folder_paths.get_filename_list("checkpoints")
|
|
result = [
|
|
{
|
|
"name": os.path.basename(ckpt),
|
|
"path": ckpt,
|
|
"full_path": folder_paths.get_full_path("checkpoints", ckpt) or ckpt,
|
|
}
|
|
for ckpt in checkpoints
|
|
]
|
|
return success_resp(result=result)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.post("/comfyapi/v1/refresh-checkpoints")
|
|
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)"""
|
|
try:
|
|
data = refresh_folder("checkpoints")
|
|
return success_resp(data=data)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.get("/comfyapi/v1/loras")
|
|
async def get_loras(request: Request):
|
|
try:
|
|
loras = loras = [
|
|
create_lora_json(obj=obj) for obj in available_networks.values()
|
|
]
|
|
return success_resp(loras=loras)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.post("/comfyapi/v1/refresh-loras")
|
|
async def refresh_loras(request: Request):
|
|
try:
|
|
data = refresh_folder("loras")
|
|
list_available_networks()
|
|
return success_resp(data=data)
|
|
except Exception as e:
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
@routes.get("/comfyapi/v1/output-images")
|
|
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", ".jpg", ".jpeg", ".webp", ".gif")):
|
|
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("/comfyapi/v1/output-images/{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.delete("/comfyapi/v1/input-images/{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.post("/comfyapi/v1/pnginfo")
|
|
async def get_png_info(request: Request):
|
|
try:
|
|
data = await request.json()
|
|
img_base64 = data.get("img_base64")
|
|
if img_base64 is None:
|
|
return error_resp(400, "img_base64 is required")
|
|
|
|
img = base64_decode_to_pil(img_base64)
|
|
metadata = extract_img_metadata(img)
|
|
return success_resp(metadata=metadata)
|
|
except Exception as e:
|
|
err = traceback.format_exc()
|
|
logging.error(err)
|
|
return error_resp(500, str(e))
|
|
|
|
|
|
def run_comfyui_extra_api():
|
|
print("extra API server started")
|