Remove clone from each model startup for ComfyUI

comfyui-truss/config.yaml

@@ -3,22 +3,12 @@ description: Deploy a ComfyUI workflow as a Truss
 model_metadata:
   example_model_input: {"workflow_values": {"positive_prompt": "An igloo on a snowy day, 4k, hd", "negative_prompt": "blurry, text, low quality", "controlnet_image": "https://storage.googleapis.com/logos-bucket-01/baseten_logo.png"}}
 model_name: ComfyUI Workflow
-python_version: py310
+base_image:
+  image: bolabaseten/comfyui-truss-base:6a7bc35
+  python_executable_path: /usr/bin/python3
 requirements:
   - websocket-client==1.6.4
-  - torch==2.1.0
-  - torchsde==0.2.6
-  - torchvision==0.16.0
-  - einops==0.7.0
-  - transformers==4.34.1
-  - safetensors==0.4.0
-  - aiohttp==3.8.6
   - accelerate==0.23.0
-  - pyyaml>=6.0.0
-  - Pillow==10.1.0
-  - scipy==1.11.3
-  - tqdm==4.66.1
-  - psutil>=5.9.4
 resources:
   cpu: "3"
   memory: 14Gi

comfyui-truss/model/helpers.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 import subprocess
 import sys
@@ -8,27 +9,29 @@
 
 import requests
 
+COMFYUI_DIR = "/app/ComfyUI"
+
 
 def download_model(model_url, destination_path):
-    print(f"Downloading model {model_url} ...")
+    logging.info(f"Downloading model {model_url} ...")
     try:
         response = requests.get(model_url, stream=True)
         response.raise_for_status()
-        print("download response: ", response)
+        logging.debug("download response: ", response)
 
         # Open the destination file and write the content in chunks
-        print("opening: ", destination_path)
+        logging.debug("opening: ", destination_path)
         with open(destination_path, "wb") as file:
-            print("writing chunks...")
+            logging.debug("writing chunks...")
             for chunk in response.iter_content(chunk_size=8192):
                 if chunk:  # Filter out keep-alive new chunks
                     file.write(chunk)
 
-            print("done writing chunks!!!!")
+            logging.debug("done writing chunks!")
 
-        print(f"Downloaded file to: {destination_path}")
+        logging.info(f"Downloaded file to: {destination_path}")
     except requests.exceptions.RequestException as e:
-        print(f"Download failed: {e}")
+        logging.error(f"Download failed: {e}")
 
 
 def download_tempfile(file_url, filename):
@@ -40,48 +43,35 @@ def download_tempfile(file_url, filename):
         temp_file.write(response.content)
         return temp_file.name, temp_file
     except Exception as e:
-        print("Error downloading and saving image:", e)
+        logging.error("Error downloading and saving image:", e)
         return None
 
 
 def setup_comfyui(original_working_directory, data_dir):
-    git_repo_url = "https://github.com/comfyanonymous/ComfyUI.git"
-    commit_hash = "248aa3e56355d75ac3d8632af769e6c700d9bfac"
-    git_clone_command = ["git", "clone", git_repo_url]
 
     try:
-        # clone the repo
-        subprocess.run(git_clone_command, check=True)
-        print("Git repository cloned successfully!")
-
-        os.chdir(os.path.join(original_working_directory, "ComfyUI"))
-
-        # Pin comfyUI to a specific commit
-        checkout_command = ["git", "checkout", commit_hash]
-        subprocess.run(checkout_command, check=True)
-
         model_json = os.path.join(original_working_directory, data_dir, "model.json")
         with open(model_json, "r") as file:
             data = json.load(file)
 
-        print(f"model json file: {data}")
+        logging.debug(f"model json file: {data}")
 
         if data and len(data) > 0:
             for model in data:
                 download_model(
                     model_url=model.get("url"),
                     destination_path=os.path.join(
-                        os.getcwd(), "models", model.get("path")
+                        COMFYUI_DIR, "models", model.get("path")
                     ),
                 )
 
-        print("Finished downloading models!")
+        logging.debug("Finished downloading models!")
 
         # run the comfy-ui server
-        subprocess.run([sys.executable, "main.py"], check=True)
+        subprocess.run([sys.executable, "main.py"], cwd=COMFYUI_DIR, check=True)
 
     except Exception as e:
-        print(e)
+        logging.error(e)
         raise Exception("Error setting up comfy UI repo")
 
 

dockerfiles/ComfyUI.dockerfile

@@ -0,0 +1,7 @@
+FROM baseten/truss-server-base:3.11-gpu-v0.7.17
+
+ARG COMMIT_HASH 6a7bc35db845179a26e62534f3d4b789151e52fe
+
+RUN git clone https://github.com/comfyanonymous/ComfyUI.git /app/ComfyUI
+
+RUN cd /app/ComfyUI; git checkout $COMMIT_HASH; pip install -r requirements.txt