Update Dockerfile

adding Postgres port env var

comfy-nodes/external_checkpoints.py

@@ -16,7 +16,7 @@ class ComfyUIDeployExternalCheckpoint:
                 ),
             },
             "optional": {
-                "default_checkpoint_name": (folder_paths.get_filename_list("checkpoints"), ),
+                "default_value": (folder_paths.get_filename_list("checkpoints"), ),
             }
         }
 
@@ -27,12 +27,12 @@ class ComfyUIDeployExternalCheckpoint:
 
     CATEGORY = "deploy"
 
-    def run(self, input_id, default_checkpoint_name=None):
+    def run(self, input_id, default_value=None):
         import requests
         import os
         import uuid
 
-        if input_id and input_id.startswith('http'):
+        if default_value.startswith('http'):
             unique_filename = str(uuid.uuid4()) + ".safetensors"
             print(unique_filename)
             print(folder_paths.folder_names_and_paths["checkpoints"][0][0])
@@ -59,7 +59,7 @@ class ComfyUIDeployExternalCheckpoint:
                     out_file.write(chunk)
             return (unique_filename,)
         else:
-            return (default_checkpoints_name,)
+            return (default_value,)
 
 
 NODE_CLASS_MAPPINGS = {

comfy-nodes/external_number.py

@@ -16,7 +16,7 @@ class ComfyUIDeployExternalNumber:
             "optional": {
                 "default_value": (
                     "FLOAT",
-                    {"multiline": True, "display": "number", "default": 0},
+                    {"multiline": True, "display": "number", "default": 0, "step": 0.01},
                 ),
             }
         }
@@ -29,9 +29,12 @@ class ComfyUIDeployExternalNumber:
     CATEGORY = "number"
 
     def run(self, input_id, default_value=None):
-        if not input_id or not input_id.strip().isdigit():
+        try:
+            float_value = float(input_id)
+            print("my number", float_value)
+            return [float_value]
+        except ValueError:
             return [default_value]
-        return [int(input_id)]
 
 
 NODE_CLASS_MAPPINGS = {"ComfyUIDeployExternalNumber": ComfyUIDeployExternalNumber}

comfy-nodes/input_websocket_image.py

@@ -0,0 +1,66 @@
+import folder_paths
+from PIL import Image, ImageOps
+import numpy as np
+import torch
+from server import PromptServer, BinaryEventTypes
+import asyncio
+
+from globals import streaming_prompt_metadata, max_output_id_length
+
+class ComfyDeployWebscoketImageInput:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input_id": (
+                    "STRING",
+                    {"multiline": False, "default": "input_id"},
+                ),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+            },
+            "optional": {
+                "default_value": ("IMAGE", ),
+                "client_id": (
+                    "STRING",
+                    {"multiline": False, "default": ""},
+                ),
+            }
+        }
+
+    OUTPUT_NODE = True
+
+    RETURN_TYPES = ("IMAGE", )
+    RETURN_NAMES = ("images",)
+
+    FUNCTION = "run"
+
+    @classmethod
+    def VALIDATE_INPUTS(s, input_id):
+        try:
+            if len(input_id.encode('ascii')) > max_output_id_length:
+                raise ValueError(f"input_id size is greater than {max_output_id_length} bytes")
+        except UnicodeEncodeError:
+            raise ValueError("input_id is not ASCII encodable")
+
+        return True
+
+    def run(self, input_id, seed, default_value=None ,client_id=None):
+        # print(streaming_prompt_metadata[client_id].inputs)
+        if client_id in streaming_prompt_metadata and input_id in streaming_prompt_metadata[client_id].inputs:
+            if isinstance(streaming_prompt_metadata[client_id].inputs[input_id], Image.Image):
+                print("Returning image from websocket input")
+                
+                image = streaming_prompt_metadata[client_id].inputs[input_id]
+                
+                image = ImageOps.exif_transpose(image)
+                image = image.convert("RGB")
+                image = np.array(image).astype(np.float32) / 255.0
+                image = torch.from_numpy(image)[None,]
+                
+                return [image]
+
+        print("Returning default value")
+        return [default_value]
+
+NODE_CLASS_MAPPINGS = {"ComfyDeployWebscoketImageInput": ComfyDeployWebscoketImageInput}
+NODE_DISPLAY_NAME_MAPPINGS = {"ComfyDeployWebscoketImageInput": "Image Websocket Input (ComfyDeploy)"}

comfy-nodes/output_websocket_image.py

@@ -0,0 +1,71 @@
+import folder_paths
+from PIL import Image, ImageOps
+import numpy as np
+import torch
+from server import PromptServer, BinaryEventTypes
+import asyncio
+
+from globals import send_image, max_output_id_length
+
+class ComfyDeployWebscoketImageOutput:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "output_id": (
+                    "STRING",
+                    {"multiline": False, "default": "output_id"},
+                ),
+                "images": ("IMAGE", ),
+                "file_type": (["WEBP", "PNG", "JPEG"], ),
+                "quality": ("INT", {"default": 80, "min": 1, "max": 100, "step": 1}),
+            },
+            "optional": {
+                    "client_id": (
+                    "STRING",
+                    {"multiline": False, "default": ""},
+                ),
+            }
+            # "hidden": {"client_id": "CLIENT_ID"},
+        }
+
+    OUTPUT_NODE = True
+
+    RETURN_TYPES = ()
+    RETURN_NAMES = ("text",)
+
+    FUNCTION = "run"
+
+    CATEGORY = "output"
+    
+    @classmethod
+    def VALIDATE_INPUTS(s, output_id):
+        try:
+            if len(output_id.encode('ascii')) > max_output_id_length:
+                raise ValueError(f"output_id size is greater than {max_output_id_length} bytes")
+        except UnicodeEncodeError:
+            raise ValueError("output_id is not ASCII encodable")
+
+        return True
+
+    def run(self, output_id, images, file_type, quality, client_id):
+        prompt_server = PromptServer.instance
+        loop = prompt_server.loop
+        
+        def schedule_coroutine_blocking(target, *args):
+            future = asyncio.run_coroutine_threadsafe(target(*args), loop)
+            return future.result()  # This makes the call blocking
+        
+        for tensor in images:
+            array = 255.0 * tensor.cpu().numpy()
+            image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
+
+            schedule_coroutine_blocking(send_image, [file_type, image, None, quality], client_id, output_id)
+            print("Image sent")
+
+        return {"ui": {}}
+        
+
+
+NODE_CLASS_MAPPINGS = {"ComfyDeployWebscoketImageOutput": ComfyDeployWebscoketImageOutput}
+NODE_DISPLAY_NAME_MAPPINGS = {"ComfyDeployWebscoketImageOutput": "Image Websocket Output (ComfyDeploy)"}

