diff --git a/src/cluster/comfyui_v1.py b/src/cluster/comfyui_v1.py index 8fcbb96..b3dfb18 100644 --- a/src/cluster/comfyui_v1.py +++ b/src/cluster/comfyui_v1.py @@ -78,9 +78,13 @@ with comfyui_image.imports(): from typing import Tuple, Any, Dict, Optional import loguru import sentry_sdk + import datetime + import tempfile + import traceback from modal import current_function_call_id from BowongModalFunctions.config import WorkerConfig - from BowongModalFunctions.models.web_model import SentryTransactionInfo + from BowongModalFunctions.models.web_model import SentryTransactionInfo, WebhookNotify + from BowongModalFunctions.utils.SentryUtils import SentryUtils from sentry_sdk.integrations.loguru import LoguruIntegration config = WorkerConfig() @@ -146,88 +150,87 @@ with comfyui_image.imports(): subprocess.run(cmd, shell=True, check=True) @modal.method() - def infer(self, workflow: Dict): - self.poll_server_health() - self.prompt_uuid = str(uuid.uuid4()) - print("Workflow JSON:") - print(json.dumps(workflow, indent=4, ensure_ascii=False)) - for node in workflow.values(): - if node.get("class_type") == "ComfyUIDeployExternalText": - if node.get("inputs").get("input_id") == "file_path": - file_to_move = node.get("inputs").get("default_value") - if file_to_move: - os.makedirs(os.path.dirname(file_to_move.replace("input_s3", "input")), exist_ok=True) - try: - shutil.copy(file_to_move, file_to_move.replace("input_s3", "input")) - except: - try: - print("Try download file from S3 manually") - # S3 Fallback - import boto3 - import yaml - with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", - encoding="utf-8", mode="r+") as myconfig: - yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) - awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], - aws_secret_access_key=yaml_config["aws_access_key"]) - awss3.meta.client.download_file(config.S3_bucket_name, config.comfyui_s3_input+"/"+file_to_move.split("input_s3/")[1], - file_to_move.replace("input_s3", "input")) - except: - raise Exception("Failed to download file from S3 manually") - node["inputs"]["default_value"] = node["inputs"]["default_value"].replace("input_s3", - "input") - with open(f"/root/{self.prompt_uuid}.json", "w", encoding="utf-8") as fi: - fi.write(json.dumps(workflow, ensure_ascii=False)) - cmd = f"comfy run --workflow /root/{self.prompt_uuid}.json --wait --timeout 1190 --verbose" - - self.file_prefix = [ - node.get("inputs") - for node in workflow.values() - if node.get("class_type") == "VHS_VideoCombine" - ][0]["filename_prefix"] - - subprocess.run(cmd, shell=True, check=True, timeout=1195) - # 获取按照文件时间创建排序的列表，默认是按时间升序 - file_list = os.listdir(output_dir) - new_file_list = sorted(file_list, key=lambda file: os.path.getctime(os.path.join(output_dir, file)), - reverse=True) - for f in new_file_list: - if f.startswith(self.file_prefix): - os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), - exist_ok=True) - try: - shutil.copy(os.path.join(output_dir, f), - os.path.join(output_dir.replace("output", "output_s3"), f)) - except: - try: - print("Try move file to S3 manually") - # S3 Fallback - import boto3 - import yaml - with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", - encoding="utf-8", mode="r+") as myconfig: - yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) - awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], - aws_secret_access_key=yaml_config["aws_access_key"]) - awss3.meta.client.upload_file(os.path.join(output_dir, f), config.S3_bucket_name, config.comfyui_s3_output+"/"+f) - except: - raise Exception("Failed to move file to S3 manually") - return f - - @modal.method() - def api(self, video_path: str, start_time: str = "00:00:03.600", filename_prefix=str(uuid.uuid4()), + async def api(self, video_path: str, start_time: str = "00:00:03.600", filename_prefix=str(uuid.uuid4()), tts_text1: str = "测试", tts_text2: str = "", tts_text3: str = "", tts_text4: str = "", anchor_id: str = "dawan", - speed: float = 1, sentry_trace: Optional[SentryTransactionInfo] = None): - # 导入必要模块 - import datetime - import tempfile - import traceback + speed: float = 1, sentry_trace: Optional[SentryTransactionInfo] = None, webhook:WebhookNotify = None): # 生成当前函数调用的唯一ID function_call_id = current_function_call_id() print(f"开始处理任务，会话ID: {self.session_id}, 调用ID: {function_call_id}") + @SentryUtils.webhook_handler(webhook=webhook, func_id=function_call_id) + async def infer(workflow: Dict): + self.poll_server_health() + self.prompt_uuid = str(uuid.uuid4()) + print("Workflow JSON:") + print(json.dumps(workflow, indent=4, ensure_ascii=False)) + for node in workflow.values(): + if node.get("class_type") == "ComfyUIDeployExternalText": + if node.get("inputs").get("input_id") == "file_path": + file_to_move = node.get("inputs").get("default_value") + if file_to_move: + os.makedirs(os.path.dirname(file_to_move.replace("input_s3", "input")), exist_ok=True) + try: + shutil.copy(file_to_move, file_to_move.replace("input_s3", "input")) + except: + try: + print("Try download file from S3 manually") + # S3 Fallback + import boto3 + import yaml + with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", + encoding="utf-8", mode="r+") as myconfig: + yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) + awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], + aws_secret_access_key=yaml_config["aws_access_key"]) + awss3.meta.client.download_file(config.S3_bucket_name, + config.comfyui_s3_input + "/" + + file_to_move.split("input_s3/")[1], + file_to_move.replace("input_s3", "input")) + except: + raise Exception("Failed to download file from S3 manually") + node["inputs"]["default_value"] = node["inputs"]["default_value"].replace("input_s3", + "input") + with open(f"/root/{self.prompt_uuid}.json", "w", encoding="utf-8") as fi: + fi.write(json.dumps(workflow, ensure_ascii=False)) + cmd = f"comfy run --workflow /root/{self.prompt_uuid}.json --wait --timeout 1190 --verbose" + + self.file_prefix = [ + node.get("inputs") + for node in workflow.values() + if node.get("class_type") == "VHS_VideoCombine" + ][0]["filename_prefix"] + + subprocess.run(cmd, shell=True, check=True, timeout=1195) + # 获取按照文件时间创建排序的列表，默认是按时间升序 + file_list = os.listdir(output_dir) + new_file_list = sorted(file_list, key=lambda file: os.path.getctime(os.path.join(output_dir, file)), + reverse=True) + for f in new_file_list: + if f.startswith(self.file_prefix): + os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), + exist_ok=True) + try: + shutil.copy(os.path.join(output_dir, f), + os.path.join(output_dir.replace("output", "output_s3"), f)) + except: + try: + print("Try move file to S3 manually") + # S3 Fallback + import boto3 + import yaml + with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", + encoding="utf-8", mode="r+") as myconfig: + yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) + awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], + aws_secret_access_key=yaml_config["aws_access_key"]) + awss3.meta.client.upload_file(os.path.join(output_dir, f), config.S3_bucket_name, + config.comfyui_s3_output + "/" + f) + except: + raise Exception("Failed to move file to S3 manually") + return f + # 使用临时文件 temp_log_path = None error_reported = False # 添加标志跟踪是否已报告错误 @@ -327,7 +330,7 @@ with comfyui_image.imports(): prep_span.set_tag("model_type", "video_generation") # 调用推理方法 - fname = self.infer.local(workflow_file) + fname = await infer(workflow_file) # 创建后处理子span with sentry_sdk.start_span(op="inference.postprocess", diff --git a/src/cluster/comfyui_v2.py b/src/cluster/comfyui_v2.py index 2c97b51..03ffd57 100644 --- a/src/cluster/comfyui_v2.py +++ b/src/cluster/comfyui_v2.py @@ -72,10 +72,14 @@ with comfyui_latentsync_1_5_image.imports(): import loguru import sentry_sdk import psutil + import datetime + import tempfile + import traceback from BowongModalFunctions.config import WorkerConfig from modal import current_function_call_id from sentry_sdk.integrations.loguru import LoguruIntegration - from BowongModalFunctions.models.web_model import SentryTransactionInfo + from BowongModalFunctions.models.web_model import SentryTransactionInfo, WebhookNotify + from BowongModalFunctions.utils.SentryUtils import SentryUtils config = WorkerConfig() vol = modal.Volume.from_name("comfyui-model", create_if_missing=True, environment_name=config.modal_environment) @@ -137,86 +141,88 @@ with comfyui_latentsync_1_5_image.imports(): subprocess.run(cmd, shell=True, check=True) @modal.method() - def infer(self, workflow: Dict): - self.poll_server_health() - self.prompt_uuid = str(uuid.uuid4()) - print("Workflow JSON:") - print(json.dumps(workflow, indent=4, ensure_ascii=False)) - for node in workflow.values(): - if node.get("class_type") == "ComfyUIDeployExternalText": - if node.get("inputs").get("input_id") == "file_path": - file_to_move = node.get("inputs").get("default_value") - if file_to_move: - os.makedirs(os.path.dirname(file_to_move.replace("input_s3", "input")), exist_ok=True) - try: - shutil.copy(file_to_move, file_to_move.replace("input_s3", "input")) - except: - try: - print("Try download file to S3 manually") - # S3 Fallback - import boto3 - import yaml - with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml",encoding="utf-8",mode="r+") as myconfig: - yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) - awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], - aws_secret_access_key=yaml_config["aws_access_key"]) - awss3.meta.client.download_file(config.S3_bucket_name, config.comfyui_s3_input + "/" + - file_to_move.split("input_s3/")[1], - file_to_move.replace("input_s3", "input")) - except: - raise Exception("Failed to download file from S3 manually") - node["inputs"]["default_value"] = node["inputs"]["default_value"].replace("input_s3", "input") - with open(f"/root/{self.prompt_uuid}.json", "w", encoding="utf-8") as fi: - fi.write(json.dumps(workflow, ensure_ascii=False)) - cmd = f"comfy run --workflow /root/{self.prompt_uuid}.json --wait --timeout 1190 --verbose" - - self.file_prefix = [ - node.get("inputs") - for node in workflow.values() - if node.get("class_type") == "VHS_VideoCombine" - ][0]["filename_prefix"] - - subprocess.run(cmd, shell=True, check=True, timeout=1195) - - file_list = os.listdir(output_dir) - # 获取按照文件时间创建排序的列表，默认是按时间升序 - new_file_list = sorted(file_list, key=lambda file: os.path.getctime(os.path.join(output_dir, file)), - reverse=True) - for f in new_file_list: - if f.startswith(self.file_prefix): - os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), exist_ok=True) - try: - shutil.copy(os.path.join(output_dir, f), os.path.join(output_dir.replace("output", "output_s3"), f)) - except: - try: - print("Try move file to S3 manually") - # S3 Fallback - import boto3 - import yaml - with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", encoding="utf-8",mode="r+") as myconfig: - yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) - awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], - aws_secret_access_key=yaml_config["aws_access_key"]) - awss3.meta.client.upload_file(os.path.join(output_dir, f), config.S3_bucket_name, config.comfyui_s3_output+"/"+f) - except: - raise Exception("Failed to move file to S3 manually") - return f - - - @modal.method() - def api(self, video_path: str, start_time: str = "00:00:03.600", filename_prefix=str(uuid.uuid4()), + async def api(self, video_path: str, start_time: str = "00:00:03.600", filename_prefix=str(uuid.uuid4()), tts_text1: str = "测试", tts_text2: str = "", tts_text3: str = "", tts_text4: str = "", anchor_id: str = "dawan", - speed: float = 1, sentry_trace: Optional[SentryTransactionInfo] = None): - # 导入必要模块 - import datetime - import tempfile - import traceback + speed: float = 1, sentry_trace: Optional[SentryTransactionInfo] = None, webhook:WebhookNotify = None): # 生成当前函数调用的唯一ID function_call_id = current_function_call_id() print(f"开始处理任务，会话ID: {self.session_id}, 调用ID: {function_call_id}") + @SentryUtils.webhook_handler(webhook=webhook, func_id=function_call_id) + async def infer(workflow: Dict): + self.poll_server_health() + self.prompt_uuid = str(uuid.uuid4()) + print("Workflow JSON:") + print(json.dumps(workflow, indent=4, ensure_ascii=False)) + for node in workflow.values(): + if node.get("class_type") == "ComfyUIDeployExternalText": + if node.get("inputs").get("input_id") == "file_path": + file_to_move = node.get("inputs").get("default_value") + if file_to_move: + os.makedirs(os.path.dirname(file_to_move.replace("input_s3", "input")), exist_ok=True) + try: + shutil.copy(file_to_move, file_to_move.replace("input_s3", "input")) + except: + try: + print("Try download file to S3 manually") + # S3 Fallback + import boto3 + import yaml + with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", + encoding="utf-8", mode="r+") as myconfig: + yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) + awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], + aws_secret_access_key=yaml_config["aws_access_key"]) + awss3.meta.client.download_file(config.S3_bucket_name, + config.comfyui_s3_input + "/" + + file_to_move.split("input_s3/")[1], + file_to_move.replace("input_s3", "input")) + except: + raise Exception("Failed to download file from S3 manually") + node["inputs"]["default_value"] = node["inputs"]["default_value"].replace("input_s3", + "input") + with open(f"/root/{self.prompt_uuid}.json", "w", encoding="utf-8") as fi: + fi.write(json.dumps(workflow, ensure_ascii=False)) + cmd = f"comfy run --workflow /root/{self.prompt_uuid}.json --wait --timeout 1190 --verbose" + + self.file_prefix = [ + node.get("inputs") + for node in workflow.values() + if node.get("class_type") == "VHS_VideoCombine" + ][0]["filename_prefix"] + + subprocess.run(cmd, shell=True, check=True, timeout=1195) + + file_list = os.listdir(output_dir) + # 获取按照文件时间创建排序的列表，默认是按时间升序 + new_file_list = sorted(file_list, key=lambda file: os.path.getctime(os.path.join(output_dir, file)), + reverse=True) + for f in new_file_list: + if f.startswith(self.file_prefix): + os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), + exist_ok=True) + try: + shutil.copy(os.path.join(output_dir, f), + os.path.join(output_dir.replace("output", "output_s3"), f)) + except: + try: + print("Try move file to S3 manually") + # S3 Fallback + import boto3 + import yaml + with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", + encoding="utf-8", mode="r+") as myconfig: + yaml_config = yaml.load(myconfig, Loader=yaml.FullLoader) + awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"], + aws_secret_access_key=yaml_config["aws_access_key"]) + awss3.meta.client.upload_file(os.path.join(output_dir, f), config.S3_bucket_name, + config.comfyui_s3_output + "/" + f) + except: + raise Exception("Failed to move file to S3 manually") + return f + # 使用临时文件 temp_log_path = None error_reported = False # 添加标志跟踪是否已报告错误 @@ -316,7 +322,7 @@ with comfyui_latentsync_1_5_image.imports(): prep_span.set_tag("model_type", "video_generation") # 调用推理方法 - fname = self.infer.local(workflow_file) + fname = await infer(workflow_file) # 创建后处理子span with sentry_sdk.start_span(op="inference.postprocess",