diff --git a/builder/modal-builder/src/template/data/restore_snapshot.py b/builder/modal-builder/src/template/data/restore_snapshot.py index 4ddbebe..8af439c 100644 --- a/builder/modal-builder/src/template/data/restore_snapshot.py +++ b/builder/modal-builder/src/template/data/restore_snapshot.py @@ -2,18 +2,20 @@ import json import requests import time import subprocess -import asyncio -import logging -# Set up the logger -logging.basicConfig(level=logging.INFO, format='%(message)s') -logger = logging.getLogger("restore_snapshot") +print("Starting ComfyUI") +command = ["python", "main.py", "--disable-auto-launch", "--disable-metadata", "--cpu"] # Start the server -# server_process = subprocess.Popen(command, cwd="/comfyui", stdout=subprocess.PIPE) +server_process = subprocess.Popen(command, cwd="/comfyui") def check_server(url, retries=50, delay=500): for i in range(retries): + # Check if the subprocess has ended + if server_process.poll() is not None: + print("Subprocess has ended") + break + try: response = requests.head(url) @@ -33,54 +35,10 @@ def check_server(url, retries=50, delay=500): ) return False -# Define the messages to look for -# success_message = "[ComfyUI-Manager] Restore snapshot done." -success_message = "To see the GUI go to: http://127.0.0.1:8188" -failure_message = "[ComfyUI-Manager] Restore snapshot failed." +root_url = "http://127.0.0.1:8188" -async def read_stream(stream, isStderr): - while True: - line = await stream.readline() - if line: - l = line.decode('utf-8').strip() +check_server(root_url, retries=1800, delay=1000) - if l == "": - continue - - if not isStderr: - logger.info(l) - - # If the output matches one of the messages, print it and break the loop - if success_message in l: - logger.info("Snapshot restore succeeded.") - break - elif failure_message in l: - logger.info("Snapshot restore failed.") - break - else: - # is error - # logger.error(l) - logger.error(l) - break - - else: - break - -async def main(): - command = "python main.py --disable-auto-launch --disable-metadata --cpu" - - server_process = await asyncio.subprocess.create_subprocess_shell(command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd="/comfyui") - - stdout_task = asyncio.create_task( - read_stream(server_process.stdout, False)) - stderr_task = asyncio.create_task( - read_stream(server_process.stderr, True)) - - await asyncio.wait([stdout_task, stderr_task]) - - print("Finished restoring snapshots.") - -asyncio.run(main()) \ No newline at end of file +# Close the server +server_process.terminate() +print("Finished installing dependencies.") \ No newline at end of file