custom_routes.py

@@ -1,38 +1,68 @@
+from io import BytesIO
 from aiohttp import web
 import os
 import requests
 import folder_paths
 import json
-import numpy as np
 import server
-import re
-import base64
 from PIL import Image
-import io
 import time
 import execution
 import random
 import traceback
 import uuid
 import asyncio
-import atexit
 import logging
-import sys
-from logging.handlers import RotatingFileHandler
 from enum import Enum
 from urllib.parse import quote
 import threading
 import hashlib
 import aiohttp
 import aiofiles
-import concurrent.futures
+from typing import List, Union, Any, Optional
+from PIL import Image
+import copy
+import struct
+
+from globals import StreamingPrompt, sockets, streaming_prompt_metadata, BaseModel
+
+class Status(Enum):
+    NOT_STARTED = "not-started"
+    RUNNING = "running"
+    SUCCESS = "success"
+    FAILED = "failed"
+    UPLOADING = "uploading"
+    
+class SimplePrompt(BaseModel):
+    status_endpoint: str
+    file_upload_endpoint: str
+    workflow_api: dict
+    status: Status = Status.NOT_STARTED
+    progress: set = set()
+    last_updated_node: Optional[str] = None,
+    uploading_nodes: set = set()
+    done: bool = False
+    is_realtime: bool = False,
+    start_time: Optional[float] = None,
 
 api = None
 api_task = None
-prompt_metadata = {}
+prompt_metadata: dict[str, SimplePrompt] = {}
 cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'
 cd_enable_run_log = os.environ.get('CD_ENABLE_RUN_LOG', 'false').lower() == 'true'
 
+def clear_current_prompt(sid):
+    prompt_server = server.PromptServer.instance
+    to_delete = list(streaming_prompt_metadata[sid].running_prompt_ids)  # Convert set to list
+    
+    print("clearning out prompt: ", to_delete)
+    for id_to_delete in to_delete:
+        delete_func = lambda a: a[1] == id_to_delete
+        prompt_server.prompt_queue.delete_queue_item(delete_func)
+        print("deleted prompt: ", id_to_delete, prompt_server.prompt_queue.get_tasks_remaining())
+        
+    streaming_prompt_metadata[sid].running_prompt_ids.clear()
+
 def post_prompt(json_data):
     prompt_server = server.PromptServer.instance
     json_data = prompt_server.trigger_on_prompt(json_data)
@@ -80,6 +110,82 @@ def randomSeed(num_digits=15):
     range_end = (10**num_digits) - 1
     return random.randint(range_start, range_end)
 
+def apply_random_seed_to_workflow(workflow_api):
+    """
+    Applies a random seed to each element in the workflow_api that has a 'seed' input.
+
+    Args:
+        workflow_api (dict): The workflow API dictionary to modify.
+    """
+    for key in workflow_api:
+        if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']:
+            if isinstance(workflow_api[key]['inputs']['seed'], list):
+                continue
+            if workflow_api[key]['class_type'] == "PromptExpansion":
+                workflow_api[key]['inputs']['seed'] = randomSeed(8);
+                continue
+            workflow_api[key]['inputs']['seed'] = randomSeed();
+
+def send_prompt(sid: str, inputs: StreamingPrompt):
+    # workflow_api = inputs.workflow_api
+    workflow_api = copy.deepcopy(inputs.workflow_api)
+    
+    # Random seed
+    apply_random_seed_to_workflow(workflow_api)
+            
+    print("getting inputs" , inputs.inputs)
+            
+    # Loop through each of the inputs and replace them
+    for key, value in workflow_api.items():
+        if 'inputs' in value:
+            if (value["class_type"] == "ComfyDeployWebscoketImageOutput"):
+                value['inputs']["client_id"] = sid
+            if (value["class_type"] == "ComfyDeployWebscoketImageInput"):
+                value['inputs']["client_id"] = sid
+                
+            if "input_id" in value['inputs'] and value['inputs']['input_id'] in inputs.inputs:
+                new_value = inputs.inputs[value['inputs']['input_id']]
+                
+                # Lets skip it if its an image
+                if isinstance(new_value, Image.Image):
+                    continue
+                
+                value['inputs']["input_id"] = new_value
+            
+                # Fix for external text default value
+                if (value["class_type"] == "ComfyUIDeployExternalText"):
+                    value['inputs']["default_value"] = new_value
+                    
+                if (value["class_type"] == "ComfyUIDeployExternalCheckpoint"):
+                    value['inputs']["default_value"] = new_value
+
+                
+    print(workflow_api)
+                
+    prompt_id = str(uuid.uuid4())
+            
+    prompt = {
+        "prompt": workflow_api,
+        "client_id": sid, #"comfy_deploy_instance", #api.client_id
+        "prompt_id": prompt_id
+    }
+    
+    try:
+        res = post_prompt(prompt)
+        inputs.running_prompt_ids.add(prompt_id)
+        prompt_metadata[prompt_id] = SimplePrompt(
+            status_endpoint=inputs.status_endpoint,
+            file_upload_endpoint=inputs.file_upload_endpoint,
+            workflow_api=workflow_api,
+            is_realtime=True
+        )
+    except Exception as e:
+        error_type = type(e).__name__
+        stack_trace_short = traceback.format_exc().strip().split('\n')[-2]
+        stack_trace = traceback.format_exc().strip()
+        print(f"error: {error_type}, {e}")
+        print(f"stack trace: {stack_trace_short}")
+       
 @server.PromptServer.instance.routes.post("/comfyui-deploy/run")
 async def comfy_deploy_run(request):
     prompt_server = server.PromptServer.instance
@@ -90,9 +196,10 @@ async def comfy_deploy_run(request):
     # The prompt id generated from comfy deploy, can be None
     prompt_id = data.get("prompt_id")
 
