Compare commits
19
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4df9d38e56 | ||
|
|
9cd626e1f6 | ||
|
|
503dca8fb6 | ||
|
|
73c149b4cb | ||
|
|
65b5b0b8c7 | ||
|
|
9d6ee85402 | ||
|
|
cdaed8a571 | ||
|
|
3129e89cce | ||
|
|
7a693eabc8 | ||
|
|
8f677e520d | ||
|
|
4c8d32c5b0 | ||
|
|
a99d2568e0 | ||
|
|
649b61c580 | ||
|
|
edff5685f9 | ||
|
|
9fc0c2b4a2 | ||
|
|
d34e2e99b1 | ||
|
|
f85043db07 | ||
|
|
894d8e1503 | ||
|
|
08d631d1eb |
@@ -12,11 +12,11 @@ class ComfyUIDeployExternalBoolean:
|
||||
"optional": {
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,11 +25,11 @@ class ComfyUIDeployExternalCheckpoint:
|
||||
"default_value": (folder_paths.get_filename_list("checkpoints"), ),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,11 +17,11 @@ class ComfyUIDeployExternalImage:
|
||||
"default_value": ("IMAGE",),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,11 +17,11 @@ class ComfyUIDeployExternalImageAlpha:
|
||||
"default_value": ("IMAGE",),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,11 +23,11 @@ class ComfyUIDeployExternalImageBatch:
|
||||
"default_value": ("IMAGE",),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,17 +25,21 @@ class ComfyUIDeployExternalLora:
|
||||
},
|
||||
"optional": {
|
||||
"default_lora_name": (folder_paths.get_filename_list("loras"),),
|
||||
"lora_save_name": ( # if `default_lora_name` is a link to download a file, we will attempt to save it with this name
|
||||
"lora_save_name": ( # if `default_lora_name` is a link to download a file, we will attempt to save it with this name
|
||||
"STRING",
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
"lora_url": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
},
|
||||
}
|
||||
@@ -47,12 +51,20 @@ class ComfyUIDeployExternalLora:
|
||||
|
||||
CATEGORY = "deploy"
|
||||
|
||||
def run(self, input_id, default_lora_name=None, lora_save_name=None, display_name=None, description=None):
|
||||
def run(
|
||||
self,
|
||||
input_id,
|
||||
default_lora_name=None,
|
||||
lora_save_name=None,
|
||||
display_name=None,
|
||||
description=None,
|
||||
lora_url=None,
|
||||
):
|
||||
import requests
|
||||
import os
|
||||
import uuid
|
||||
|
||||
if default_lora_name.startswith("http"):
|
||||
if lora_url and lora_url.startswith("http"):
|
||||
if lora_save_name:
|
||||
existing_loras = folder_paths.get_filename_list("loras")
|
||||
# Check if lora_save_name exists in the list
|
||||
@@ -67,9 +79,9 @@ class ComfyUIDeployExternalLora:
|
||||
folder_paths.folder_names_and_paths["loras"][0][0], lora_save_name
|
||||
)
|
||||
print(destination_path)
|
||||
print("Downloading external lora - " + input_id + " to " + destination_path)
|
||||
print("Downloading external lora - " + lora_url + " to " + destination_path)
|
||||
response = requests.get(
|
||||
input_id,
|
||||
lora_url,
|
||||
headers={"User-Agent": "Mozilla/5.0"},
|
||||
allow_redirects=True,
|
||||
)
|
||||
|
||||
@@ -20,11 +20,11 @@ class ComfyUIDeployExternalNumber:
|
||||
),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,11 +20,11 @@ class ComfyUIDeployExternalNumberInt:
|
||||
),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,11 +23,11 @@ class ComfyUIDeployExternalNumberSlider:
|
||||
),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,11 +20,11 @@ class ComfyUIDeployExternalText:
|
||||
),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -21,11 +21,11 @@ class ComfyUIDeployExternalTextList:
|
||||
"optional": {
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -764,14 +764,14 @@ class ComfyUIDeployExternalVideo:
|
||||
"optional": {
|
||||
"meta_batch": ("VHS_BatchManager",),
|
||||
"vae": ("VAE",),
|
||||
"default_value": (sorted(files),),
|
||||
"default_video": (sorted(files),),
|
||||
"display_name": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "Name of the node (optional)"},
|
||||
{"multiline": False, "default": ""},
|
||||
),
|
||||
"description": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "Description of the node (optional)"},
|
||||
{"multiline": True, "default": ""},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
@@ -834,7 +834,7 @@ class ComfyUIDeployExternalVideo:
|
||||
):
|
||||
out_file.write(chunk)
|
||||
else:
|
||||
video = kwargs.get("default_value", "")
|
||||
video = kwargs.get("default_video", None)
|
||||
if video is None:
|
||||
raise "No default video given and no external video provided"
|
||||
video_path = folder_paths.get_annotated_filepath(video.strip('"'))
|
||||
|
||||
+167
-73
@@ -1,4 +1,5 @@
|
||||
from io import BytesIO
|
||||
from pprint import pprint
|
||||
from aiohttp import web
|
||||
import os
|
||||
import requests
|
||||
@@ -56,35 +57,65 @@ 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):
|
||||
import time
|
||||
|
||||
async def async_request_with_retry(method, url, disable_timeout=False, token=None, **kwargs):
|
||||
global client_session
|
||||
await ensure_client_session()
|
||||
retry_delay = 1 # Start with 1 second delay
|
||||
initial_timeout = 5 # 5 seconds timeout for the initial connection
|
||||
|
||||
start_time = time.time()
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
# Set a timeout for the initial connection
|
||||
timeout = ClientTimeout(total=None, connect=initial_timeout)
|
||||
kwargs['timeout'] = timeout
|
||||
if not disable_timeout:
|
||||
timeout = ClientTimeout(total=None, connect=initial_timeout)
|
||||
kwargs['timeout'] = timeout
|
||||
|
||||
if token is not None:
|
||||
if 'headers' not in kwargs:
|
||||
kwargs['headers'] = {}
|
||||
kwargs['headers']['Authorization'] = f"Bearer {token}"
|
||||
|
||||
request_start = time.time()
|
||||
async with client_session.request(method, url, **kwargs) as response:
|
||||
request_end = time.time()
|
||||
logger.info(f"Request attempt {attempt + 1} took {request_end - request_start:.2f} seconds")
|
||||
|
||||
if response.status != 200:
|
||||
error_body = await response.text()
|
||||
logger.error(f"Request failed with status {response.status} and body {error_body}")
|
||||
# raise Exception(f"Request failed with status {response.status}")
|
||||
|
||||
response.raise_for_status()
|
||||
if method.upper() == 'GET':
|
||||
await response.read()
|
||||
|
||||
total_time = time.time() - start_time
|
||||
logger.info(f"Request succeeded after {total_time:.2f} seconds (attempt {attempt + 1}/{max_retries})")
|
||||
return response
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(f"Request timed out after {initial_timeout} seconds (attempt {attempt + 1}/{max_retries})")
|
||||
except ClientError as e:
|
||||
end_time = time.time()
|
||||
logger.error(f"Request failed (attempt {attempt + 1}/{max_retries}): {e}")
|
||||
logger.error(f"Time taken for failed attempt: {end_time - request_start:.2f} seconds")
|
||||
logger.error(f"Total time elapsed: {end_time - start_time:.2f} seconds")
|
||||
|
||||
# Log the response body for ClientError as well
|
||||
if hasattr(e, 'response') and e.response is not None:
|
||||
error_body = await e.response.text()
|
||||
logger.error(f"Error response body: {error_body}")
|
||||
|
||||
if attempt == max_retries - 1:
|
||||
logger.error(f"Request failed after {max_retries} attempts: {e}")
|
||||
# raise
|
||||
logger.warning(f"Request failed (attempt {attempt + 1}/{max_retries}): {e}")
|
||||
raise
|
||||
|
||||
# Wait before retrying
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay *= retry_delay_multiplier # Exponential backoff
|
||||
retry_delay *= retry_delay_multiplier
|
||||
|
||||
# If all retries fail, raise an exception
|
||||
raise Exception(f"Request failed after {max_retries} attempts")
|
||||
total_time = time.time() - start_time
|
||||
raise Exception(f"Request failed after {max_retries} attempts and {total_time:.2f} seconds")
|
||||
|
||||
from logging import basicConfig, getLogger
|
||||
|
||||
@@ -230,27 +261,30 @@ 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 isinstance(workflow_api[key]['inputs']['seed'], list):
|
||||
continue
|
||||
if workflow_api[key]['class_type'] == "PromptExpansion":
|
||||
workflow_api[key]['inputs']['seed'] = randomSeed(8)
|
||||
logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to PromptExpansion")
|
||||
continue
|
||||
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
|
||||
workflow_api[key]['inputs']['seed'] = randomSeed()
|
||||
logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to {workflow_api[key]['class_type']}")
|
||||
if 'inputs' in workflow_api[key]:
|
||||
if '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)
|
||||
logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to PromptExpansion")
|
||||
continue
|
||||
workflow_api[key]['inputs']['seed'] = randomSeed()
|
||||
logger.info(f"Applied random seed {workflow_api[key]['inputs']['seed']} to {workflow_api[key]['class_type']}")
|
||||
|
||||
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):
|
||||
# Loop through each of the inputs and replace them
|
||||
@@ -285,7 +319,7 @@ def apply_inputs_to_workflow(workflow_api: Any, inputs: Any, sid: str = None):
|
||||
value['inputs']["images"] = new_value
|
||||
|
||||
if value["class_type"] == "ComfyUIDeployExternalLora":
|
||||
value["inputs"]["default_lora_name"] = new_value
|
||||
value["inputs"]["lora_url"] = new_value
|
||||
|
||||
if value["class_type"] == "ComfyUIDeployExternalSlider":
|
||||
value["inputs"]["default_value"] = new_value
|
||||
@@ -332,6 +366,14 @@ def send_prompt(sid: str, inputs: StreamingPrompt):
|
||||
|
||||
@server.PromptServer.instance.routes.post("/comfyui-deploy/run")
|
||||
async def comfy_deploy_run(request):
|
||||
# Extract the bearer token from the Authorization header
|
||||
auth_header = request.headers.get('Authorization')
|
||||
token = None
|
||||
if auth_header:
|
||||
parts = auth_header.split()
|
||||
if len(parts) == 2 and parts[0].lower() == 'bearer':
|
||||
token = parts[1]
|
||||
|
||||
data = await request.json()
|
||||
|
||||
# In older version, we use workflow_api, but this has inputs already swapped in nextjs frontend, which is tricky
|
||||
@@ -347,13 +389,14 @@ async def comfy_deploy_run(request):
|
||||
prompt = {
|
||||
"prompt": workflow_api,
|
||||
"client_id": "comfy_deploy_instance", #api.client_id
|
||||
"prompt_id": prompt_id
|
||||
"prompt_id": prompt_id,
|
||||
}
|
||||
|
||||
prompt_metadata[prompt_id] = SimplePrompt(
|
||||
status_endpoint=data.get('status_endpoint'),
|
||||
file_upload_endpoint=data.get('file_upload_endpoint'),
|
||||
workflow_api=workflow_api
|
||||
workflow_api=workflow_api,
|
||||
token=token
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -391,7 +434,7 @@ async def comfy_deploy_run(request):
|
||||
|
||||
return web.json_response(res, status=status)
|
||||
|
||||
async def stream_prompt(data):
|
||||
async def stream_prompt(data, token):
|
||||
# In older version, we use workflow_api, but this has inputs already swapped in nextjs frontend, which is tricky
|
||||
workflow_api = data.get("workflow_api_raw")
|
||||
# The prompt id generated from comfy deploy, can be None
|
||||
@@ -411,7 +454,8 @@ async def stream_prompt(data):
|
||||
prompt_metadata[prompt_id] = SimplePrompt(
|
||||
status_endpoint=data.get('status_endpoint'),
|
||||
file_upload_endpoint=data.get('file_upload_endpoint'),
|
||||
workflow_api=workflow_api
|
||||
workflow_api=workflow_api,
|
||||
token=token
|
||||
)
|
||||
|
||||
# log('info', "Begin prompt", prompt=prompt)
|
||||
@@ -460,6 +504,14 @@ comfy_message_queues: Dict[str, asyncio.Queue] = {}
|
||||
async def stream_response(request):
|
||||
response = web.StreamResponse(status=200, reason='OK', headers={'Content-Type': 'text/event-stream'})
|
||||
await response.prepare(request)
|
||||
|
||||
# Extract the bearer token from the Authorization header
|
||||
auth_header = request.headers.get('Authorization')
|
||||
token = None
|
||||
if auth_header:
|
||||
parts = auth_header.split()
|
||||
if len(parts) == 2 and parts[0].lower() == 'bearer':
|
||||
token = parts[1]
|
||||
|
||||
pending = True
|
||||
data = await request.json()
|
||||
@@ -471,7 +523,7 @@ async def stream_response(request):
|
||||
log('info', 'Streaming prompt')
|
||||
|
||||
try:
|
||||
result = await stream_prompt(data=data)
|
||||
result = await stream_prompt(data=data, token=token)
|
||||
await response.write(f"event: event_update\ndata: {json.dumps(result)}\n\n".encode('utf-8'))
|
||||
# await response.write(.encode('utf-8'))
|
||||
await response.drain() # Ensure the buffer is flushed
|
||||
@@ -600,9 +652,10 @@ async def upload_file_endpoint(request):
|
||||
with open(file_path, 'rb') as f:
|
||||
headers = {
|
||||
"Content-Type": file_type,
|
||||
# "x-amz-acl": "public-read",
|
||||
"Content-Length": str(file_size)
|
||||
# "Content-Length": str(file_size)
|
||||
}
|
||||
if content.get('include_acl') is True:
|
||||
headers["x-amz-acl"] = "public-read"
|
||||
upload_response = await async_request_with_retry('PUT', upload_url, data=f, headers=headers)
|
||||
if upload_response.status == 200:
|
||||
return web.json_response({
|
||||
@@ -985,17 +1038,22 @@ async def send_json_override(self, event, data, sid=None):
|
||||
# await update_run_with_output(prompt_id, data)
|
||||
|
||||
if event == 'executed' and 'node' in data and 'output' in data:
|
||||
node_meta = None
|
||||
if prompt_id in prompt_metadata:
|
||||
node = data.get('node')
|
||||
class_type = prompt_metadata[prompt_id].workflow_api[node]['class_type']
|
||||
logger.info(f"Executed {class_type} {data}")
|
||||
node_meta = {
|
||||
"node_id": node,
|
||||
"node_class": class_type,
|
||||
}
|
||||
if class_type == "PreviewImage":
|
||||
logger.info("Skipping preview image")
|
||||
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'), node_meta=node_meta)
|
||||
# await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
|
||||
# update_run_with_output(prompt_id, data.get('output'))
|
||||
|
||||
@@ -1010,6 +1068,7 @@ async def update_run_live_status(prompt_id, live_status, calculated_progress: fl
|
||||
return
|
||||
|
||||
status_endpoint = prompt_metadata[prompt_id].status_endpoint
|
||||
token = prompt_metadata[prompt_id].token
|
||||
|
||||
if (status_endpoint is None):
|
||||
return
|
||||
@@ -1033,7 +1092,7 @@ async def update_run_live_status(prompt_id, live_status, calculated_progress: fl
|
||||
})
|
||||
|
||||
# requests.post(status_endpoint, json=body)
|
||||
await async_request_with_retry('POST', status_endpoint, json=body)
|
||||
await async_request_with_retry('POST', status_endpoint, token=token, json=body)
|
||||
|
||||
|
||||
async def update_run(prompt_id: str, status: Status):
|
||||
@@ -1064,7 +1123,8 @@ async def update_run(prompt_id: str, status: Status):
|
||||
try:
|
||||
# requests.post(status_endpoint, json=body)
|
||||
if (status_endpoint is not None):
|
||||
await async_request_with_retry('POST', status_endpoint, json=body)
|
||||
token = prompt_metadata[prompt_id].token
|
||||
await async_request_with_retry('POST', status_endpoint, token=token, json=body)
|
||||
|
||||
if (status_endpoint is not None) and cd_enable_run_log and (status == Status.SUCCESS or status == Status.FAILED):
|
||||
try:
|
||||
@@ -1094,7 +1154,7 @@ async def update_run(prompt_id: str, status: Status):
|
||||
]
|
||||
}
|
||||
|
||||
await async_request_with_retry('POST', status_endpoint, json=body)
|
||||
await async_request_with_retry('POST', status_endpoint, token=token, json=body)
|
||||
# requests.post(status_endpoint, json=body)
|
||||
except Exception as log_error:
|
||||
logger.info(f"Error reading log file: {log_error}")
|
||||
@@ -1115,7 +1175,7 @@ async def update_run(prompt_id: str, status: Status):
|
||||
})
|
||||
|
||||
|
||||
async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output"):
|
||||
async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output", item=None):
|
||||
"""
|
||||
Uploads file to S3 bucket using S3 client object
|
||||
:return: None
|
||||
@@ -1145,33 +1205,47 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
|
||||
logger.info(f"Uploading file {file}")
|
||||
|
||||
file_upload_endpoint = prompt_metadata[prompt_id].file_upload_endpoint
|
||||
|
||||
token = prompt_metadata[prompt_id].token
|
||||
filename = quote(filename)
|
||||
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}"
|
||||
async with aiofiles.open(file, 'rb') as f:
|
||||
data = await f.read()
|
||||
size = str(len(data))
|
||||
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)
|
||||
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()
|
||||
start_time = time.time() # Start timing here
|
||||
logger.info(f"Target URL: {target_url}")
|
||||
result = await async_request_with_retry("GET", target_url, disable_timeout=True, token=token)
|
||||
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 = await result.json()
|
||||
|
||||
logger.info(f"Result: {ok}")
|
||||
|
||||
start_time = time.time() # Start timing here
|
||||
|
||||
with open(file, 'rb') as f:
|
||||
data = f.read()
|
||||
start_time = time.time() # Start timing here
|
||||
headers = {
|
||||
# "x-amz-acl": "public-read",
|
||||
"Content-Type": content_type,
|
||||
"Content-Length": str(len(data)),
|
||||
# "Content-Length": size,
|
||||
}
|
||||
|
||||
if ok.get('include_acl') is True:
|
||||
headers["x-amz-acl"] = "public-read"
|
||||
|
||||
# response = requests.put(ok.get("url"), headers=headers, data=data)
|
||||
response = await async_request_with_retry('PUT', ok.get("url"), headers=headers, data=data)
|
||||
logger.info(f"Upload file response status: {response.status}, status text: {response.reason}")
|
||||
end_time = time.time() # End timing after the request is complete
|
||||
logger.info("Upload time: {:.2f} seconds".format(end_time - start_time))
|
||||
|
||||
if item is not None:
|
||||
file_download_url = ok.get("download_url")
|
||||
if file_download_url is not None:
|
||||
item["url"] = file_download_url
|
||||
item["upload_duration"] = end_time - start_time
|
||||
if ok.get("is_public") is not None:
|
||||
item["is_public"] = ok.get("is_public")
|
||||
|
||||
def have_pending_upload(prompt_id):
|
||||
if prompt_id in prompt_metadata and len(prompt_metadata[prompt_id].uploading_nodes) > 0:
|
||||
@@ -1264,8 +1338,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,29 +1354,46 @@ 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
|
||||
)
|
||||
content_type=file_type,
|
||||
item=item
|
||||
))
|
||||
|
||||
# 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):
|
||||
async def upload_in_background(prompt_id: str, data, node_id=None, have_upload=True, node_meta=None):
|
||||
try:
|
||||
await handle_upload(prompt_id, data, 'images', "content_type", "image/png")
|
||||
await handle_upload(prompt_id, data, 'files', "content_type", "image/png")
|
||||
# This will also be mp4
|
||||
await handle_upload(prompt_id, data, 'gifs', "format", "image/gif")
|
||||
await handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream")
|
||||
upload_tasks = [
|
||||
handle_upload(prompt_id, data, 'images', "content_type", "image/png"),
|
||||
handle_upload(prompt_id, data, 'files', "content_type", "image/png"),
|
||||
handle_upload(prompt_id, data, 'gifs', "format", "image/gif"),
|
||||
handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream")
|
||||
]
|
||||
|
||||
await asyncio.gather(*upload_tasks)
|
||||
|
||||
status_endpoint = prompt_metadata[prompt_id].status_endpoint
|
||||
token = prompt_metadata[prompt_id].token
|
||||
if have_upload:
|
||||
if status_endpoint is not None:
|
||||
body = {
|
||||
"run_id": prompt_id,
|
||||
"output_data": data,
|
||||
"node_meta": node_meta,
|
||||
}
|
||||
# pprint(body)
|
||||
await async_request_with_retry('POST', status_endpoint, token=token, json=body)
|
||||
await update_file_status(prompt_id, data, False, node_id=node_id)
|
||||
except Exception as e:
|
||||
await handle_error(prompt_id, data, e)
|
||||
|
||||
async def update_run_with_output(prompt_id, data, node_id=None):
|
||||
async def update_run_with_output(prompt_id, data, node_id=None, node_meta=None):
|
||||
if prompt_id not in prompt_metadata:
|
||||
return
|
||||
|
||||
@@ -1311,7 +1404,8 @@ async def update_run_with_output(prompt_id, data, node_id=None):
|
||||
|
||||
body = {
|
||||
"run_id": prompt_id,
|
||||
"output_data": data
|
||||
"output_data": data,
|
||||
"node_meta": node_meta,
|
||||
}
|
||||
have_upload_media = 'images' in data or 'files' in data or 'gifs' in data or 'mesh' in data
|
||||
if bypass_upload and have_upload_media:
|
||||
@@ -1325,15 +1419,15 @@ async def update_run_with_output(prompt_id, data, node_id=None):
|
||||
if have_upload_media:
|
||||
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_media))
|
||||
asyncio.create_task(upload_in_background(prompt_id, data, node_id=node_id, have_upload=have_upload_media, node_meta=node_meta))
|
||||
# await 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)
|
||||
|
||||
# requests.post(status_endpoint, json=body)
|
||||
if status_endpoint is not None:
|
||||
await async_request_with_retry('POST', status_endpoint, json=body)
|
||||
elif status_endpoint is not None:
|
||||
token = prompt_metadata[prompt_id].token
|
||||
await async_request_with_retry('POST', status_endpoint, token=token, json=body)
|
||||
|
||||
await send('outputs_uploaded', {
|
||||
"prompt_id": prompt_id
|
||||
|
||||
@@ -29,6 +29,8 @@ class SimplePrompt(BaseModel):
|
||||
status_endpoint: Optional[str]
|
||||
file_upload_endpoint: Optional[str]
|
||||
|
||||
token: Optional[str]
|
||||
|
||||
workflow_api: dict
|
||||
status: Status = Status.NOT_STARTED
|
||||
progress: set = set()
|
||||
|
||||
Reference in New Issue
Block a user