backend_base/comfyclient_pipeline.py

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