-    for key in workflow_api:
-        if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']:
-            workflow_api[key]['inputs']['seed'] = randomSeed()
+    apply_random_seed_to_workflow(workflow_api)
+    # for key in workflow_api:
+    #     if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']:
+    #         workflow_api[key]['inputs']['seed'] = randomSeed()
 
     prompt = {
         "prompt": workflow_api,
@@ -100,11 +207,11 @@ async def comfy_deploy_run(request):
         "prompt_id": prompt_id
     }
 
-    prompt_metadata[prompt_id] = {
-        'status_endpoint': data.get('status_endpoint'),
-        'file_upload_endpoint': data.get('file_upload_endpoint'),
-        'workflow_api': workflow_api
-    }
+    prompt_metadata[prompt_id] = SimplePrompt(
+        status_endpoint=data.get('status_endpoint'),
+        file_upload_endpoint=data.get('file_upload_endpoint'),
+        workflow_api=workflow_api
+    )
 
     try:
         res = post_prompt(prompt)
@@ -148,7 +255,6 @@ async def comfy_deploy_run(request):
 
     return web.json_response(res, status=status)
 
-sockets = dict()
 
 def get_comfyui_path_from_file_path(file_path):
     file_path_parts = file_path.split("\\")
@@ -179,63 +285,22 @@ async def compute_sha256_checksum(filepath):
             sha256.update(chunk)
     return sha256.hexdigest()
 
-# def hash_chunk(start_end, filepath):
-#     """Hash a specific chunk of the file."""
-#     start, end = start_end
-#     sha256 = hashlib.sha256()
-#     with open(filepath, 'rb') as f:
-#         f.seek(start)
-#         chunk = f.read(end - start)
-#         sha256.update(chunk)
-#     return sha256.digest()  # Return the digest of the chunk
-
-# async def compute_sha256_checksum(filepath):
-#     file_size = os.path.getsize(filepath)
-#     parts = 1  # Or any other division based on file size or desired concurrency
-#     part_size = file_size // parts
-#     start_end_ranges = [(i * part_size, min((i + 1) * part_size, file_size)) for i in range(parts)]
-    
-#     print(start_end_ranges, file_size)
-
-#     loop = asyncio.get_running_loop()
-
-#     # Use ThreadPoolExecutor to process chunks in parallel
-#     with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
-#         futures = [loop.run_in_executor(executor, hash_chunk, start_end, filepath) for start_end in start_end_ranges]
-#         chunk_hashes = await asyncio.gather(*futures)
-
-#     # Combine the hashes sequentially
-#     final_sha256 = hashlib.sha256()
-#     for chunk_hash in chunk_hashes:
-#         final_sha256.update(chunk_hash)
-
-#     return final_sha256.hexdigest()
-
-# def hash_chunk(filepath):
-#     chunk_size = 1024 * 256  # 256KB per chunk
-#     sha256 = hashlib.sha256()
-#     with open(filepath, 'rb') as f:
-#         while True:
-#             chunk = f.read(chunk_size)
-#             if not chunk:
-#                 break  # End of file
-#             sha256.update(chunk)
-#     return sha256.hexdigest()
-
-# async def compute_sha256_checksum(filepath):
-#     print("computing sha256 checksum")
-#     filepath = get_comfyui_path_from_file_path(filepath)
-    
-#     loop = asyncio.get_running_loop()
-    
-#     with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
-#         task = loop.run_in_executor(executor, hash_chunk, filepath)
-    
-#     return await task
+@server.PromptServer.instance.routes.get('/comfyui-deploy/models')
+async def get_installed_models(request):
+    # Directly return the list of paths as JSON
+    new_dict = {}
+    for key, value in folder_paths.folder_names_and_paths.items():
+        # Convert set to list for JSON compatibility
+        # for path in value[0]:
+        file_list = folder_paths.get_filename_list(key)
+        value_json_compatible = (value[0], list(value[1]), file_list)
+        new_dict[key] = value_json_compatible
+    # print(new_dict)
+    return web.json_response(new_dict)
 
 # This is start uploading the files to Comfy Deploy
 @server.PromptServer.instance.routes.post('/comfyui-deploy/upload-file')
-async def upload_file(request):
+async def upload_file_endpoint(request):
     data = await request.json()
 
     file_path = data.get("file_path")
@@ -344,6 +409,16 @@ async def get_file_hash(request):
         return web.json_response({
             "error": str(e)
         }, status=500)
+        
+async def update_realtime_run_status(realtime_id: str, status_endpoint: str, status: Status):
+    body = {
+        "run_id": realtime_id,
+        "status": status.value,
+    }
+    # requests.post(status_endpoint, json=body)
+    async with aiohttp.ClientSession() as session:
+        async with session.post(status_endpoint, json=body) as response:
+            pass
 
 @server.PromptServer.instance.routes.get('/comfyui-deploy/ws')
 async def websocket_handler(request):
@@ -357,36 +432,135 @@ async def websocket_handler(request):
         sid = uuid.uuid4().hex
 
     sockets[sid] = ws
+    
+    auth_token = request.rel_url.query.get('token', None)
+    get_workflow_endpoint_url = request.rel_url.query.get('workflow_endpoint', None)
+    realtime_id = request.rel_url.query.get('realtime_id', None)
+    status_endpoint = request.rel_url.query.get('status_endpoint', None)
+    
+    if auth_token is not None and get_workflow_endpoint_url is not None:
+        async with aiohttp.ClientSession() as session:
+            headers = {'Authorization': f'Bearer {auth_token}'}
+            async with session.get(get_workflow_endpoint_url, headers=headers) as response:
+                if response.status == 200:
+                    workflow = await response.json()
+                   
+                    print("Loaded workflow version ",workflow["version"])
+                    
+                    streaming_prompt_metadata[sid] = StreamingPrompt(
+                        workflow_api=workflow["workflow_api"], 
+                        auth_token=auth_token,
+                        inputs={},
+                        status_endpoint=status_endpoint,
+                        file_upload_endpoint=request.rel_url.query.get('file_upload_endpoint', None),
+                    )
+                    
+                    await update_realtime_run_status(realtime_id, status_endpoint, Status.RUNNING)
+                    # await send("workflow_api", workflow_api, sid)
+                else:
+                    error_message = await response.text()
+                    print(f"Failed to fetch workflow endpoint. Status: {response.status}, Error: {error_message}")
+                    # await send("error", {"message": error_message}, sid)
 
     try:
         # Send initial state to the new client
         await send("status", { 'sid': sid }, sid)
 
