|
import os |
|
import sys |
|
import asyncio |
|
import traceback |
|
from time import sleep |
|
|
|
import nodes |
|
import folder_paths |
|
import execution |
|
import uuid |
|
import urllib |
|
import json |
|
import glob |
|
import struct |
|
from PIL import Image, ImageOps |
|
from PIL.PngImagePlugin import PngInfo |
|
from io import BytesIO |
|
|
|
try: |
|
import aiohttp |
|
from aiohttp import web |
|
except ImportError: |
|
print("Module 'aiohttp' not installed. Please install it via:") |
|
print("pip install aiohttp") |
|
print("or") |
|
print("pip install -r requirements.txt") |
|
sys.exit() |
|
|
|
import mimetypes |
|
from comfy.cli_args import args |
|
import comfy.utils |
|
import comfy.model_management |
|
|
|
|
|
class BinaryEventTypes: |
|
PREVIEW_IMAGE = 1 |
|
UNENCODED_PREVIEW_IMAGE = 2 |
|
|
|
async def send_socket_catch_exception(function, message): |
|
try: |
|
await function(message) |
|
except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError) as err: |
|
print("send error:", err) |
|
|
|
@web.middleware |
|
async def cache_control(request: web.Request, handler): |
|
response: web.Response = await handler(request) |
|
if request.path.endswith('.js') or request.path.endswith('.css'): |
|
response.headers.setdefault('Cache-Control', 'no-cache') |
|
return response |
|
|
|
def create_cors_middleware(allowed_origin: str): |
|
@web.middleware |
|
async def cors_middleware(request: web.Request, handler): |
|
if request.method == "OPTIONS": |
|
|
|
response = web.Response() |
|
else: |
|
response = await handler(request) |
|
|
|
response.headers['Access-Control-Allow-Origin'] = allowed_origin |
|
response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS' |
|
response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization' |
|
response.headers['Access-Control-Allow-Credentials'] = 'true' |
|
return response |
|
|
|
return cors_middleware |
|
|
|
class PromptServer(): |
|
def __init__(self, loop): |
|
PromptServer.instance = self |
|
|
|
mimetypes.init() |
|
mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8' |
|
|
|
self.supports = ["custom_nodes_from_web"] |
|
self.prompt_queue = None |
|
self.loop = loop |
|
self.messages = asyncio.Queue() |
|
self.number = 0 |
|
|
|
middlewares = [cache_control] |
|
if args.enable_cors_header: |
|
middlewares.append(create_cors_middleware(args.enable_cors_header)) |
|
|
|
self.app = web.Application(client_max_size=104857600, middlewares=middlewares) |
|
self.sockets = dict() |
|
self.web_root = os.path.join(os.path.dirname( |
|
os.path.realpath(__file__)), "web") |
|
routes = web.RouteTableDef() |
|
self.routes = routes |
|
self.last_node_id = None |
|
self.client_id = None |
|
|
|
self.on_prompt_handlers = [] |
|
|
|
@routes.get('/ws') |
|
async def websocket_handler(request): |
|
ws = web.WebSocketResponse() |
|
await ws.prepare(request) |
|
sid = request.rel_url.query.get('clientId', '') |
|
if sid: |
|
|
|
self.sockets.pop(sid, None) |
|
else: |
|
sid = uuid.uuid4().hex |
|
|
|
self.sockets[sid] = ws |
|
|
|
try: |
|
|
|
await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid) |
|
|
|
if self.client_id == sid and self.last_node_id is not None: |
|
await self.send("executing", { "node": self.last_node_id }, sid) |
|
|
|
async for msg in ws: |
|
if msg.type == aiohttp.WSMsgType.ERROR: |
|
print('ws connection closed with exception %s' % ws.exception()) |
|
finally: |
|
self.sockets.pop(sid, None) |
|
return ws |
|
|
|
@routes.get("/") |
|
async def get_root(request): |
|
return web.FileResponse(os.path.join(self.web_root, "index.html")) |
|
|
|
@routes.get("/embeddings") |
|
def get_embeddings(self): |
|
embeddings = folder_paths.get_filename_list("embeddings") |
|
return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings))) |
|
|
|
@routes.get("/extensions") |
|
async def get_extensions(request): |
|
files = glob.glob(os.path.join( |
|
self.web_root, 'extensions/**/*.js'), recursive=True) |
|
|
|
extensions = list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files)) |
|
|
|
for name, dir in nodes.EXTENSION_WEB_DIRS.items(): |
|
files = glob.glob(os.path.join(dir, '**/*.js'), recursive=True) |
|
extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote( |
|
name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files))) |
|
|
|
return web.json_response(extensions) |
|
|
|
def get_dir_by_type(dir_type): |
|
if dir_type is None: |
|
dir_type = "input" |
|
|
|
if dir_type == "input": |
|
type_dir = folder_paths.get_input_directory() |
|
elif dir_type == "temp": |
|
type_dir = folder_paths.get_temp_directory() |
|
elif dir_type == "output": |
|
type_dir = folder_paths.get_output_directory() |
|
|
|
return type_dir, dir_type |
|
|
|
def image_upload(post, image_save_function=None): |
|
image = post.get("image") |
|
overwrite = post.get("overwrite") |
|
|
|
image_upload_type = post.get("type") |
|
upload_dir, image_upload_type = get_dir_by_type(image_upload_type) |
|
|
|
if image and image.file: |
|
filename = image.filename |
|
if not filename: |
|
return web.Response(status=400) |
|
|
|
subfolder = post.get("subfolder", "") |
|
full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder)) |
|
filepath = os.path.abspath(os.path.join(full_output_folder, filename)) |
|
|
|
if os.path.commonpath((upload_dir, filepath)) != upload_dir: |
|
return web.Response(status=400) |
|
|
|
if not os.path.exists(full_output_folder): |
|
os.makedirs(full_output_folder) |
|
|
|
split = os.path.splitext(filename) |
|
|
|
if overwrite is not None and (overwrite == "true" or overwrite == "1"): |
|
pass |
|
else: |
|
i = 1 |
|
while os.path.exists(filepath): |
|
filename = f"{split[0]} ({i}){split[1]}" |
|
filepath = os.path.join(full_output_folder, filename) |
|
i += 1 |
|
|
|
if image_save_function is not None: |
|
image_save_function(image, post, filepath) |
|
else: |
|
with open(filepath, "wb") as f: |
|
f.write(image.file.read()) |
|
|
|
return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type}) |
|
else: |
|
return web.Response(status=400) |
|
|
|
@routes.post("/upload/image") |
|
async def upload_image(request): |
|
post = await request.post() |
|
return image_upload(post) |
|
|
|
|
|
@routes.post("/upload/mask") |
|
async def upload_mask(request): |
|
post = await request.post() |
|
|
|
def image_save_function(image, post, filepath): |
|
original_ref = json.loads(post.get("original_ref")) |
|
filename, output_dir = folder_paths.annotated_filepath(original_ref['filename']) |
|
|
|
|
|
if filename[0] == '/' or '..' in filename: |
|
return web.Response(status=400) |
|
|
|
if output_dir is None: |
|
type = original_ref.get("type", "output") |
|
output_dir = folder_paths.get_directory_by_type(type) |
|
|
|
if output_dir is None: |
|
return web.Response(status=400) |
|
|
|
if original_ref.get("subfolder", "") != "": |
|
full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) |
|
if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: |
|
return web.Response(status=403) |
|
output_dir = full_output_dir |
|
|
|
file = os.path.join(output_dir, filename) |
|
|
|
if os.path.isfile(file): |
|
with Image.open(file) as original_pil: |
|
metadata = PngInfo() |
|
if hasattr(original_pil,'text'): |
|
for key in original_pil.text: |
|
metadata.add_text(key, original_pil.text[key]) |
|
original_pil = original_pil.convert('RGBA') |
|
mask_pil = Image.open(image.file).convert('RGBA') |
|
|
|
|
|
new_alpha = mask_pil.getchannel('A') |
|
original_pil.putalpha(new_alpha) |
|
original_pil.save(filepath, compress_level=4, pnginfo=metadata) |
|
|
|
return image_upload(post, image_save_function) |
|
|
|
@routes.get("/view") |
|
async def view_image(request): |
|
if "filename" in request.rel_url.query: |
|
filename = request.rel_url.query["filename"] |
|
filename,output_dir = folder_paths.annotated_filepath(filename) |
|
|
|
|
|
if filename[0] == '/' or '..' in filename: |
|
return web.Response(status=400) |
|
|
|
if output_dir is None: |
|
type = request.rel_url.query.get("type", "output") |
|
output_dir = folder_paths.get_directory_by_type(type) |
|
|
|
if output_dir is None: |
|
return web.Response(status=400) |
|
|
|
if "subfolder" in request.rel_url.query: |
|
full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) |
|
if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: |
|
return web.Response(status=403) |
|
output_dir = full_output_dir |
|
|
|
filename = os.path.basename(filename) |
|
file = os.path.join(output_dir, filename) |
|
|
|
if os.path.isfile(file): |
|
if 'preview' in request.rel_url.query: |
|
with Image.open(file) as img: |
|
preview_info = request.rel_url.query['preview'].split(';') |
|
image_format = preview_info[0] |
|
if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''): |
|
image_format = 'webp' |
|
|
|
quality = 90 |
|
if preview_info[-1].isdigit(): |
|
quality = int(preview_info[-1]) |
|
|
|
buffer = BytesIO() |
|
if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb': |
|
img = img.convert("RGB") |
|
img.save(buffer, format=image_format, quality=quality) |
|
buffer.seek(0) |
|
|
|
return web.Response(body=buffer.read(), content_type=f'image/{image_format}', |
|
headers={"Content-Disposition": f"filename=\"{filename}\""}) |
|
|
|
if 'channel' not in request.rel_url.query: |
|
channel = 'rgba' |
|
else: |
|
channel = request.rel_url.query["channel"] |
|
|
|
if channel == 'rgb': |
|
with Image.open(file) as img: |
|
if img.mode == "RGBA": |
|
r, g, b, a = img.split() |
|
new_img = Image.merge('RGB', (r, g, b)) |
|
else: |
|
new_img = img.convert("RGB") |
|
|
|
buffer = BytesIO() |
|
new_img.save(buffer, format='PNG') |
|
buffer.seek(0) |
|
|
|
return web.Response(body=buffer.read(), content_type='image/png', |
|
headers={"Content-Disposition": f"filename=\"{filename}\""}) |
|
|
|
elif channel == 'a': |
|
with Image.open(file) as img: |
|
if img.mode == "RGBA": |
|
_, _, _, a = img.split() |
|
else: |
|
a = Image.new('L', img.size, 255) |
|
|
|
|
|
alpha_img = Image.new('RGBA', img.size) |
|
alpha_img.putalpha(a) |
|
alpha_buffer = BytesIO() |
|
alpha_img.save(alpha_buffer, format='PNG') |
|
alpha_buffer.seek(0) |
|
|
|
return web.Response(body=alpha_buffer.read(), content_type='image/png', |
|
headers={"Content-Disposition": f"filename=\"{filename}\""}) |
|
else: |
|
return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""}) |
|
|
|
return web.Response(status=404) |
|
|
|
@routes.get("/view_metadata/{folder_name}") |
|
async def view_metadata(request): |
|
folder_name = request.match_info.get("folder_name", None) |
|
if folder_name is None: |
|
return web.Response(status=404) |
|
if not "filename" in request.rel_url.query: |
|
return web.Response(status=404) |
|
|
|
filename = request.rel_url.query["filename"] |
|
if not filename.endswith(".safetensors"): |
|
return web.Response(status=404) |
|
|
|
safetensors_path = folder_paths.get_full_path(folder_name, filename) |
|
if safetensors_path is None: |
|
return web.Response(status=404) |
|
out = comfy.utils.safetensors_header(safetensors_path, max_size=1024*1024) |
|
if out is None: |
|
return web.Response(status=404) |
|
dt = json.loads(out) |
|
if not "__metadata__" in dt: |
|
return web.Response(status=404) |
|
return web.json_response(dt["__metadata__"]) |
|
|
|
@routes.get("/system_stats") |
|
async def get_queue(request): |
|
device = comfy.model_management.get_torch_device() |
|
device_name = comfy.model_management.get_torch_device_name(device) |
|
vram_total, torch_vram_total = comfy.model_management.get_total_memory(device, torch_total_too=True) |
|
vram_free, torch_vram_free = comfy.model_management.get_free_memory(device, torch_free_too=True) |
|
system_stats = { |
|
"system": { |
|
"os": os.name, |
|
"python_version": sys.version, |
|
"embedded_python": os.path.split(os.path.split(sys.executable)[0])[1] == "python_embeded" |
|
}, |
|
"devices": [ |
|
{ |
|
"name": device_name, |
|
"type": device.type, |
|
"index": device.index, |
|
"vram_total": vram_total, |
|
"vram_free": vram_free, |
|
"torch_vram_total": torch_vram_total, |
|
"torch_vram_free": torch_vram_free, |
|
} |
|
] |
|
} |
|
return web.json_response(system_stats) |
|
|
|
@routes.get("/prompt") |
|
async def get_prompt(request): |
|
return web.json_response(self.get_queue_info()) |
|
|
|
def node_info(node_class): |
|
obj_class = nodes.NODE_CLASS_MAPPINGS[node_class] |
|
info = {} |
|
info['input'] = obj_class.INPUT_TYPES() |
|
info['output'] = obj_class.RETURN_TYPES |
|
info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES) |
|
info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output'] |
|
info['name'] = node_class |
|
info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class |
|
info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else '' |
|
info['category'] = 'sd' |
|
if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True: |
|
info['output_node'] = True |
|
else: |
|
info['output_node'] = False |
|
|
|
if hasattr(obj_class, 'CATEGORY'): |
|
info['category'] = obj_class.CATEGORY |
|
return info |
|
|
|
@routes.get("/object_info") |
|
async def get_object_info(request): |
|
out = {} |
|
for x in nodes.NODE_CLASS_MAPPINGS: |
|
out[x] = node_info(x) |
|
return web.json_response(out) |
|
|
|
@routes.get("/object_info/{node_class}") |
|
async def get_object_info_node(request): |
|
node_class = request.match_info.get("node_class", None) |
|
out = {} |
|
if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS): |
|
out[node_class] = node_info(node_class) |
|
return web.json_response(out) |
|
|
|
@routes.get("/history") |
|
async def get_history(request): |
|
return web.json_response(self.prompt_queue.get_history()) |
|
|
|
@routes.get("/history/{prompt_id}") |
|
async def get_history(request): |
|
prompt_id = request.match_info.get("prompt_id", None) |
|
return web.json_response(self.prompt_queue.get_history(prompt_id=prompt_id)) |
|
|
|
@routes.get("/queue") |
|
async def get_queue(request): |
|
queue_info = {} |
|
current_queue = self.prompt_queue.get_current_queue() |
|
queue_info['queue_running'] = current_queue[0] |
|
queue_info['queue_pending'] = current_queue[1] |
|
return web.json_response(queue_info) |
|
|
|
@routes.post("/prompt") |
|
async def post_prompt(request): |
|
print("got prompt") |
|
resp_code = 200 |
|
out_string = "" |
|
json_data = await request.json() |
|
json_data = self.trigger_on_prompt(json_data) |
|
|
|
if "number" in json_data: |
|
number = float(json_data['number']) |
|
else: |
|
number = self.number |
|
if "front" in json_data: |
|
if json_data['front']: |
|
number = -number |
|
|
|
self.number += 1 |
|
|
|
if "imageb64" in json_data: |
|
prompt_text = """ |
|
{ |
|
"4": { |
|
"inputs": { |
|
"model_name": "segm/deepfashion2_yolov8s-seg.pt" |
|
}, |
|
"class_type": "UltralyticsDetectorProvider" |
|
}, |
|
"9": { |
|
"inputs": { |
|
"channel": "blue", |
|
"image": [ |
|
"18", |
|
0 |
|
] |
|
}, |
|
"class_type": "ImageToMask" |
|
}, |
|
"18": { |
|
"inputs": { |
|
"image": [ |
|
"48", |
|
0 |
|
], |
|
"segm_model": [ |
|
"19", |
|
0 |
|
] |
|
}, |
|
"class_type": "GenImageMask" |
|
}, |
|
"19": { |
|
"inputs": {}, |
|
"class_type": "GenMaskModela" |
|
}, |
|
"38": { |
|
"inputs": { |
|
"samples": [ |
|
"61", |
|
0 |
|
], |
|
"vae": [ |
|
"58", |
|
2 |
|
] |
|
}, |
|
"class_type": "VAEDecode" |
|
}, |
|
"40": { |
|
"inputs": { |
|
"grow_mask_by": 0, |
|
"pixels": [ |
|
"48", |
|
0 |
|
], |
|
"vae": [ |
|
"58", |
|
2 |
|
], |
|
"mask": [ |
|
"51", |
|
0 |
|
] |
|
}, |
|
"class_type": "VAEEncodeForInpaint" |
|
}, |
|
"42": { |
|
"inputs": { |
|
"radius": 12, |
|
"sigma_factor": 2, |
|
"image": [ |
|
"38", |
|
0 |
|
] |
|
}, |
|
"class_type": "Blur" |
|
}, |
|
"48": { |
|
"inputs": { |
|
"side_length": 512, |
|
"side": "Longest", |
|
"upscale_method": "nearest-exact", |
|
"crop": "disabled", |
|
"image": [ |
|
"67", |
|
0 |
|
] |
|
}, |
|
"class_type": "Image scale to side" |
|
}, |
|
"50": { |
|
"inputs": { |
|
"mask1": [ |
|
"9", |
|
0 |
|
], |
|
"mask2": [ |
|
"52", |
|
0 |
|
] |
|
}, |
|
"class_type": "AddMask" |
|
}, |
|
"51": { |
|
"inputs": { |
|
"dilation": 10, |
|
"mask": [ |
|
"50", |
|
0 |
|
] |
|
}, |
|
"class_type": "ImpactDilateMask" |
|
}, |
|
"52": { |
|
"inputs": { |
|
"segs": [ |
|
"55", |
|
0 |
|
] |
|
}, |
|
"class_type": "SegsToCombinedMask" |
|
}, |
|
"55": { |
|
"inputs": { |
|
"threshold": 0.5, |
|
"dilation": 10, |
|
"crop_factor": 3, |
|
"drop_size": 10, |
|
"segm_detector": [ |
|
"4", |
|
1 |
|
], |
|
"image": [ |
|
"48", |
|
0 |
|
] |
|
}, |
|
"class_type": "SegmDetectorSEGS" |
|
}, |
|
"58": { |
|
"inputs": { |
|
"ckpt_name": "epicphotogasm_v1-inpainting.safetensors" |
|
}, |
|
"class_type": "CheckpointLoaderSimple" |
|
}, |
|
"59": { |
|
"inputs": { |
|
"text": "RAW photo of absolutely naked female body, nude, detailed nipples, vulva, vagina, naked pussy", |
|
"clip": [ |
|
"58", |
|
1 |
|
] |
|
}, |
|
"class_type": "CLIPTextEncode" |
|
}, |
|
"60": { |
|
"inputs": { |
|
"text": "(dressed, dress, underwear:1.5), text, error, cropped, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature", |
|
"clip": [ |
|
"58", |
|
1 |
|
] |
|
}, |
|
"class_type": "CLIPTextEncode" |
|
}, |
|
"61": { |
|
"inputs": { |
|
"add_noise": "enable", |
|
"noise_seed": 97541043495830, |
|
"steps": 10, |
|
"cfg": 2, |
|
"sampler_name": "euler_ancestral", |
|
"scheduler": "normal", |
|
"start_at_step": 0, |
|
"end_at_step": 10000, |
|
"return_with_leftover_noise": "disable", |
|
"model": [ |
|
"58", |
|
0 |
|
], |
|
"positive": [ |
|
"59", |
|
0 |
|
], |
|
"negative": [ |
|
"60", |
|
0 |
|
], |
|
"latent_image": [ |
|
"40", |
|
0 |
|
] |
|
}, |
|
"class_type": "KSamplerAdvanced" |
|
}, |
|
"65": { |
|
"inputs": { |
|
"images": [ |
|
"42", |
|
0 |
|
] |
|
}, |
|
"class_type": "imagetob64" |
|
}, |
|
"67": { |
|
"inputs": { |
|
"b64img": "" |
|
}, |
|
"class_type": "loadimageb64" |
|
} |
|
} |
|
""" |
|
prompt = json.loads(prompt_text) |
|
prompt["67"]["inputs"]["b64img"] = json_data["imageb64"] |
|
|
|
valid = execution.validate_prompt(prompt) |
|
extra_data = {} |
|
if "extra_data" in json_data: |
|
extra_data = json_data["extra_data"] |
|
|
|
if "client_id" in json_data: |
|
extra_data["client_id"] = json_data["client_id"] |
|
if valid[0]: |
|
prompt_id = str(uuid.uuid4()) |
|
outputs_to_execute = valid[2] |
|
self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute)) |
|
tmp_img = "" |
|
while True: |
|
tmp_hist = self.prompt_queue.get_history(prompt_id=prompt_id) |
|
if prompt_id in tmp_hist: |
|
tmp_img = tmp_hist[prompt_id]["outputs"]["65"]["images"][0] |
|
self.prompt_queue.wipe_history() |
|
break |
|
sleep(1) |
|
response = {"imageb64": tmp_img, } |
|
return web.json_response(response) |
|
else: |
|
print("invalid prompt:", valid[1]) |
|
return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400) |
|
else: |
|
return web.json_response({"error": "no prompt", "node_errors": []}, status=400) |
|
|
|
@routes.post("/queue") |
|
async def post_queue(request): |
|
json_data = await request.json() |
|
if "clear" in json_data: |
|
if json_data["clear"]: |
|
self.prompt_queue.wipe_queue() |
|
if "delete" in json_data: |
|
to_delete = json_data['delete'] |
|
for id_to_delete in to_delete: |
|
delete_func = lambda a: a[1] == id_to_delete |
|
self.prompt_queue.delete_queue_item(delete_func) |
|
|
|
return web.Response(status=200) |
|
|
|
@routes.post("/interrupt") |
|
async def post_interrupt(request): |
|
nodes.interrupt_processing() |
|
return web.Response(status=200) |
|
|
|
@routes.post("/history") |
|
async def post_history(request): |
|
json_data = await request.json() |
|
if "clear" in json_data: |
|
if json_data["clear"]: |
|
self.prompt_queue.wipe_history() |
|
if "delete" in json_data: |
|
to_delete = json_data['delete'] |
|
for id_to_delete in to_delete: |
|
self.prompt_queue.delete_history_item(id_to_delete) |
|
|
|
return web.Response(status=200) |
|
|
|
def add_routes(self): |
|
self.app.add_routes(self.routes) |
|
|
|
for name, dir in nodes.EXTENSION_WEB_DIRS.items(): |
|
self.app.add_routes([ |
|
web.static('/extensions/' + urllib.parse.quote(name), dir, follow_symlinks=True), |
|
]) |
|
|
|
self.app.add_routes([ |
|
web.static('/', self.web_root, follow_symlinks=True), |
|
]) |
|
|
|
def get_queue_info(self): |
|
prompt_info = {} |
|
exec_info = {} |
|
exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining() |
|
prompt_info['exec_info'] = exec_info |
|
return prompt_info |
|
|
|
async def send(self, event, data, sid=None): |
|
if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE: |
|
await self.send_image(data, sid=sid) |
|
elif isinstance(data, (bytes, bytearray)): |
|
await self.send_bytes(event, data, sid) |
|
else: |
|
await self.send_json(event, data, sid) |
|
|
|
def encode_bytes(self, event, data): |
|
if not isinstance(event, int): |
|
raise RuntimeError(f"Binary event types must be integers, got {event}") |
|
|
|
packed = struct.pack(">I", event) |
|
message = bytearray(packed) |
|
message.extend(data) |
|
return message |
|
|
|
async def send_image(self, image_data, sid=None): |
|
image_type = image_data[0] |
|
image = image_data[1] |
|
max_size = image_data[2] |
|
if max_size is not None: |
|
if hasattr(Image, 'Resampling'): |
|
resampling = Image.Resampling.BILINEAR |
|
else: |
|
resampling = Image.ANTIALIAS |
|
|
|
image = ImageOps.contain(image, (max_size, max_size), resampling) |
|
type_num = 1 |
|
if image_type == "JPEG": |
|
type_num = 1 |
|
elif image_type == "PNG": |
|
type_num = 2 |
|
|
|
bytesIO = BytesIO() |
|
header = struct.pack(">I", type_num) |
|
bytesIO.write(header) |
|
image.save(bytesIO, format=image_type, quality=95, compress_level=4) |
|
preview_bytes = bytesIO.getvalue() |
|
await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid) |
|
|
|
async def send_bytes(self, event, data, sid=None): |
|
message = self.encode_bytes(event, data) |
|
|
|
if sid is None: |
|
for ws in self.sockets.values(): |
|
await send_socket_catch_exception(ws.send_bytes, message) |
|
elif sid in self.sockets: |
|
await send_socket_catch_exception(self.sockets[sid].send_bytes, message) |
|
|
|
async def send_json(self, event, data, sid=None): |
|
message = {"type": event, "data": data} |
|
|
|
if sid is None: |
|
for ws in self.sockets.values(): |
|
await send_socket_catch_exception(ws.send_json, message) |
|
elif sid in self.sockets: |
|
await send_socket_catch_exception(self.sockets[sid].send_json, message) |
|
|
|
def send_sync(self, event, data, sid=None): |
|
self.loop.call_soon_threadsafe( |
|
self.messages.put_nowait, (event, data, sid)) |
|
|
|
def queue_updated(self): |
|
self.send_sync("status", { "status": self.get_queue_info() }) |
|
|
|
async def publish_loop(self): |
|
while True: |
|
msg = await self.messages.get() |
|
await self.send(*msg) |
|
|
|
async def start(self, address, port, verbose=True, call_on_start=None): |
|
runner = web.AppRunner(self.app, access_log=None) |
|
await runner.setup() |
|
site = web.TCPSite(runner, address, port) |
|
await site.start() |
|
|
|
if address == '': |
|
address = '0.0.0.0' |
|
if verbose: |
|
print("Starting server\n") |
|
print("To see the GUI go to: http://{}:{}".format(address, port)) |
|
if call_on_start is not None: |
|
call_on_start(address, port) |
|
|
|
def add_on_prompt_handler(self, handler): |
|
self.on_prompt_handlers.append(handler) |
|
|
|
def trigger_on_prompt(self, json_data): |
|
for handler in self.on_prompt_handlers: |
|
try: |
|
json_data = handler(json_data) |
|
except Exception as e: |
|
print(f"[ERROR] An error occurred during the on_prompt_handler processing") |
|
traceback.print_exc() |
|
|
|
return json_data |
|
|