Comfy API example (#628)

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}'")

06_gpu_and_ml/stable_diffusion/comfy_ui.py → 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,
         ]

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"
+      }
+    }
+  }