Remove clone from each model startup for ComfyUI (#146)

First pass at cleaning up this example. The should probably make it easy
to add the models to the cache but will just take it step by step and
hopefully invite other people to contribute.

@htrivedi99 would you mind taking a look please? I tested the dev
deployment.

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