FIX

heygem.py

@@ -0,0 +1,12 @@
+import site
+import subprocess
+import sys
+pkg_path = site.getsitepackages()[0]
+sys.path.append(f"{pkg_path}/heygem/code")
+
+p = subprocess.Popen(f"cd {pkg_path}/heygem/code && "
+                     "python3 app_local.py", shell=True)
+p.communicate()
+
+
+

server_with_s3.py

@@ -80,7 +80,7 @@ image = (
 
 app = modal.App(name="highlight-comfyui-s3", image=image)
 vol = modal.Volume.from_name("comfyui-model", create_if_missing=True)
-bucket_dict = modal.Dict.from_name("aws_s3_bucket", create_if_missing=True)
+bucket_dict = modal.Dict.from_name("aws_s3_bucket", create_if_missing=False)
 bucket_input = str(bucket_dict.get("INPUT"))
 bucket_output = str(bucket_dict.get("OUTPUT"))
 secret = modal.Secret.from_name("aws-s3-secret")
@@ -141,7 +141,23 @@ class ComfyUI:
                     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)
-                        shutil.copy(file_to_move, file_to_move.replace("input_s3", "input"))
+                        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 config:
+                                    yaml_config = yaml.load(config, 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(bucket_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 f:
             f.write(json.dumps(workflow, ensure_ascii=False))
@@ -164,7 +180,22 @@ class ComfyUI:
         for f in new_file_list:
             if f.startswith(file_prefix):
                 os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), exist_ok=True)
-                shutil.copy(os.path.join(output_dir, f), os.path.join(output_dir.replace("output", "output_s3"), f))
+                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 config:
+                            yaml_config = yaml.load(config, 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), bucket_output, f)
+                    except:
+                        raise Exception("Failed to move file to S3 manually")
                 return f
 
     @modal.fastapi_endpoint(method="POST")
@@ -175,6 +206,8 @@ class ComfyUI:
             new_workflow_file = item["prompt"]
             # run inference on the currently running container
             fname = self.infer.local(new_workflow_file)
+            if fname is None:
+                raise RuntimeError("Output File not found")
             j = {"file_name": fname}
             print("json", j)
             return Response(content=json.dumps(j), media_type="application/json")
@@ -183,9 +216,15 @@ class ComfyUI:
         finally:
             print("Purge Garbage By Restart Comfy")
             cmd = "comfy stop"
-            subprocess.run(cmd, shell=True, check=True)
+            try:
+                subprocess.run(cmd, shell=True, check=True)
+            except:
+                pass
             cmd = "comfy launch --background"
-            subprocess.run(cmd, shell=True, check=True)
+            try:
+                subprocess.run(cmd, shell=True, check=True)
+            except:
+                raise Exception("Failed to launch ComfyUI")
 
 
     def poll_server_health(self):
@@ -201,8 +240,16 @@ class ComfyUI:
             # if no response in 5 seconds, stop the container; Modal will schedule queued inputs on a new container
             print(f"Server health check failed: {str(e)} restarting ComfyUI...")
             try:
-                cmd = "comfy stop && comfy launch --background"
-                subprocess.run(cmd, shell=True, check=True)
+                cmd = "comfy stop"
+                try:
+                    subprocess.run(cmd, shell=True, check=True)
+                except:
+                    pass
+                cmd = "comfy launch --background"
+                try:
+                    subprocess.run(cmd, shell=True, check=True)
+                except:
+                    raise Exception("Failed to launch ComfyUI")
             except:
                 modal.experimental.stop_fetching_inputs()
                 raise Exception("ComfyUI server is not healthy, restart failed, stopping container")

server_with_s3_auth.py

@@ -82,7 +82,7 @@ image = (
 
 app = modal.App(name="highlight-comfyui-s3", image=image)
 vol = modal.Volume.from_name("comfyui-model", create_if_missing=True)
-bucket_dict = modal.Dict.from_name("aws_s3_bucket", create_if_missing=True)
+bucket_dict = modal.Dict.from_name("aws_s3_bucket", create_if_missing=False)
 bucket_input = str(bucket_dict.get("INPUT"))
 bucket_output = str(bucket_dict.get("OUTPUT"))
 secret = modal.Secret.from_name("aws-s3-secret")
@@ -144,10 +144,24 @@ class ComfyUI:
                     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)
-                        shutil.copy(file_to_move, file_to_move.replace("input_s3", "input"))
+                        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 config:
+                                    yaml_config = yaml.load(config, 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(bucket_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 f:
-            f.write(json.dumps(workflow, ensure_ascii=False))
+        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 890 --verbose"
         subprocess.run(cmd, shell=True, check=True, timeout=895)
 
@@ -167,7 +181,21 @@ class ComfyUI:
         for f in new_file_list:
             if f.startswith(file_prefix):
                 os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), exist_ok=True)