-        if cd_enable_log:
+        # Make sure when its connected via client, the full log is not being sent
+        if cd_enable_log and get_workflow_endpoint_url is None:
             await send_first_time_log(sid)
             
         async for msg in ws:
+            if msg.type == aiohttp.WSMsgType.TEXT:
+                try:
+                    data = json.loads(msg.data)
+                    print(data)
+                    event_type = data.get('event')
+                    if event_type == 'input':
+                        print("Got input: ", data.get("inputs"))
+                        input = data.get('inputs')
+                        streaming_prompt_metadata[sid].inputs.update(input)
+                    elif event_type == 'queue_prompt':
+                        clear_current_prompt(sid)
+                        send_prompt(sid, streaming_prompt_metadata[sid])
+                    else:
+                        # Handle other event types
+                        pass
+                except json.JSONDecodeError:
+                    print('Failed to decode JSON from message')
+            
+            if msg.type == aiohttp.WSMsgType.BINARY:
+                data = msg.data
+                event_type, = struct.unpack("<I", data[:4])
+                if event_type == 0:  # Image input
+                    image_type_code, = struct.unpack("<I", data[4:8])
+                    input_id_bytes = data[8:32]  # Extract the next 24 bytes for the input ID
+                    input_id = input_id_bytes.decode('ascii').strip()  # Decode the input ID from ASCII
+                    print(event_type)
+                    print(image_type_code)
+                    print(input_id)
+                    image_data = data[32:]  # The rest is the image data
+                    if image_type_code == 1:
+                        image_type = "JPEG"
+                    elif image_type_code == 2:
+                        image_type = "PNG"
+                    elif image_type_code == 3:
+                        image_type = "WEBP"
+                    else:
+                        print("Unknown image type code:", image_type_code)
+                        return
+                    image = Image.open(BytesIO(image_data))
+                    # Check if the input ID already exists and replace the input with the new one
+                    if input_id in streaming_prompt_metadata[sid].inputs:
+                        # If the input exists, we assume it's an image and attempt to close it to free resources
+                        try:
+                            existing_image = streaming_prompt_metadata[sid].inputs[input_id]
+                            if hasattr(existing_image, 'close'):
+                                existing_image.close()
+                        except Exception as e:
+                            print(f"Error closing previous image for input ID {input_id}: {e}")
+                    streaming_prompt_metadata[sid].inputs[input_id] = image
+                    # clear_current_prompt(sid)
+                    # send_prompt(sid, streaming_prompt_metadata[sid])
+                    print(f"Received {image_type} image of size {image.size} with input ID {input_id}")
+            
             if msg.type == aiohttp.WSMsgType.ERROR:
                 print('ws connection closed with exception %s' % ws.exception())
     finally:
         sockets.pop(sid, None)
+        
+        if realtime_id is not None:
+            await update_realtime_run_status(realtime_id, status_endpoint, Status.SUCCESS)
     return ws
 
 @server.PromptServer.instance.routes.get('/comfyui-deploy/check-status')
 async def comfy_deploy_check_status(request):
-    prompt_server = server.PromptServer.instance
     prompt_id = request.rel_url.query.get('prompt_id', None)
-    if prompt_id in prompt_metadata and 'status' in prompt_metadata[prompt_id]:
+    if prompt_id in prompt_metadata:
         return web.json_response({
-            "status": prompt_metadata[prompt_id]['status'].value
+            "status": prompt_metadata[prompt_id].status.value
         })
     else:
         return web.json_response({
             "message": "prompt_id not found"
         })
+        
+@server.PromptServer.instance.routes.get('/comfyui-deploy/check-ws-status')
+async def comfy_deploy_check_ws_status(request):
+    client_id = request.rel_url.query.get('client_id', None)
+    if client_id in streaming_prompt_metadata:
+        remaining_queue = 0  # Initialize remaining queue count
+        for prompt_id in streaming_prompt_metadata[client_id].running_prompt_ids:
+            prompt_status = prompt_metadata[prompt_id].status
+            if prompt_status not in [Status.FAILED, Status.SUCCESS]:
+                remaining_queue += 1  # Increment for each prompt still running
+        return web.json_response({"remaining_queue": remaining_queue})
+    else:
+        return web.json_response({"message": "client_id not found"}, status=404)
 
 async def send(event, data, sid=None):
     try:
+        # message = {"event": event, "data": data}
         if sid:
             ws = sockets.get(sid)
             if ws != None and not ws.closed:  # Check if the WebSocket connection is open and not closing
@@ -402,53 +576,74 @@ async def send(event, data, sid=None):
 logging.basicConfig(level=logging.INFO)
 
 prompt_server = server.PromptServer.instance
-
 send_json = prompt_server.send_json
+
 async def send_json_override(self, event, data, sid=None):
     # print("INTERNAL:", event, data, sid)
     prompt_id = data.get('prompt_id')
+    
+    target_sid = sid
+    if target_sid == "comfy_deploy_instance":
+        target_sid = None
 
     # now we send everything
     await asyncio.wait([
-        asyncio.create_task(send(event, data)),
+        asyncio.create_task(send(event, data, sid=target_sid)),
         asyncio.create_task(self.send_json_original(event, data, sid))
     ])
 
     if event == 'execution_start':
         update_run(prompt_id, Status.RUNNING)
+        
+        if prompt_id in prompt_metadata:
+            prompt_metadata[prompt_id].start_time = time.perf_counter()
 
     # the last executing event is none, then the workflow is finished
     if event == 'executing' and data.get('node') is None:
         mark_prompt_done(prompt_id=prompt_id)
         if not have_pending_upload(prompt_id):
             update_run(prompt_id, Status.SUCCESS)
+            if prompt_id in prompt_metadata:
+                current_time = time.perf_counter()
+                if prompt_metadata[prompt_id].start_time is not None:
+                    elapsed_time = current_time - prompt_metadata[prompt_id].start_time
+                    print(f"Elapsed time: {elapsed_time} seconds")
+                    await send("elapsed_time", {
+                        "prompt_id": prompt_id,
+                        "elapsed_time": elapsed_time
+                    }, sid=sid)
 
     if event == 'executing' and data.get('node') is not None:
         node = data.get('node')
         
         if prompt_id in prompt_metadata:
