better

custom_routes.py

@@ -33,7 +33,7 @@ client_session = None
 #     global client_session
 #     if client_session is None:
 #         client_session = aiohttp.ClientSession()
-        
+
 async def ensure_client_session():
     global client_session
     if client_session is None:
@@ -43,7 +43,7 @@ async def cleanup():
     global client_session
     if client_session:
         await client_session.close()
-        
+
 def exit_handler():
     print("Exiting the application. Initiating cleanup...")
     loop = asyncio.get_event_loop()
@@ -56,20 +56,24 @@ retry_delay_multiplier = float(os.environ.get('RETRY_DELAY_MULTIPLIER', '2'))
 
 print(f"max_retries: {max_retries}, retry_delay_multiplier: {retry_delay_multiplier}")
 
-async def async_request_with_retry(method, url, **kwargs):
+async def async_request_with_retry(method, url, disable_timeout=False, **kwargs):
     global client_session
     await ensure_client_session()
+    # async with aiohttp.ClientSession() as client_session:
     retry_delay = 1  # Start with 1 second delay
     initial_timeout = 5  # 5 seconds timeout for the initial connection
 
     for attempt in range(max_retries):
         try:
             # Set a timeout for the initial connection
-            timeout = ClientTimeout(total=None, connect=initial_timeout)
+            if not disable_timeout:
-            kwargs['timeout'] = timeout
+                timeout = ClientTimeout(total=None, connect=initial_timeout)
+                kwargs['timeout'] = timeout
 
             async with client_session.request(method, url, **kwargs) as response:
                 response.raise_for_status()
+                if method.upper() == 'GET':
+                    await response.read()
                 return response
         except asyncio.TimeoutError:
             logger.warning(f"Request timed out after {initial_timeout} seconds (attempt {attempt + 1}/{max_retries})")
@@ -78,7 +82,7 @@ async def async_request_with_retry(method, url, **kwargs):
                 logger.error(f"Request failed after {max_retries} attempts: {e}")
                 # raise
             logger.warning(f"Request failed (attempt {attempt + 1}/{max_retries}): {e}")
-        
+
         # Wait before retrying
         await asyncio.sleep(retry_delay)
         retry_delay *= retry_delay_multiplier  # Exponential backoff
@@ -112,7 +116,7 @@ def log(level, message, **kwargs):
         getattr(logger, level)(message, **kwargs)
     else:
         getattr(logger, level)(f"{message} {kwargs}")
-        
+
 # For a span, you might need to create a context manager
 from contextlib import contextmanager
 
@@ -230,29 +234,32 @@ def apply_random_seed_to_workflow(workflow_api):
         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 'inputs' in workflow_api[key]:
-            if isinstance(workflow_api[key]['inputs']['seed'], list):
+            if 'seed' in workflow_api[key]['inputs']:
-                continue
+                if isinstance(workflow_api[key]['inputs']['seed'], list):
-            if workflow_api[key]['class_type'] == "PromptExpansion":
+                    continue
-                workflow_api[key]['inputs']['seed'] = randomSeed(8)
+                if workflow_api[key]['class_type'] == "PromptExpansion":
-                logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to PromptExpansion")
+                    workflow_api[key]['inputs']['seed'] = randomSeed(8)
-                continue
+                    logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to PromptExpansion")
-            if workflow_api[key]['class_type'] == "RandomNoise":
+                    continue
-                workflow_api[key]['inputs']['noise_seed'] = randomSeed()
+                workflow_api[key]['inputs']['seed'] = randomSeed()
-                logger.info(f"Applied random noise_seed {workflow_api[key]['inputs']['noise_seed']} to RandomNoise")
+                logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to {workflow_api[key]['class_type']}")
-                continue
-            if workflow_api[key]['class_type'] == "KSamplerAdvanced":
-                workflow_api[key]['inputs']['noise_seed'] = randomSeed()
-                logger.info(f"Applied random noise_seed {workflow_api[key]['inputs']['noise_seed']} to KSamplerAdvanced")
-                continue
-            if workflow_api[key]['class_type'] == "SamplerCustom":
-                workflow_api[key]['inputs']['noise_seed'] = randomSeed()
-                logger.info(f"Applied random noise_seed {workflow_api[key]['inputs']['noise_seed']} to SamplerCustom")
-                continue
-            workflow_api[key]['inputs']['seed'] = randomSeed()
-            logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to {workflow_api[key]['class_type']}")
 