-                shutil.copy(os.path.join(output_dir, f), os.path.join(output_dir.replace("output", "output_s3"), f))
+                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 config:
+                            yaml_config = yaml.load(config, 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), bucket_output, f)
+                    except:
+                        raise Exception("Failed to move file to S3 manually")
                 return f
 
     @modal.fastapi_endpoint(method="POST")
@@ -184,6 +212,8 @@ class ComfyUI:
             new_workflow_file = item["prompt"]
             # run inference on the currently running container
             fname = self.infer.local(new_workflow_file)
+            if fname is None:
+                raise RuntimeError("Output File not found")
             j = {"file_name": fname}
             print("json", j)
             return Response(content=json.dumps(j), media_type="application/json")
@@ -192,10 +222,15 @@ class ComfyUI:
         finally:
             print("Purge Garbage By Restart Comfy")
             cmd = "comfy stop"
-            subprocess.run(cmd, shell=True, check=True)
+            try:
+                subprocess.run(cmd, shell=True, check=True)
+            except:
+                pass
             cmd = "comfy launch --background"
-            subprocess.run(cmd, shell=True, check=True)
-
+            try:
+                subprocess.run(cmd, shell=True, check=True)
+            except:
+                raise Exception("Failed to launch ComfyUI")
 
     def poll_server_health(self):
         import socket
@@ -210,8 +245,16 @@ class ComfyUI:
             # if no response in 5 seconds, stop the container; Modal will schedule queued inputs on a new container
             print(f"Server health check failed: {str(e)} restarting ComfyUI...")
             try:
-                cmd = "comfy stop && comfy launch --background"
-                subprocess.run(cmd, shell=True, check=True)
+                cmd = "comfy stop"
+                try:
+                    subprocess.run(cmd, shell=True, check=True)
+                except:
+                    pass
+                cmd = "comfy launch --background"
+                try:
+                    subprocess.run(cmd, shell=True, check=True)
+                except:
+                    raise Exception("Failed to launch ComfyUI")
             except:
                 modal.experimental.stop_fetching_inputs()
                 raise Exception("ComfyUI server is not healthy, restart failed, stopping container")

server_with_s3_auth_heygem.py

