diff --git a/06_gpu_and_ml/comfyui/comfy_api.py b/06_gpu_and_ml/comfyui/comfy_api.py new file mode 100644 index 000000000..7d2fe43b0 --- /dev/null +++ b/06_gpu_and_ml/comfyui/comfy_api.py @@ -0,0 +1,148 @@ +# --- +# lambda-test: false +# --- +# +# # Make API calls to a ComfyUI server +# +# This example shows you how to execute ComfyUI workflows via ComfyUI's API. +# +# ![example comfyui workspace](./comfyui-hero.png) +import json +import os +import pathlib +import urllib +import uuid + +import modal + +stub = modal.Stub(name="example-comfy-api") +image = modal.Image.debian_slim(python_version="3.10").pip_install( + "websocket-client==1.6.4" +) + +# This workflow JSON has been exported by running `comfy_ui.py` and downloading the JSON +# using the web UI. +comfyui_workflow_data_path = assets_path = ( + pathlib.Path(__file__).parent / "comfy_ui_workflow.json" +) + + +def fetch_image( + filename: str, subfolder: str, folder_type: str, server_address: str +) -> bytes: + data = {"filename": filename, "subfolder": subfolder, "type": folder_type} + url_values = urllib.parse.urlencode(data) + with urllib.request.urlopen( + "https://{}/view?{}".format(server_address, url_values) + ) as response: + return response.read() + + +def run_workflow( + ws, prompt: str, server_address: str, client_id: str +) -> list[bytes]: + p = {"prompt": prompt, "client_id": client_id} + data = json.dumps(p).encode("utf-8") + req = urllib.request.Request( + "https://{}/prompt".format(server_address), data=data + ) + response_data = json.loads(urllib.request.urlopen(req).read()) + prompt_id = response_data["prompt_id"] + output_images = {} + + while True: + out = ws.recv() + if isinstance(out, str): + print(f"recieved str msg from websocket. ws msg: {out}") + try: + message = json.loads(out) + except json.JSONDecodeError: + print(f"expected valid JSON but got: {out}") + raise + print(f"received msg from ws: {message}") + if message["type"] == "executing": + data = message["data"] + if data["node"] is None and data["prompt_id"] == prompt_id: + break # Execution is done! + else: + continue # previews are binary data + + # Fetch workflow execution history, which contains references to our completed images. + with urllib.request.urlopen( + f"https://{server_address}/history/{prompt_id}" + ) as response: + output = json.loads(response.read()) + history = output[prompt_id].get("outputs") if prompt_id in output else None + if not history: + raise RuntimeError( + f"Unexpected missing ComfyUI history for {prompt_id}" + ) + for node_id in history: + node_output = history[node_id] + if "images" in node_output: + images_output = [] + for image in node_output["images"]: + image_data = fetch_image( + filename=image["filename"], + subfolder=image["subfolder"], + folder_type=image["type"], + server_address=server_address, + ) + images_output.append(image_data) + output_images[node_id] = images_output + + return output_images + + +@stub.function(image=image) +def query_comfy_via_api(workflow_data: dict, prompt: str, server_address: str): + import websocket + + # Modify workflow to use requested prompt. + workflow_data["2"]["inputs"]["text"] = prompt + + # Make a websocket connection to the ComfyUI server. The server will + # will stream workflow execution updates over this websocket. + ws = websocket.WebSocket() + client_id = str(uuid.uuid4()) + ws_address = f"wss://{server_address}/ws?clientId={client_id}" + print(f"Connecting to websocket at {ws_address} ...") + ws.connect(ws_address) + print(f"Connected at {ws_address}. Running workflow via API") + images = run_workflow(ws, workflow_data, server_address, client_id) + image_list = [] + for node_id in images: + for image_data in images[node_id]: + image_list.append(image_data) + return image_list + + +@stub.local_entrypoint() +def main(prompt: str = "bag of wooden blocks") -> None: + workflow_data = json.loads(comfyui_workflow_data_path.read_text()) + + # Run the ComfyUI server app and make an API call to it. + # The ComfyUI server app will shutdown on exit of this context manager. + from comfy_ui import stub as comfyui_stub + + with comfyui_stub.run( + show_progress=False, # hide server app's modal progress logs + stdout=open(os.devnull, "w"), # hide server app's application logs + ) as comfyui_app: + print(f"{comfyui_app.app_id=}") + comfyui_url = comfyui_app.web.web_url + + server_address = comfyui_url.split("://")[1] # strip protocol + + image_list = query_comfy_via_api.remote( + workflow_data=workflow_data, + prompt=prompt, + server_address=server_address, + ) + + for i, img_bytes in enumerate(image_list): + filename = f"comfyui_{i}.png" + with open(filename, "wb") as f: + f.write(img_bytes) + f.close() + print(f"saved '{filename}'") diff --git a/06_gpu_and_ml/stable_diffusion/comfy_ui.py b/06_gpu_and_ml/comfyui/comfy_ui.py similarity index 95% rename from 06_gpu_and_ml/stable_diffusion/comfy_ui.py rename to 06_gpu_and_ml/comfyui/comfy_ui.py index cbff0a92d..9c24297d4 100644 --- a/06_gpu_and_ml/stable_diffusion/comfy_ui.py +++ b/06_gpu_and_ml/comfyui/comfy_ui.py @@ -68,9 +68,11 @@ def download_checkpoint(): f"cd /root && git checkout {comfyui_commit_sha}", "cd /root && pip install xformers!=0.0.18 -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121", ) - # Use fork until https://github.com/valohai/asgiproxy/pull/11 is merged. + # Use fork of https://github.com/valohai/asgiproxy with bugfixes. .pip_install( - "git+https://github.com/modal-labs/asgiproxy.git", "httpx", "tqdm" + "git+https://github.com/modal-labs/asgiproxy.git@ef25fe52cf226f9a635e87616e7c049e451e2bd8", + "httpx", + "tqdm", ) .run_function(download_checkpoint) ) @@ -95,6 +97,7 @@ def spawn_comfyui_in_background(): "python", "main.py", "--dont-print-server", + "--multi-user", "--port", PORT, ] diff --git a/06_gpu_and_ml/comfyui/comfy_ui_workflow.json b/06_gpu_and_ml/comfyui/comfy_ui_workflow.json new file mode 100644 index 000000000..8518f054b --- /dev/null +++ b/06_gpu_and_ml/comfyui/comfy_ui_workflow.json @@ -0,0 +1,181 @@ +{ + "1": { + "inputs": { + "ckpt_name": "dreamlike-photoreal-2.0.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "2": { + "inputs": { + "text": "a bag of wooden blocks", + "clip": [ + "1", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Positive)" + } + }, + "3": { + "inputs": { + "text": "bag of noodles", + "clip": [ + "1", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Negative)" + } + }, + "4": { + "inputs": { + "seed": 350088449706888, + "steps": 12, + "cfg": 8, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 1, + "model": [ + "1", + 0 + ], + "positive": [ + "2", + 0 + ], + "negative": [ + "3", + 0 + ], + "latent_image": [ + "5", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "5": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "6": { + "inputs": { + "samples": [ + "8", + 0 + ], + "vae": [ + "1", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "7": { + "inputs": { + "images": [ + "6", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Preview Image" + } + }, + "8": { + "inputs": { + "add_noise": "enable", + "noise_seed": 350088449706888, + "steps": 30, + "cfg": 8, + "sampler_name": "euler", + "scheduler": "karras", + "start_at_step": 12, + "end_at_step": 10000, + "return_with_leftover_noise": "disable", + "model": [ + "1", + 0 + ], + "positive": [ + "2", + 0 + ], + "negative": [ + "3", + 0 + ], + "latent_image": [ + "10", + 0 + ] + }, + "class_type": "KSamplerAdvanced", + "_meta": { + "title": "KSampler (Advanced)" + } + }, + "10": { + "inputs": { + "upscale_method": "nearest-exact", + "scale_by": 2, + "samples": [ + "4", + 0 + ] + }, + "class_type": "LatentUpscaleBy", + "_meta": { + "title": "Upscale Latent By" + } + }, + "11": { + "inputs": { + "samples": [ + "4", + 0 + ], + "vae": [ + "1", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "12": { + "inputs": { + "images": [ + "11", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Preview Image" + } + } + } diff --git a/06_gpu_and_ml/stable_diffusion/comfyui-hero.png b/06_gpu_and_ml/comfyui/comfyui-hero.png similarity index 100% rename from 06_gpu_and_ml/stable_diffusion/comfyui-hero.png rename to 06_gpu_and_ml/comfyui/comfyui-hero.png