Skip to content

Commit 33fb820

Browse files
authored
moving memory snapshot to separate example (#1202)
1 parent 11bcb75 commit 33fb820

File tree

6 files changed

+96
-23
lines changed

6 files changed

+96
-23
lines changed

06_gpu_and_ml/comfyui/comfyapp.py

Lines changed: 5 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
# python 06_gpu_and_ml/comfyui/comfyclient.py --modal-workspace $(modal profile current) --prompt "Surreal dreamscape with floating islands, upside-down waterfalls, and impossible geometric structures, all bathed in a soft, ethereal light"
2424
# ```
2525

26-
# ![example comfyui image](./flux_gen_image.jpeg)
26+
# ![example comfyui image](https://modal-cdn.com/cdnbot/flux_gen_imagesenr_0w3_209b7170.webp)
2727

2828
# The first inference will take ~1m since the container needs to launch the ComfyUI server and load Flux into memory. Successive calls on a warm container should take a few seconds.
2929

@@ -64,12 +64,6 @@
6464
# Add .run_commands(...) calls for any other custom nodes you want to download
6565
)
6666

67-
# We'll also add our own custom node that patches core ComfyUI so that we can use Modal's [memory snapshot](https://modal.com/docs/guide/memory-snapshot) feature to speed up cold starts (more on that on [running as an API](https://modal.com/docs/examples/comfyapp#running-comfyui-as-an-api)).
68-
image = image.add_local_dir(
69-
local_path=Path(__file__).parent / "memory_snapshot_helper",
70-
remote_path="/root/comfy/ComfyUI/custom_nodes/memory_snapshot_helper",
71-
copy=True,
72-
)
7367
# See [this post](https://modal.com/blog/comfyui-custom-nodes) for more examples
7468
# on how to install popular custom nodes like ComfyUI Impact Pack and ComfyUI IPAdapter Plus.
7569

@@ -164,29 +158,17 @@ def ui():
164158
scaledown_window=300, # 5 minute container keep alive after it processes an input
165159
gpu="L40S",
166160
volumes={"/cache": vol},
167-
enable_memory_snapshot=True, # snapshot container state for faster cold starts
168161
)
169162
@modal.concurrent(max_inputs=5) # run 5 inputs per container
170163
class ComfyUI:
171164
port: int = 8000
172165

173-
@modal.enter(snap=True)
166+
@modal.enter()
174167
def launch_comfy_background(self):
168+
# launch the ComfyUI server exactly once when the container starts
175169
cmd = f"comfy launch --background -- --port {self.port}"
176170
subprocess.run(cmd, shell=True, check=True)
177171

178-
@modal.enter(snap=False)
179-
def restore_snapshot(self):
180-
# initialize GPU for ComfyUI after snapshot restore
181-
# note: requires patching core ComfyUI, see the memory_snapshot_helper directory for more details
182-
import requests
183-
184-
response = requests.post(f"http://127.0.0.1:{self.port}/cuda/set_device")
185-
if response.status_code != 200:
186-
print("Failed to set CUDA device")
187-
else:
188-
print("Successfully set CUDA device")
189-
190172
@modal.method()
191173
def infer(self, workflow_path: str = "/root/workflow_api.json"):
192174
# sometimes the ComfyUI server stops responding (we think because of memory leaks), so this makes sure it's still up
@@ -256,10 +238,10 @@ def poll_server_health(self) -> Dict:
256238

257239
# This serves the `workflow_api.json` in this repo. When deploying your own workflows, make sure you select the "Export (API)" option in the ComfyUI menu:
258240

259-
# ![comfyui menu](./comfyui_menu.jpeg)
241+
# ![comfyui menu](https://modal-cdn.com/cdnbot/comfyui_menugo5j8ahx_27d72c45.webp)
260242

261243
# ## More resources
262-
# - [Alternative approach](https://modal.com/blog/comfyui-mem-snapshots) for deploying ComfyUI with memory snapshots
244+
# - Use [memory snapshots](https://modal.com/docs/guide/memory-snapshot) to speed up cold starts (check out the `memory_snapshot` directory on [Github](https://github.com/modal-labs/modal-examples/tree/main/06_gpu_and_ml/comfyui))
263245
# - Run a ComfyUI workflow as a [Python script](https://modal.com/blog/comfyui-prototype-to-production)
264246