-def apply_inputs_to_workflow(workflow_api: Any, inputs: Any, sid: str = None):
+            if 'noise_seed' in workflow_api[key]['inputs']:
+                if workflow_api[key]['class_type'] == "RandomNoise":
+                    workflow_api[key]['inputs']['noise_seed'] = randomSeed()
+                    logger.info(f"Applied random noise_seed {workflow_api[key]['inputs']['noise_seed']} to RandomNoise")
+                    continue
+                if workflow_api[key]['class_type'] == "KSamplerAdvanced":
+                    workflow_api[key]['inputs']['noise_seed'] = randomSeed()
+                    logger.info(f"Applied random noise_seed {workflow_api[key]['inputs']['noise_seed']} to KSamplerAdvanced")
+                    continue
+                if workflow_api[key]['class_type'] == "SamplerCustom":
+                    workflow_api[key]['inputs']['noise_seed'] = randomSeed()
+                    logger.info(f"Applied random noise_seed {workflow_api[key]['inputs']['noise_seed']} to SamplerCustom")
+                    continue
+
+def apply_inputs_to_workflow(workflow_api: Any, inputs: Any, sid: str | None = None):
     # Loop through each of the inputs and replace them
     for key, value in workflow_api.items():
         if 'inputs' in value:
@@ -848,19 +855,19 @@ async def send(event, data, sid=None):
     except Exception as e:
         logger.info(f"Exception: {e}")
         traceback.print_exc()
-        
+
 @server.PromptServer.instance.routes.get('/comfydeploy/{tail:.*}')
 @server.PromptServer.instance.routes.post('/comfydeploy/{tail:.*}')
 async def proxy_to_comfydeploy(request):
     # Get the base URL
     base_url = f'https://www.comfydeploy.com/{request.match_info["tail"]}'
-    
+
     # Get all query parameters
     query_params = request.query_string
-    
+
     # Construct the full target URL with query parameters
     target_url = f"{base_url}?{query_params}" if query_params else base_url
-    
+
     # print(f"Proxying request to: {target_url}")
 
     try:
@@ -994,7 +1001,7 @@ async def send_json_override(self, event, data, sid=None):
                 return
         else:
             logger.info(f"Executed {data}")
-            
+
         await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
         # await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
         # update_run_with_output(prompt_id, data.get('output'))
@@ -1150,18 +1157,18 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
     prompt_id = quote(prompt_id)
     content_type = quote(content_type)
 
-    target_url = f"{file_upload_endpoint}?file_name={filename}&run_id={prompt_id}&type={content_type}"
+    target_url = f"{file_upload_endpoint}?file_name={filename}&run_id={prompt_id}&type={content_type}&version=v2"
 
     start_time = time.time()  # Start timing here
-    result = requests.get(target_url)
+    result = await async_request_with_retry("GET", target_url, disable_timeout=True)
     end_time = time.time()  # End timing after the request is complete
     logger.info("Time taken for getting file upload endpoint: {:.2f} seconds".format(end_time - start_time))
-    ok = result.json()
+    ok = await result.json()
 
     start_time = time.time()  # Start timing here
 
-    with open(file, 'rb') as f:
+    async with aiofiles.open(file, 'rb') as f:
-        data = f.read()
+        data = await f.read()
         headers = {
             # "x-amz-acl": "public-read",
             "Content-Type": content_type,
@@ -1264,8 +1271,10 @@ async def update_file_status(prompt_id: str, data, uploading, have_error=False,
 
 async def handle_upload(prompt_id: str, data, key: str, content_type_key: str, default_content_type: str):
     items = data.get(key, [])
+    upload_tasks = []
+
     for item in items:
-        # # Skipping temp files
+        # Skipping temp files
         if item.get("type") == "temp":
             continue
 
@@ -1278,22 +1287,28 @@ async def handle_upload(prompt_id: str, data, key: str, content_type_key: str, d
         elif file_extension == '.webp':
             file_type = 'image/webp'
 
-        await upload_file(
+        upload_tasks.append(upload_file(
             prompt_id,
             item.get("filename"),
             subfolder=item.get("subfolder"),
             type=item.get("type"),
             content_type=file_type
-        )
+        ))
+
+    # Execute all upload tasks concurrently
+    await asyncio.gather(*upload_tasks)
 
 # Upload files in the background
 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")
+        upload_tasks = [
-        await handle_upload(prompt_id, data, 'files', "content_type", "image/png")
+            handle_upload(prompt_id, data, 'images', "content_type", "image/png"),
-        # This will also be mp4
+            handle_upload(prompt_id, data, 'files', "content_type", "image/png"),
-        await handle_upload(prompt_id, data, 'gifs', "format", "image/gif")
+            handle_upload(prompt_id, data, 'gifs', "format", "image/gif"),
-        await handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream")
+            handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream")
+        ]
+
+        await asyncio.gather(*upload_tasks)
 
         if have_upload:
             await update_file_status(prompt_id, data, False, node_id=node_id)
@@ -1392,4 +1407,4 @@ if cd_enable_log:
 @server.PromptServer.instance.routes.get("/comfyui-deploy/filename_list_cache")
 async def get_filename_list_cache(_):
     from folder_paths import filename_list_cache
-    return web.json_response({'filename_list': filename_list_cache})
+    return web.json_response({'filename_list': filename_list_cache})