Merge pull request comfyanonymous#23 from goroshevsky/feature-webhook

feat: Add webhook url and extra params for serverless API
Eleo-ai · Dec 22, 2023 · 9147eb5 · 9147eb5
2 parents deefae2 + 62eef4d
commit 9147eb5
Show file tree

Hide file tree

Showing 13 changed files with 67 additions and 2 deletions.
diff --git a/build/COPY_ROOT/opt/serverless/docs/example_payloads/bound_image2image.json b/build/COPY_ROOT/opt/serverless/docs/example_payloads/bound_image2image.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "ckpt_name": "v1-5-pruned-emaonly.ckpt",
         "include_text": "photograph of a victorian woman, arms outstretched with angel wings. cloudy sky, meadow grass",
         "exclude_text": "watermark, text",

diff --git a/build/COPY_ROOT/opt/serverless/docs/example_payloads/bound_text2image.json b/build/COPY_ROOT/opt/serverless/docs/example_payloads/bound_text2image.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "steps": 20,
         "ckpt_name": "v1-5-pruned-emaonly.ckpt",
         "sampler_name": "euler",

diff --git a/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_controlnet_t2i_adapters.json b/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_controlnet_t2i_adapters.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "workflow_json": {
             "3": {
               "inputs": {

diff --git a/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_image2image.json b/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_image2image.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "workflow_json": {
             "3": {
               "inputs": {

diff --git a/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_text2image.json b/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_text2image.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "workflow_json": {
             "3": {
               "inputs": {

diff --git a/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_upscale.json b/build/COPY_ROOT/opt/serverless/docs/example_payloads/raw_upscale.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "workflow_json": {
             "3": {
               "inputs": {

diff --git a/build/COPY_ROOT/opt/serverless/docs/swagger/openapi.yaml b/build/COPY_ROOT/opt/serverless/docs/swagger/openapi.yaml
@@ -42,6 +42,14 @@ components:
                             type: string
                             required: false
                             description: Alternatively set AWS_BUCKET_NAME environment variable
+                        webhook_url:
+                            type: string
+                            required: false
+                            description: Webhook URL to invoke after a successful run or an error. Alternatively set WEBHOOK_URL environment variable
+                        webhook_extra_params:
+                            type: object
+                            required: false
+                            description: Extra params for webhook request
         RawWorkflow:
             description: Downloads URLs to input directory with no additional processing - Your application must modify the workflow as needed.
             allOf:

diff --git a/build/COPY_ROOT/opt/serverless/handlers/basehandler.py b/build/COPY_ROOT/opt/serverless/handlers/basehandler.py
@@ -43,7 +43,7 @@ def get_value(self, key, default = None):
             raise IndexError(f"{key} required but not set")
         elif key not in self.payload:
             return default
-        elif Network.is_url(self.payload[key]) and not key.startswith("aws_"):
+        elif Network.is_url(self.payload[key]) and not (key.startswith("aws_") or key.startswith("webhook_")):
             return self.get_url_content(self.payload[key])
         else:
             return self.payload[key]
@@ -88,6 +88,7 @@ def queue_job(self, timeout = 30):
                 time.sleep(0.5)
 
             if not self.is_server_ready():
+                self.invoke_webhook(success=False, error=f"Server not ready after timeout ({timeout}s)")
                 raise requests.RequestException(f"Server not ready after timeout ({timeout}s)")
 
             print ("Posting job to local server...")
@@ -96,12 +97,16 @@ def queue_job(self, timeout = 30):
             if "prompt_id" in response:
                 return response["prompt_id"]
             elif "node_errors" in response:
+                self.invoke_webhook(success=False, error=response["node_errors"])
                 raise requests.RequestException(response["node_errors"])
             elif "error" in response:
+                self.invoke_webhook(success=False, error=response["error"])
                 raise requests.RequestException(response["error"])
         except requests.RequestException:
+            self.invoke_webhook(success=False, error="Unknown error")
             raise
         except:
+            self.invoke_webhook(success=False, error="Unknown error")
             raise requests.RequestException("Failed to queue prompt")
 
     def get_job_status(self):
@@ -118,6 +123,7 @@ def get_job_status(self):
                 if self.comfyui_job_id in job:
                     return "pending"
         except:
+            self.invoke_webhook(success=False, error="Failed to queue job")
             raise requests.RequestException("Failed to queue job")
 
     def image_to_base64(self, path):
@@ -178,6 +184,25 @@ def get_s3_settings(self):
         settings["connect_attempts"] = 1
         return settings
 
+    def invoke_webhook(self, success = False, result = {}, error = ""):
+        webhook_url = self.get_value("webhook_url", os.environ.get("WEBHOOK_URL"))
+        webhook_extra_params = self.get_value("webhook_extra_params", {})
+
+        if Network.is_url(webhook_url):
+            data = {}
+            data["job_id"] = self.comfyui_job_id
+            data["request_id"] = self.request_id
+            data["success"] = success
+            if result:
+                data["result"] = result
+            if error:
+                data["error"] = error    
+            if webhook_extra_params:
+                data["extra_params"] = webhook_extra_params
+            Network.invoke_webhook(webhook_url, data)
+        else:
+            print("webhook_url is NOT valid!")    
+
     def handle(self):
         self.comfyui_job_id = self.queue_job(30)
 
@@ -188,4 +213,6 @@ def handle(self):
                 print (f"Waiting for {status} job to complete")
                 time.sleep(0.5)
 
-        return self.get_result(self.comfyui_job_id)
+        result = self.get_result(self.comfyui_job_id)
+        self.invoke_webhook(success=True, result=result)
+        return result
diff --git a/build/COPY_ROOT/opt/serverless/handlers/image2image.py b/build/COPY_ROOT/opt/serverless/handlers/image2image.py
@@ -61,6 +61,8 @@ def apply_modifiers(self):
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "ckpt_name": "v1-5-pruned-emaonly.ckpt",
         "include_text": "photograph of a victorian woman, arms outstretched with angel wings. cloudy sky, meadow grass",
         "exclude_text": "watermark, text",

diff --git a/build/COPY_ROOT/opt/serverless/handlers/rawworkflow.py b/build/COPY_ROOT/opt/serverless/handlers/rawworkflow.py
@@ -34,6 +34,8 @@ def apply_modifiers(self):
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "workflow_json": {
             "3": {
               "inputs": {

diff --git a/build/COPY_ROOT/opt/serverless/handlers/text2image.py b/build/COPY_ROOT/opt/serverless/handlers/text2image.py
@@ -65,6 +65,8 @@ def apply_modifiers(self):
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "steps": 20,
         "ckpt_name": "v1-5-pruned-emaonly.ckpt",
         "sampler_name": "euler",

diff --git a/build/COPY_ROOT/opt/serverless/providers/runpod/test_input.json b/build/COPY_ROOT/opt/serverless/providers/runpod/test_input.json
@@ -5,6 +5,8 @@
         "aws_secret_access_key": "your-s3-secret-access-key",
         "aws_endpoint_url": "https://my-endpoint.backblaze.com",
         "aws_bucket_name": "your-bucket",
+        "webhook_url": "your-webhook-url",
+        "webhook_extra_params": {},
         "workflow_json": {
             "3": {
               "inputs": {

diff --git a/build/COPY_ROOT/opt/serverless/utils/network.py b/build/COPY_ROOT/opt/serverless/utils/network.py
@@ -36,3 +36,13 @@ def download_file(url, target_dir, request_id):
 
         print(f"Downloaded {url} to {filepath}")
         return filepath
+
+    @staticmethod
+    def invoke_webhook(url, data):
+        try:
+            response = requests.post(url, json=data)
+            print(f"Invoke webhook {url} with data {data} - status {response.status_code}")
+            return response
+        except requests.exceptions.RequestException as e:
+            print(f"Error making POST request: {e}")
+            return None