From 745d3f9d35d137e784199faf64ed661d7be2c26b Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Thu, 11 Jan 2024 03:47:08 +0800
Subject: [PATCH] fix: install snapshots

---
 builder/modal-builder/src/template/app.py     | 14 +--
 .../src/template/data/restore_snapshot.py     | 85 +++++++++++++++++++
 2 files changed, 94 insertions(+), 5 deletions(-)
 create mode 100644 builder/modal-builder/src/template/data/restore_snapshot.py

diff --git a/builder/modal-builder/src/template/app.py b/builder/modal-builder/src/template/app.py
index 96a7a70..5334ee9 100644
--- a/builder/modal-builder/src/template/app.py
+++ b/builder/modal-builder/src/template/app.py
@@ -28,6 +28,7 @@ web_app = FastAPI()
 print(config)
 print("deploy_test ", deploy_test)
 stub = Stub(name=config["name"])
+# print(stub.app_id)
 
 if not deploy_test:
     # dockerfile_image = Image.from_dockerfile(f"{current_directory}/Dockerfile", context_mount=Mount.from_local_dir(f"{current_directory}/data", remote_path="/data"))
@@ -36,6 +37,10 @@ if not deploy_test:
     dockerfile_image = (
         modal.Image.debian_slim()
         .apt_install("git", "wget")
+        .pip_install(
+            "git+https://github.com/modal-labs/asgiproxy.git", "httpx", "tqdm"
+        )
+        .apt_install("libgl1-mesa-glx", "libglib2.0-0")
         .run_commands(
             # Basic comfyui setup
             "git clone https://github.com/comfyanonymous/ComfyUI.git /comfyui",
@@ -50,7 +55,6 @@ if not deploy_test:
             "cd /comfyui/custom_nodes && git clone https://github.com/BennyKok/comfyui-deploy.git",
         )
         .copy_local_file(f"{current_directory}/data/extra_model_paths.yaml", "/comfyui")
-        .copy_local_file(f"{current_directory}/data/snapshot.json", "/comfyui/custom_nodes/ComfyUI-Manager/startup-scripts/restore-snapshot.json")
 
         .copy_local_file(f"{current_directory}/data/start.sh", "/start.sh")
         .run_commands("chmod +x /start.sh")
@@ -61,9 +65,9 @@ if not deploy_test:
 
         .run_commands("python install_deps.py")
 
-        .pip_install(
-            "git+https://github.com/modal-labs/asgiproxy.git", "httpx", "tqdm"
-        )
+        .copy_local_file(f"{current_directory}/data/restore_snapshot.py", "/")
+        .copy_local_file(f"{current_directory}/data/snapshot.json", "/comfyui/custom_nodes/ComfyUI-Manager/startup-scripts/restore-snapshot.json")
+        .run_commands("python restore_snapshot.py")
     )
 
 # Time to wait between API check attempts in milliseconds
@@ -295,4 +299,4 @@ def comfyui_app():
         },
     )()
 
-    return make_simple_proxy_app(ProxyContext(config))
+    return make_simple_proxy_app(ProxyContext(config))
\ No newline at end of file
diff --git a/builder/modal-builder/src/template/data/restore_snapshot.py b/builder/modal-builder/src/template/data/restore_snapshot.py
new file mode 100644
index 0000000..27d898c
--- /dev/null
+++ b/builder/modal-builder/src/template/data/restore_snapshot.py
@@ -0,0 +1,85 @@
+import json
+import requests
+import time
+import subprocess
+import asyncio
+import logging
+
+# Set up the logger
+logging.basicConfig(level=logging.INFO)
+
+# Start the server
+# server_process = subprocess.Popen(command, cwd="/comfyui", stdout=subprocess.PIPE)
+
+def check_server(url, retries=50, delay=500):
+    for i in range(retries):
+        try:
+            response = requests.head(url)
+
+            # If the response status code is 200, the server is up and running
+            if response.status_code == 200:
+                print(f"builder - API is reachable")
+                return True
+        except requests.RequestException as e:
+            # If an exception occurs, the server may not be ready
+            pass
+
+        # Wait for the specified delay before retrying
+        time.sleep(delay / 1000)
+
+    print(
+        f"builder- Failed to connect to server at {url} after {retries} attempts."
+    )
+    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."
+
+async def read_stream(stream, isStderr):
+    while True:
+        line = await stream.readline()
+        if line:
+            l = line.decode('utf-8').strip()
+
+            if l == "":
+                continue
+
+            if not isStderr:
+                logging.info(l)
+
+                # If the output matches one of the messages, print it and break the loop
+                if success_message in l:
+                    logging.info("Snapshot restore succeeded.")
+                    break
+                elif failure_message in l:
+                    logging.info("Snapshot restore failed.")
+                    break
+            else:
+                # is error
+                # logger.error(l)
+                logging.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