265247
# - When to use [A1111 vs ComfyUI](https://modal.com/blog/a1111-vs-comfyui)
-62.8 KB
Binary file not shown.
-250 KB
Binary file not shown.
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
# Simple ComfyUI example using memory snapshot to speed up cold starts.
2+
3+
# CAUTION: Some custom nodes may not work with memory snapshots, especially if they make calls to torch (i.e. require a GPU) on initialization.
4+
# Run `modal deploy memory_snapshot_example.py` to deploy with memory snapshot enabled.
5+
6+
# Image building and model downloading is directly taken from the core example: https://modal.com/docs/examples/comfyapp
7+
# The notable changes are copying the custom node in the image and the cls object
8+
import subprocess
9+
from pathlib import Path
10+
11+
import modal
12+
13+
image = (
14+
modal.Image.debian_slim(python_version="3.11")
15+
.apt_install("git")
16+
.pip_install("fastapi[standard]==0.115.4")
17+
.pip_install("comfy-cli==1.3.8")
18+
.run_commands("comfy --skip-prompt install --fast-deps --nvidia --version 0.3.10")
19+
)
20+
21+
# Add custom node that patches core ComfyUI so that we can use Modal's [memory snapshot](https://modal.com/docs/guide/memory-snapshot)
22+
image = image.add_local_dir(
23+
local_path=Path(__file__).parent / "memory_snapshot_helper",
24+
remote_path="/root/comfy/ComfyUI/custom_nodes/memory_snapshot_helper",
25+
copy=True,
26+
)
27+
28+
29+
def hf_download():
30+
from huggingface_hub import hf_hub_download
31+
32+
flux_model = hf_hub_download(
33+
repo_id="Comfy-Org/flux1-schnell",
34+
filename="flux1-schnell-fp8.safetensors",
35+
cache_dir="/cache",
36+
)
37+
38+
subprocess.run(
39+
f"ln -s {flux_model} /root/comfy/ComfyUI/models/checkpoints/flux1-schnell-fp8.safetensors",
40+
shell=True,
41+
check=True,
42+
)
43+
44+
45+
vol = modal.Volume.from_name("hf-hub-cache", create_if_missing=True)
46+
47+
image = (
48+
image.pip_install("huggingface_hub[hf_transfer]==0.30.0")
49+
.env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
50+
.run_function(
51+
hf_download,
52+
volumes={"/cache": vol},
53+
)
54+
)
55+
56+
57+
app = modal.App(name="example-comfyui-memory-snapshot", image=image)
58+
59+
60+
@app.cls(
61+
max_containers=1,
62+
gpu="L40S",
63+
volumes={"/cache": vol},
64+
enable_memory_snapshot=True, # snapshot container state for faster cold starts
65+
)
66+
@modal.concurrent(max_inputs=10)
67+
class ComfyUIMemorySnapshot:
68+
port: int = 8000
69+
70+
# Snapshot ComfyUI server launch state, which includes import torch and custom node initialization (GPU not available during this step)
71+
@modal.enter(snap=True)
72+
def launch_comfy_background(self):
73+
cmd = f"comfy launch --background -- --port {self.port}"
74+
subprocess.run(cmd, shell=True, check=True)
75+
76+
# Restore ComfyUI server state. Re-enables the CUDA device for inference.
77+
@modal.enter(snap=False)
78+
def restore_snapshot(self):
79+
import requests
80+
81+
response = requests.post(f"http://127.0.0.1:{self.port}/cuda/set_device")
82+
if response.status_code != 200:
83+
print("Failed to set CUDA device")
84+
else:
85+
print("Successfully set CUDA device")
86+
87+
@modal.web_server(port, startup_timeout=60)
88+
def ui(self):
89+
subprocess.Popen(
90+
f"comfy launch -- --listen 0.0.0.0 --port {self.port}", shell=True
91+
)

0 commit comments

Comments
 (0)