@@ -0,0 +1,274 @@
+import json
+import os
+import shutil
+import subprocess
+import sys
+import traceback
+import uuid
+from pathlib import Path
+from typing import Dict
+
+import modal
+from fastapi import Depends, HTTPException, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+image = (  # build up a Modal Image to run ComfyUI, step by step
+    modal.Image.from_registry(  # start from basic Linux with Python
+        "ubuntu/python:3.12-24.04_stable"
+    ).apt_install("python3.8-dev")
+    .apt_install("git")  # install git to clone ComfyUI
+    .apt_install("gcc")
+    .apt_install("libportaudio2")
+    .pip_install("fastapi[standard]==0.115.4")  # install web dependencies
+    .pip_install("comfy-cli==1.3.8")  # install comfy-cli
+    .pip_install("cos-python-sdk-v5")
+    .pip_install("sqlalchemy")
+    .pip_install("ultralytics")
+    .pip_install("tencentcloud-sdk-python")
+    .pip_install("pymysql")
+    .pip_install("Pillow")
+    .pip_install("ffmpy")
+    .pip_install("opencv-python")
+    .pip_install("av")
+    .pip_install("imageio")
+    .pip_install("loguru")
+    .pip_install("conformer==0.3.2",extra_options="--no-dependencies")
+    .pip_install("einops>0.6.1",extra_options="--no-dependencies")
+    .pip_install("openai-whisper")
+    .run_commands(  # use comfy-cli to install ComfyUI and its dependencies
+        "comfy --skip-prompt install --nvidia --version 0.3.10"
+    )
+)
+
+image = (
+    image.run_commands("comfy node install https://github.com/M1kep/ComfyLiterals")
+    .run_commands("comfy node install https://github.com/evanspearman/ComfyMath")
+    .run_commands("comfy node install https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved")
+    .run_commands("comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-Bowong.git")
+    .run_commands("comfy node install https://github.com/crystian/ComfyUI-Crystools")
+    .run_commands("comfy node install https://github.com/pythongosssss/ComfyUI-Custom-Scripts")
+    .run_commands("comfy node install https://github.com/BennyKok/comfyui-deploy")
+    .run_commands("comfy node install https://github.com/yolain/ComfyUI-Easy-Use")
+    .run_commands("comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-VideoHelperSuite.git")
+    .run_commands("comfy node install https://github.com/jags111/efficiency-nodes-comfyui")
+    .run_commands("comfy node install https://github.com/WASasquatch/was-node-suite-comfyui")
+    .run_commands("comfy node install https://github.com/rgthree/rgthree-comfy")
+    .run_commands("comfy node install https://github.com/cubiq/ComfyUI_essentials")
+    .run_commands("comfy node install https://github.com/melMass/comfy_mtb")
+    .run_commands("echo 3 && comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI_SparkTTS.git")
+    .run_commands("echo 1 && comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-CustomNode.git")
+    .run_commands("comfy node install https://e.coding.net/g-ldyi2063/dev/cosyvoice_comfyui.git")
+    .run_commands("echo 4 && comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-LatentSync-Node.git")
+    .run_commands(
+        "mkdir -p /root/comfy/ComfyUI/models/ComfyUI-CustomNode/model && rm -rf /root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/model && ln -s /root/comfy/ComfyUI/models/ComfyUI-CustomNode/model /root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/model"
+    ).run_commands(
+        "mkdir -p /root/comfy/ComfyUI/models/ComfyUI-LatentSync-Node/checkpoints && rm -rf /root/comfy/ComfyUI/custom_nodes/ComfyUI-LatentSync-Node/checkpoints && ln -s /root/comfy/ComfyUI/models/ComfyUI-LatentSync-Node/checkpoints /root/comfy/ComfyUI/custom_nodes/ComfyUI-LatentSync-Node/checkpoints"
+    ).run_commands(
+        "mkdir -p /root/comfy/ComfyUI/models/.cache/torch/hub/checkpoints && mkdir -p /root/.cache/torch/hub/checkpoints && rm -rf /root/.cache/torch/hub/checkpoints && ln -s /root/comfy/ComfyUI/models/.cache/torch/hub/checkpoints /root/.cache/torch/hub/checkpoints"
+    ).run_commands(
+        "mkdir -p /root/comfy/ComfyUI/models/stabilityai && ln -s /root/comfy/ComfyUI/models/stabilityai /root/comfy/ComfyUI/stabilityai"
+    ).run_commands(
+        "mkdir -p /root/comfy/ComfyUI/models/CosyVoice-ComfyUI/pretrained_models/CosyVoice-300M && mkdir -p /root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/CosyVoice-300M && rm -rf /root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/CosyVoice-300M && ln -s /root/comfy/ComfyUI/models/CosyVoice-ComfyUI/pretrained_models/CosyVoice-300M /root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/CosyVoice-300M"
+    ).run_commands(
+        "rm -rf /root/comfy/ComfyUI/models"
+    ).run_commands(
+        "apt update && apt install -y ffmpeg && ffmpeg -version"
+    ).add_local_file("config.yaml", "/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", copy=True
+    ).add_local_file("config.py", "/root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/tools/config.py", copy=True
+    ).add_local_file("heygem-1.0-py3-none-any.whl","/root/comfy/heygem-1.0-py3-none-any.whl", copy=True
+    ).add_local_file("heygem.py", "/root/comfy/heygem.py", copy=True
+    ).pip_install("/root/comfy/heygem-1.0-py3-none-any.whl")
+    .workdir("/root/comfy")
+
+    # Add .run_commands(...) calls for any other custom nodes you want to download
+)
+
+app = modal.App(name="highlight-comfyui-s3", image=image)
+vol = modal.Volume.from_name("comfyui-model", create_if_missing=True)
+bucket_dict = modal.Dict.from_name("aws_s3_bucket", create_if_missing=False)
+bucket_input = str(bucket_dict.get("INPUT"))
+bucket_output = str(bucket_dict.get("OUTPUT"))
+secret = modal.Secret.from_name("aws-s3-secret")
+# completed workflows write output images to this directory
+output_dir = "/root/comfy/ComfyUI/output"
+auth_scheme = HTTPBearer()
+
+
+@app.cls(
+    allow_concurrent_inputs=1,  # allow 10 concurrent API calls
+    max_containers=200,
+    min_containers=0,
+    buffer_containers=0,
+    scaledown_window=120,
+    # 5 minute container keep alive after it processes an input; increasing this value is a great way to reduce ComfyUI cold start times
+    timeout=900,
+    gpu=["L4", "T4"],
+    cpu=(2,16),
+    # memory=(32768, 32768), # (内存预留量, 内存使用上限)
+    memory=(20480,40960),
+    retries=1,
+    enable_memory_snapshot=True,
+    secrets=[secret, modal.Secret.from_name("web_auth_token")],
+    volumes={
+        "/root/comfy/ComfyUI/models": vol,
+        "/root/comfy/ComfyUI/input_s3": modal.CloudBucketMount(
+            bucket_name=bucket_input,
+            # bucket_endpoint_url="https://s3.%s.amazonaws.com" % os.environ["AWS_REGION"],
+            secret=secret,
+            key_prefix="/"
+        ),
+        "/root/comfy/ComfyUI/output_s3": modal.CloudBucketMount(
+            bucket_name=bucket_output,
+            # bucket_endpoint_url="https://s3.%s.amazonaws.com" % os.environ["AWS_REGION"],
+            secret=secret,
+            key_prefix="/"
+        ),
+    },
+)
+class ComfyUI:
+    @modal.enter()
+    def launch_comfy_background(self):
+        # starts the ComfyUI server in the background exactly once when the first input is received
+        self.session_id = str(uuid.uuid4())
+        cmd = ("echo client_uuid: {}"
+               " && rm -rf /root/comfy/ComfyUI/user"
+               " && mkdir -p /root/comfy/ComfyUI/output_s3/logs/{}"
+               " && ln -s /root/comfy/ComfyUI/output_s3/logs/{} /root/comfy/ComfyUI/user"
+               " && comfy launch --background"
+               " && nohup python heygem.py >> /root/comfy/ComfyUI/output_s3/logs/{}/heygem.log 2>&1 &").format(self.session_id,self.session_id, self.session_id, self.session_id)
+        subprocess.run(cmd, shell=True, check=True)
+
+    @modal.method()
+    def infer(self, workflow_json: str = ""):
+        self.poll_server_health()
+        self.prompt_uuid = str(uuid.uuid4())
+        # runs the comfy run --workflow command as a subprocess
+        workflow = json.loads(workflow_json)
+        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 config:
+                                    yaml_config = yaml.load(config, 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(bucket_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 890 --verbose"
+        subprocess.run(cmd, shell=True, check=True, timeout=895)
+
+        # looks up the name of the output image file based on the workflow
+        file_prefix = [
+            node.get("inputs")
+            for node in workflow.values()
+            if node.get("class_type") == "VHS_VideoCombine"
+        ][0]["filename_prefix"]
+
+        # returns the image as bytes
+        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)
+        # print("file_list", new_file_list)
+        for f in new_file_list:
+            if f.startswith(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 config:
+                            yaml_config = yaml.load(config, 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), bucket_output, f)
+                    except:
+                        raise Exception("Failed to move file to S3 manually")
+                return f
+
+    @modal.fastapi_endpoint(method="POST")
+    def api(self, item: Dict, token: HTTPAuthorizationCredentials = Depends(auth_scheme)):
+        if token.credentials != os.environ["AUTH_TOKEN"]:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Incorrect bearer token",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        from fastapi import Response
+        try:
+            # save this updated workflow to a new file
+            new_workflow_file = item["prompt"]
+            # run inference on the currently running container
+            fname = self.infer.local(new_workflow_file)
+            if fname is None:
+                raise RuntimeError("Output File not found")
+            j = {"file_name": fname}
+            print("json", j)
+            return Response(content=json.dumps(j), media_type="application/json")
+        except Exception as e:
+            raise RuntimeError(str(e))
+        finally:
+            print("Purge Garbage By Restart Comfy")
+            cmd = "comfy stop"
+            try:
+                subprocess.run(cmd, shell=True, check=True)
+            except:
+                pass
+            cmd = "comfy launch --background"
+            try:
+                subprocess.run(cmd, shell=True, check=True)
+            except:
+                raise Exception("Failed to launch ComfyUI")
+
+    def poll_server_health(self):
+        import socket
+        import urllib
+
+        try:
+            # dummy request to check if the server is healthy
+            req = urllib.request.Request("http://127.0.0.1:8188/system_stats")
+            urllib.request.urlopen(req, timeout=5)
+            print("ComfyUI server is healthy")
+        except (socket.timeout, urllib.error.URLError) as e:
+            # if no response in 5 seconds, stop the container; Modal will schedule queued inputs on a new container
+            print(f"Server health check failed: {str(e)} restarting ComfyUI...")
+            try:
+                cmd = "comfy stop"
+                try:
+                    subprocess.run(cmd, shell=True, check=True)
+                except:
+                    pass
+                cmd = "comfy launch --background"
+                try:
+                    subprocess.run(cmd, shell=True, check=True)
+                except:
+                    raise Exception("Failed to launch ComfyUI")
+            except:
+                modal.experimental.stop_fetching_inputs()
+                raise Exception("ComfyUI server is not healthy, restart failed, stopping container")
+
+    # @modal.web_endpoint(method="POST", label="tk")
+    # def tk_api(self, item: Dict):
+    #     pass
+    #
+