Spaces:
Configuration error
Configuration error
File size: 7,455 Bytes
aebc1d4 |
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 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
import os
import json
import websocket
import uuid
import random
import httpx
import time
import numpy as np
import ldm_patched.modules.model_management as model_management
from io import BytesIO
from PIL import Image
from . import utils
def upload_mask(mask):
with BytesIO() as output:
mask.save(output)
output.seek(0)
files = {'mask': ('mask.jpg', output)}
data = {'overwrite': 'true', 'type': 'example_type'}
response = httpx.post("http://{}/upload/mask".format(server_address), files=files, data=data)
return response.json()
def queue_prompt(prompt):
p = {"prompt": prompt, "client_id": client_id}
data = json.dumps(p).encode('utf-8')
try:
with httpx.Client() as client:
response = client.post("http://{}/prompt".format(server_address), data=data)
return json.loads(response.read())
except httpx.RequestError as e:
print(f"httpx.RequestError: {e}")
return None
def get_image(filename, subfolder, folder_type):
params = httpx.QueryParams({
"filename": filename,
"subfolder": subfolder,
"type": folder_type
})
with httpx.Client() as client:
response = client.get(f"http://{server_address}/view", params=params)
return response.read()
def get_history(prompt_id):
with httpx.Client() as client:
response = client.get("http://{}/history/{}".format(server_address, prompt_id))
return json.loads(response.read())
def get_images(ws, prompt, callback=None):
prompt_id = queue_prompt(prompt)['prompt_id']
print('[ComfyClient] Request and get ComfyTask_id:{}'.format(prompt_id))
output_images = {}
current_node = ''
last_node = None
preview_image = []
last_step = None
current_step = None
current_total_steps = None
while True:
model_management.throw_exception_if_processing_interrupted()
try:
out = ws.recv()
except ConnectionResetError as e:
print(f'[ComfyClient] The connect was exception, restart and try again: {e}')
ws = websocket.WebSocket()
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
out = ws.recv()
if isinstance(out, str):
message = json.loads(out)
current_type = message['type']
#print(f'current_message={message}')
if message['type'] == 'executing':
data = message['data']
if data['node'] is None and data['prompt_id'] == prompt_id:
break
else:
current_node = data['node']
elif message['type'] == 'progress':
current_step = message["data"]["value"]
current_total_steps = message["data"]["max"]
else:
if current_type == 'progress':
if prompt[current_node]['class_type'] in ['KSampler', 'SamplerCustomAdvanced', 'TiledKSampler'] and callback is not None:
if current_step == last_step:
preview_image.append(out[8:])
else:
if last_step is not None:
callback(last_step, current_total_steps, Image.open(BytesIO(preview_image[0])))
preview_image = []
preview_image.append(out[8:])
last_step = current_step
if prompt[current_node]['class_type'] == 'SaveImageWebsocket':
images_output = output_images.get(prompt[current_node]['_meta']['title'], [])
images_output.append(out[8:])
output_images[prompt[current_node]['_meta']['title']] = images_output[0]
continue
output_images = {k: np.array(Image.open(BytesIO(v))) for k, v in output_images.items()}
print(f'[ComfyClient] The ComfyTask:{prompt_id} has finished: {len(output_images)}')
return output_images
def images_upload(images):
result = {}
if images is None:
return result
for k,np_image in images.items():
pil_image = Image.fromarray(np_image)
with BytesIO() as output:
pil_image.save(output, format="PNG")
output.seek(0)
files = {'image': (f'image_{client_id}_{random.randint(1000, 9999)}.png', output)}
data = {'overwrite': 'true', 'type': 'input'}
response = httpx.post("http://{}/upload/image".format(server_address), files=files, data=data)
result.update({k: response.json()["name"]})
print(f'[ComfyClient] The ComfyTask:upload_input_images has finished: {len(result)}')
return result
def process_flow(flow_name, params, images, callback=None):
global ws
flow_file = os.path.join(WORKFLOW_DIR, f'{flow_name}_api.json')
if ws is None or ws.status != 101:
if ws is not None:
print(f'[ComfyClient] websocket status: {ws.status}, timeout:{ws.timeout}s.')
ws.close()
try:
ws = websocket.WebSocket()
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
except ConnectionRefusedError as e:
print(f'[ComfyClient] The connect_to_server has failed, sleep and try again: {e}')
time.sleep(8)
try:
ws = websocket.WebSocket()
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
except ConnectionRefusedError as e:
print(f'[ComfyClient] The connect_to_server has failed, restart and try again: {e}')
time.sleep(12)
ws = websocket.WebSocket()
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
images_map = images_upload(images)
params.update_params(images_map)
with open(flow_file, 'r', encoding="utf-8") as workflow_api_file:
flowdata = json.load(workflow_api_file)
print(f'[ComfyClient] Ready ComfyTask to process: workflow={flow_name}')
for k,v in params.params.items():
print(f' {k} = {v}')
try:
prompt_str = params.convert2comfy(flowdata)
if not utils.echo_off:
print(f'[ComfyClient] ComfyTask prompt: {prompt_str}')
images = get_images(ws, prompt_str, callback=callback)
except websocket.WebSocketException as e:
print(f'[ComfyClient] The connection has been closed, restart and try again: {e}')
ws = None
images_keys = sorted(images.keys(), reverse=True)
imgs = [images[key] for key in images_keys]
return imgs
def interrupt():
try:
with httpx.Client() as client:
response = client.post("http://{}/interrupt".format(server_address))
return
except httpx.RequestError as e:
print(f"httpx.RequestError: {e}")
return
def free(all=False):
p = {"unload_models": all==True, "free_memory": True}
data = json.dumps(p).encode('utf-8')
try:
with httpx.Client() as client:
response = client.post("http://{}/free".format(server_address), data=data)
return
except httpx.RequestError as e:
print(f"httpx.RequestError: {e}")
return
WORKFLOW_DIR = 'workflows'
COMFYUI_ENDPOINT_IP = '127.0.0.1'
COMFYUI_ENDPOINT_PORT = '8187'
server_address = f'{COMFYUI_ENDPOINT_IP}:{COMFYUI_ENDPOINT_PORT}'
client_id = str(uuid.uuid4())
ws = None
|