-            if 'progress' not in prompt_metadata[prompt_id]:
-                prompt_metadata[prompt_id]["progress"] = set()
-                
-            prompt_metadata[prompt_id]["progress"].add(node)
-            calculated_progress = len(prompt_metadata[prompt_id]["progress"]) / len(prompt_metadata[prompt_id]['workflow_api'])
+            # if 'progress' not in prompt_metadata[prompt_id]:
+            #     prompt_metadata[prompt_id]["progress"] = set()
+              
+            prompt_metadata[prompt_id].progress.add(node)
+            calculated_progress = len(prompt_metadata[prompt_id].progress) / len(prompt_metadata[prompt_id].workflow_api)
             # print("calculated_progress", calculated_progress)
-            
-            if 'last_updated_node' in prompt_metadata[prompt_id] and prompt_metadata[prompt_id]['last_updated_node'] == node:
+         
+            if prompt_metadata[prompt_id].last_updated_node is not None and prompt_metadata[prompt_id].last_updated_node == node:
                 return
-            prompt_metadata[prompt_id]['last_updated_node'] = node
-            class_type = prompt_metadata[prompt_id]['workflow_api'][node]['class_type']
+            prompt_metadata[prompt_id].last_updated_node = node
+            class_type = prompt_metadata[prompt_id].workflow_api[node]['class_type']
             print("updating run live status", class_type)
+            await send("live_status", {
+                "prompt_id": prompt_id,
+                "current_node": class_type,
+                "progress": calculated_progress,
+            }, sid=sid)
             await update_run_live_status(prompt_id, "Executing " + class_type, calculated_progress)
 
     if event == 'execution_cached' and data.get('nodes') is not None:
         if prompt_id in prompt_metadata:
-            if 'progress' not in prompt_metadata[prompt_id]:
-                prompt_metadata[prompt_id]["progress"] = set()
-            
+            # if 'progress' not in prompt_metadata[prompt_id]:
+            #     prompt_metadata[prompt_id].progress = set()
+           
             if 'nodes' in data:
                 for node in data.get('nodes', []):
-                    prompt_metadata[prompt_id]["progress"].add(node)
+                    prompt_metadata[prompt_id].progress.add(node)
             # prompt_metadata[prompt_id]["progress"].update(data.get('nodes'))
 
     if event == 'execution_error':
@@ -462,14 +657,6 @@ async def send_json_override(self, event, data, sid=None):
         # await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
         # update_run_with_output(prompt_id, data.get('output'))
 
-
-class Status(Enum):
-    NOT_STARTED = "not-started"
-    RUNNING = "running"
-    SUCCESS = "success"
-    FAILED = "failed"
-    UPLOADING = "uploading"
-
 # Global variable to keep track of the last read line number
 last_read_line_number = 0
 
