From 62eef4d46fd57a7cee9d6807aacef215b58ad06b Mon Sep 17 00:00:00 2001 From: Vadim Goroshevsky Date: Thu, 21 Dec 2023 15:52:46 +0200 Subject: [PATCH] feat: Add webhook url and extra params for serverless API --- .../example_payloads/bound_image2image.json | 2 ++ .../example_payloads/bound_text2image.json | 2 ++ .../raw_controlnet_t2i_adapters.json | 2 ++ .../example_payloads/raw_image2image.json | 2 ++ .../docs/example_payloads/raw_text2image.json | 2 ++ .../docs/example_payloads/raw_upscale.json | 2 ++ .../opt/serverless/docs/swagger/openapi.yaml | 8 +++++ .../opt/serverless/handlers/basehandler.py | 31 +++++++++++++++++-- .../opt/serverless/handlers/image2image.py | 2 ++ .../opt/serverless/handlers/rawworkflow.py | 2 ++ .../opt/serverless/handlers/text2image.py | 2 ++ .../providers/runpod/test_input.json | 2 ++ .../COPY_ROOT/opt/serverless/utils/network.py | 10 ++++++ 13 files changed, 67 insertions(+), 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 index 1367e936..b7964daf 100644 --- 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 index 19c93ade..b2c92158 100644 --- 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 index 4ca1418a..a047ed90 100644 --- 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 index 5c8a76c9..a65bd5fd 100644 --- 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 index 84244ea6..14a36cf0 100644 --- 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 index db407f71..ae5fcbed 100644 --- 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 index 2a01b5fe..50ba2664 100644 --- 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 index 395332ef..51997088 100644 --- 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) \ No newline at end of file + result = self.get_result(self.comfyui_job_id) + self.invoke_webhook(success=True, result=result) + return result \ No newline at end of file diff --git a/build/COPY_ROOT/opt/serverless/handlers/image2image.py b/build/COPY_ROOT/opt/serverless/handlers/image2image.py index 7feedd90..152b1645 100644 --- 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 index d63e615b..a03dd738 100644 --- 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 index 68e8191e..d64f4230 100644 --- 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 index 5c8a76c9..a65bd5fd 100644 --- 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 index 355ddf4c..b8afbd80 100644 --- 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