@@ -477,9 +664,12 @@ async def update_run_live_status(prompt_id, live_status, calculated_progress: fl
     if prompt_id not in prompt_metadata:
         return
     
+    if prompt_metadata[prompt_id].is_realtime is True:
+        return
+    
     print("progress", calculated_progress)
     
-    status_endpoint = prompt_metadata[prompt_id]['status_endpoint']
+    status_endpoint = prompt_metadata[prompt_id].status_endpoint
     body = {
         "run_id": prompt_id,
         "live_status": live_status,
@@ -491,19 +681,25 @@ async def update_run_live_status(prompt_id, live_status, calculated_progress: fl
             pass
 
 
-def update_run(prompt_id, status: Status):
+def update_run(prompt_id: str, status: Status):
     global last_read_line_number
 
     if prompt_id not in prompt_metadata:
         return
+    
+    # if prompt_metadata[prompt_id].start_time is None and status == Status.RUNNING:
+    # if its realtime prompt we need to skip that.
+    if prompt_metadata[prompt_id].is_realtime is True:
+        prompt_metadata[prompt_id].status = status
+        return
 
-    if ('status' not in prompt_metadata[prompt_id] or prompt_metadata[prompt_id]['status'] != status):
+    if (prompt_metadata[prompt_id].status != status):
 
         # when the status is already failed, we don't want to update it to success
-        if ('status' in prompt_metadata[prompt_id] and prompt_metadata[prompt_id]['status'] == Status.FAILED):
+        if ('status' in prompt_metadata[prompt_id] and prompt_metadata[prompt_id].status == Status.FAILED):
             return
 
-        status_endpoint = prompt_metadata[prompt_id]['status_endpoint']
+        status_endpoint = prompt_metadata[prompt_id].status_endpoint
         body = {
             "run_id": prompt_id,
             "status": status.value,
@@ -544,13 +740,12 @@ def update_run(prompt_id, status: Status):
                 except Exception as log_error:
                     print(f"Error reading log file: {log_error}")
                 
-
         except Exception as e:
             error_type = type(e).__name__
             stack_trace = traceback.format_exc().strip()
             print(f"Error occurred while updating run: {e} {stack_trace}")
         finally:
-            prompt_metadata[prompt_id]['status'] = status
+            prompt_metadata[prompt_id].status = status
             
 
 async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output"):
@@ -582,7 +777,7 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
 
     print("uploading file", file)
 
-    file_upload_endpoint = prompt_metadata[prompt_id]['file_upload_endpoint']
+    file_upload_endpoint = prompt_metadata[prompt_id].file_upload_endpoint
 
     filename = quote(filename)
     prompt_id = quote(prompt_id)
@@ -606,23 +801,37 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
                 print("upload file response", response.status)
 
 def have_pending_upload(prompt_id):
-    if 'prompt_id' in prompt_metadata and 'uploading_nodes' in prompt_metadata[prompt_id] and len(prompt_metadata[prompt_id]['uploading_nodes']) > 0:
-        print("have pending upload ", len(prompt_metadata[prompt_id]['uploading_nodes']))
+    if prompt_id in prompt_metadata and len(prompt_metadata[prompt_id].uploading_nodes) > 0:
+        print("have pending upload ", len(prompt_metadata[prompt_id].uploading_nodes))
         return True
 
     print("no pending upload")
     return False
 
 def mark_prompt_done(prompt_id):
+    """
+    Mark the prompt as done in the prompt metadata.
+
+    Args:
+        prompt_id (str): The ID of the prompt to mark as done.
+    """
     if prompt_id in prompt_metadata:
-        prompt_metadata[prompt_id]["done"] = True
+        prompt_metadata[prompt_id].done = True
         print("Prompt done")
 
-def is_prompt_done(prompt_id):
-    if prompt_id in prompt_metadata and "done" in prompt_metadata[prompt_id]:
-        if prompt_metadata[prompt_id]["done"] == True:
-            return True
-    
+def is_prompt_done(prompt_id: str):
+    """
+    Check if the prompt with the given ID is marked as done.
+
+    Args:
+        prompt_id (str): The ID of the prompt to check.
+
+    Returns:
+        bool: True if the prompt is marked as done, False otherwise.
+    """
+    if prompt_id in prompt_metadata and prompt_metadata[prompt_id].done is True:
+        return True
+
     return False
 
 # Use to handle upload error and send back to ComfyDeploy
@@ -644,17 +853,17 @@ async def handle_error(prompt_id, data, e: Exception):
     print(f"Error occurred while uploading file: {e}")
 
 # Mark the current prompt requires upload, and block it from being marked as success
-async def update_file_status(prompt_id, data, uploading, have_error=False, node_id=None):
-    if 'uploading_nodes' not in prompt_metadata[prompt_id]:
-        prompt_metadata[prompt_id]['uploading_nodes'] = set()
+async def update_file_status(prompt_id: str, data, uploading, have_error=False, node_id=None):
+    # if 'uploading_nodes' not in prompt_metadata[prompt_id]:
+    #     prompt_metadata[prompt_id]['uploading_nodes'] = set()
 
     if node_id is not None:
         if uploading:
-            prompt_metadata[prompt_id]['uploading_nodes'].add(node_id)
+            prompt_metadata[prompt_id].uploading_nodes.add(node_id)
         else:
-            prompt_metadata[prompt_id]['uploading_nodes'].discard(node_id)
+            prompt_metadata[prompt_id].uploading_nodes.discard(node_id)
 
-    print(prompt_metadata[prompt_id]['uploading_nodes'])
+    print(prompt_metadata[prompt_id].uploading_nodes)
     # Update the remote status
 
     if have_error:
@@ -666,12 +875,12 @@ async def update_file_status(prompt_id, data, uploading, have_error=False, node_
 
     # if there are still nodes that are uploading, then we set the status to uploading
     if uploading:
-        if prompt_metadata[prompt_id]['status'] != Status.UPLOADING:
+        if prompt_metadata[prompt_id].status != Status.UPLOADING:
             update_run(prompt_id, Status.UPLOADING)
             await send("uploading", {
                 "prompt_id": prompt_id,
             })
-    
+   
     # if there are no nodes that are uploading, then we set the status to success
     elif not uploading and not have_pending_upload(prompt_id) and is_prompt_done(prompt_id=prompt_id):
         update_run(prompt_id, Status.SUCCESS)
@@ -680,20 +889,20 @@ async def update_file_status(prompt_id, data, uploading, have_error=False, node_
             "prompt_id": prompt_id,
         })
 
-async def handle_upload(prompt_id, data, key, content_type_key, default_content_type):
+async def handle_upload(prompt_id: str, data, key: str, content_type_key: str, default_content_type: str):
     items = data.get(key, [])
     for item in items:
         await upload_file(
-            prompt_id, 
-            item.get("filename"), 
-            subfolder=item.get("subfolder"), 
-            type=item.get("type"), 
+            prompt_id,
+            item.get("filename"),
+            subfolder=item.get("subfolder"),
+            type=item.get("type"),
             content_type=item.get(content_type_key, default_content_type)
         )
 
 
 # Upload files in the background
-async def upload_in_background(prompt_id, data, node_id=None, have_upload=True):
+async def upload_in_background(prompt_id: str, data, node_id=None, have_upload=True):
     try:
         await handle_upload(prompt_id, data, 'images', "content_type", "image/png")
         await handle_upload(prompt_id, data, 'files', "content_type", "image/png")
@@ -706,32 +915,37 @@ async def upload_in_background(prompt_id, data, node_id=None, have_upload=True):
         await handle_error(prompt_id, data, e)
 
 async def update_run_with_output(prompt_id, data, node_id=None):
-    if prompt_id in prompt_metadata:
-        status_endpoint = prompt_metadata[prompt_id]['status_endpoint']
+    if prompt_id not in prompt_metadata:
+        return
+        
+    if prompt_metadata[prompt_id].is_realtime is True:
+        return
+    
+    status_endpoint = prompt_metadata[prompt_id].status_endpoint
 
-        body = {
-            "run_id": prompt_id,
-            "output_data": data
-        }
+    body = {
+        "run_id": prompt_id,
+        "output_data": data
+    }
 
-        try:
-            have_upload = 'images' in data or 'files' in data or 'gifs' in data
-            print("\nhave_upload", have_upload, node_id)
+    try:
+        have_upload = 'images' in data or 'files' in data or 'gifs' in data
+        print("\nhave_upload", have_upload, node_id)
 
-            if have_upload:
-                await update_file_status(prompt_id, data, True, node_id=node_id)
+        if have_upload:
+            await update_file_status(prompt_id, data, True, node_id=node_id)
 
-            asyncio.create_task(upload_in_background(prompt_id, data, node_id=node_id, have_upload=have_upload))
+        asyncio.create_task(upload_in_background(prompt_id, data, node_id=node_id, have_upload=have_upload))
 
-        except Exception as e:
-            await handle_error(prompt_id, data, e)
-            
+    except Exception as e:
+        await handle_error(prompt_id, data, e)
+        
 
-        requests.post(status_endpoint, json=body)
+    requests.post(status_endpoint, json=body)
 
-        await send('outputs_uploaded', {
-            "prompt_id": prompt_id
-        })
+    await send('outputs_uploaded', {
+        "prompt_id": prompt_id
+    })
 
 prompt_server.send_json_original = prompt_server.send_json
 prompt_server.send_json = send_json_override.__get__(prompt_server, server.PromptServer)

docker-compose.yaml

@@ -0,0 +1,55 @@
+version: '3.9'
+services:
+
+  comfy-deploy:
+    build:
+      context: .
+      dockerfile: ./local/Dockerfile
+    restart: unless-stopped
+    volumes:
+      - ./local/scripts/entrypoint.sh:/comfyui-deploy/web/deploy_entrypoint.sh
+    entrypoint: /comfyui-deploy/web/deploy_entrypoint.sh
+    ports:
+      - 3000:3000
+    depends_on:
+      - postgres
+      - pg_proxy
+      - localstack
+    environment:
+      VSCODE_DEV_CONTAINER: true
+
+
+### comfy-deploy services
+
+  postgres:
+      image: "postgres:15.2-alpine"
+      environment:
+        POSTGRES_USER: postgres
+        POSTGRES_PASSWORD: postgres
+        POSTGRES_DB: verceldb
+        POSTGRES_PORT: 5480
+      expose:
+        - 5480
+
+
+  pg_proxy:
+    image: ghcr.io/neondatabase/wsproxy:latest
+    environment:
+      APPEND_PORT: "postgres:5480"
+      ALLOW_ADDR_REGEX: ".*"
+      LOG_TRAFFIC: "true"
+    expose:
+      - 80
+    depends_on:
+      - postgres
+
+
+  localstack:
+    image: localstack/localstack:latest
+    environment:
+      SERVICES: s3
+    ports:
+      - 4566:4566
+    volumes:
+      - ../localstack/aws:/etc/localstack/init/ready.d
+      - ../localstack/aws:/app/web/aws

globals.py

@@ -0,0 +1,97 @@
+import struct
+
+import aiohttp
+
+from typing import List, Union, Any, Optional
+from PIL import Image, ImageOps
+from io import BytesIO
+
+from pydantic import BaseModel as PydanticBaseModel
+
+class BaseModel(PydanticBaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+class StreamingPrompt(BaseModel):
+    workflow_api: Any
+    auth_token: str
+    inputs: dict[str, Union[str, bytes, Image.Image]]
+    running_prompt_ids: set[str] = set()
+    status_endpoint: str
+    file_upload_endpoint: str
+
+sockets = dict()
+streaming_prompt_metadata: dict[str, StreamingPrompt] = {}
+
+class BinaryEventTypes:
+    PREVIEW_IMAGE = 1
+    UNENCODED_PREVIEW_IMAGE = 2
+    
+max_output_id_length = 24
+
+async def send_image(image_data, sid=None, output_id:str = None):
+    max_length = max_output_id_length
+    output_id = output_id[:max_length]
+    padded_output_id = output_id.ljust(max_length, '\x00')
+    encoded_output_id = padded_output_id.encode('ascii', 'replace')
+    
+    image_type = image_data[0]
+    image = image_data[1]
+    max_size = image_data[2]
+    quality = image_data[3]
+    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
+    elif image_type == "WEBP":
+        type_num = 3
+
+    bytesIO = BytesIO()
+    header = struct.pack(">I", type_num)
+    # 4 bytes for the type
+    bytesIO.write(header)
+    # 10 bytes for the output_id
+    position_before = bytesIO.tell()
+    bytesIO.write(encoded_output_id)
+    position_after = bytesIO.tell()
+    bytes_written = position_after - position_before
+    print(f"Bytes written: {bytes_written}")
+    
+    image.save(bytesIO, format=image_type, quality=quality, compress_level=1)
+    preview_bytes = bytesIO.getvalue()
+    await send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)
+        
+async def send_socket_catch_exception(function, message):
+    try:
+        await function(message)
+    except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError) as err:
+        print("send error:", err)
+
+def encode_bytes(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_bytes(event, data, sid=None):
+    message = encode_bytes(event, data)
+    
+    print("sending image to ", event, sid)
+
+    if sid is None:
+        _sockets = list(sockets.values())
+        for ws in _sockets:
+            await send_socket_catch_exception(ws.send_bytes, message)
+    elif sid in sockets:
+        await send_socket_catch_exception(sockets[sid].send_bytes, message)

local/Dockerfile

@@ -0,0 +1,18 @@
+FROM node:21-bullseye AS comfy_deploy
+
+RUN apt-get update && apt-get install -y python3 make g++
+
+RUN npm install -g bun
+
+COPY ./web /web
+
+
+WORKDIR /web
+
+RUN cp .env.example .env.local
+
+RUN bunx node-gyp
+
+RUN bun i
+
+ENTRYPOINT [ "bun", "dev" ]

local/scripts/entrypoint.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+echo "comfy deploy container starting.."
+
+echo "Running migrations.."
+bun migrate-local
+
+echo "Starting comfy deploy.."
+bun dev

requirements.txt

@@ -1 +1,2 @@
-aiofiles
+aiofiles
+pydantic

web-plugin/index.js

@@ -1,7 +1,7 @@
 import { app } from "./app.js";
 import { api } from "./api.js";
 import { ComfyWidgets, LGraphNode } from "./widgets.js";
-import { generateDependencyGraph } from "https://esm.sh/[email protected].19";
+import { generateDependencyGraph } from "https://esm.sh/[email protected].22";
 
 const loadingIcon = `<svg xmlns="http://www.w3.org/2000/svg" width="32" height="32" viewBox="0 0 24 24"><g fill="none" stroke="#888888" stroke-linecap="round" stroke-width="2"><path stroke-dasharray="60" stroke-dashoffset="60" stroke-opacity=".3" d="M12 3C16.9706 3 21 7.02944 21 12C21 16.9706 16.9706 21 12 21C7.02944 21 3 16.9706 3 12C3 7.02944 7.02944 3 12 3Z"><animate fill="freeze" attributeName="stroke-dashoffset" dur="1.3s" values="60;0"/></path><path stroke-dasharray="15" stroke-dashoffset="15" d="M12 3C16.9706 3 21 7.02944 21 12"><animate fill="freeze" attributeName="stroke-dashoffset" dur="0.3s" values="15;0"/><animateTransform attributeName="transform" dur="1.5s" repeatCount="indefinite" type="rotate" values="0 12 12;360 12 12"/></path></g></svg>`;
 
@@ -59,8 +59,8 @@ const ext = {
             return;
           }
 
-          // Adding a delay to wait for the intial graph to load
-          await new Promise((resolve) => setTimeout(resolve, 2000));
+          // // Adding a delay to wait for the intial graph to load
+          // await new Promise((resolve) => setTimeout(resolve, 2000));
 
           workflow?.nodes.forEach((x) => {
             if (x?.type === "ComfyDeploy") {
@@ -285,12 +285,37 @@ function addButton() {
       return;
     }
 
+    let deployMeta = graph.findNodesByType("ComfyDeploy");
+
+    if (deployMeta.length == 0) {
+      const text = await inputDialog.input(
+        "Create your deployment",
+        "Workflow name",
+      );
+      if (!text) return;
+      console.log(text);
+      app.graph.beforeChange();
+      var node = LiteGraph.createNode("ComfyDeploy");
+      node.configure({
+        widgets_values: [text],
+      });
+      node.pos = [0, 0];
+      app.graph.add(node);
+      app.graph.afterChange();
+      deployMeta = [node];
+    }
+
+    const deployMetaNode = deployMeta[0];
+
+    const workflow_name = deployMetaNode.widgets[0].value;
+    const workflow_id = deployMetaNode.widgets[1].value;
+
     const ok = await confirmDialog.confirm(
       `Confirm deployment`,
       `
       <div>
 
-      A new version will be deployed, do you confirm? 
+      A new version of <button style="font-size: 18px;">${workflow_name}</button> will be deployed, do you confirm? 
       <br><br>
 
       <button style="font-size: 18px;">${displayName}</button>
@@ -332,31 +357,6 @@ function addButton() {
 
     const title = deploy.querySelector("#button-title");
 
-    let deployMeta = graph.findNodesByType("ComfyDeploy");
-
-    if (deployMeta.length == 0) {
-      const text = await inputDialog.input(
-        "Create your deployment",
-        "Workflow name",
-      );
-      if (!text) return;
-      console.log(text);
-      app.graph.beforeChange();
-      var node = LiteGraph.createNode("ComfyDeploy");
-      node.configure({
-        widgets_values: [text],
-      });
-      node.pos = [0, 0];
-      app.graph.add(node);
-      app.graph.afterChange();
-      deployMeta = [node];
-    }
-
-    const deployMetaNode = deployMeta[0];
-
-    const workflow_name = deployMetaNode.widgets[0].value;
-    const workflow_id = deployMetaNode.widgets[1].value;
-
     const prompt = await app.graphToPrompt();
     let deps = undefined;
 
@@ -474,7 +474,7 @@ function addButton() {
         <div style="position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%);">${loadingIcon}</div>
         <iframe 
         style="z-index: 10; min-width: 600px; max-width: 1024px; min-height: 600px; border: none; background-color: transparent;"
-        src="${endpoint}/dependency-graph?deps=${encodeURIComponent(
+        src="https://www.comfydeploy.com/dependency-graph?deps=${encodeURIComponent(
           JSON.stringify(deps),
         )}" />`,
         // createDynamicUIHtml(deps),
@@ -880,7 +880,10 @@ export class ConfigDialog extends ComfyDialog {
             justifyContent: "flex-end",
             width: "100%",
           },
-          onclick: () => this.save(),
+          onclick: () => {
+            this.save();
+            this.close();
+          },
         },
         [
           $el("button", {
@@ -891,7 +894,10 @@ export class ConfigDialog extends ComfyDialog {
           $el("button", {
             type: "button",
             textContent: "Save",
-            onclick: () => this.save(),
+            onclick: () => {
+              this.save();
+              this.close();
+            },
           }),
         ],
       ),
@@ -905,20 +911,26 @@ export class ConfigDialog extends ComfyDialog {
   }
 
   save(api_key, displayName) {
-    if (!displayName) displayName = getData().displayName;
-
     const deployOption = this.container.querySelector("#deployOption").value;
     localStorage.setItem("comfy_deploy_env", deployOption);
 
     const endpoint = this.container.querySelector("#endpoint").value;
     const apiKey = api_key ?? this.container.querySelector("#apiKey").value;
+
+    if (!displayName) {
+      if (apiKey != getData().apiKey) {
+        displayName = "Custom";
+      } else {
+        displayName = getData().displayName;
+      }
+    }
+
     saveData({
       endpoint,
       apiKey,
       displayName,
       environment: deployOption,
     });
-    this.close();
   }
 
   show() {
@@ -941,8 +953,10 @@ export class ConfigDialog extends ComfyDialog {
           data.endpoint
         }">
       </label>
-      <label style="color: white;">
-        API Key: ${data.displayName ?? ""}
+      <div style="color: white;">
+        API Key: User / Org <button style="font-size: 18px;">${
+          data.displayName ?? ""
+        }</button>
         <input id="apiKey" style="margin-top: 8px; width: 100%; height:40px; box-sizing: border-box; padding: 0px 6px;" type="password" value="${
           data.apiKey
         }">
@@ -951,12 +965,15 @@ export class ConfigDialog extends ComfyDialog {
             data.apiKey ? "Re-login with ComfyDeploy" : "Login with ComfyDeploy"
           }
         </button>
-      </label>
+      </div>
       </div>
     `;
 
     const button = this.container.querySelector("#loginButton");
     button.onclick = () => {
+      this.save();
+      const data = getData();
+
       const uuid =
         Math.random().toString(36).substring(2, 15) +
         Math.random().toString(36).substring(2, 15);
@@ -973,17 +990,20 @@ export class ConfigDialog extends ComfyDialog {
       this.poll = setInterval(() => {
         fetch(data.endpoint + "/api/auth-response/" + uuid)
           .then((response) => response.json())
-          .then((json) => {
+          .then(async (json) => {
             if (json.api_key) {
               this.save(json.api_key, json.name);
+              this.close();
               this.container.querySelector("#apiKey").value = json.api_key;
-              infoDialog.show();
+              // infoDialog.show();
               clearInterval(this.poll);
               clearTimeout(this.timeout);
-              infoDialog.showMessage(
+              // Refresh dialog
+              const a = await confirmDialog.confirm(
                 "Authenticated",
-                "You will be able to upload workflow to " + json.name,
+                `<div>You will be able to upload workflow to <button style="font-size: 18px; width: fit;">${json.name}</button></div>`,
               );
+              configDialog.show();
             }
           })
           .catch((error) => {

web/src/db/db.ts

@@ -9,10 +9,10 @@ if (process.env.VERCEL_ENV !== "production") {
   // Set the WebSocket proxy to work with the local instance
   if (isDevContainer) {
     // Running inside a VS Code devcontainer
-    neonConfig.wsProxy = (host) => "host.docker.internal:5481/v1";
+    neonConfig.wsProxy = (host) => "pg_proxy:80/v1";
   } else {
     // Not running inside a VS Code devcontainer
-    neonConfig.wsProxy = (host) => `${host}:5481/v1`;
+    neonConfig.wsProxy = (host) => "pg_proxy:80/v1";
   }
   // Disable all authentication and encryption
   neonConfig.useSecureWebSocket = false;