Skip to content

Commit

Permalink
Add serverless framework (BETA)
Browse files Browse the repository at this point in the history
  • Loading branch information
robballantyne committed Nov 11, 2023
1 parent 8fa1b8b commit 7c9b02a
Show file tree
Hide file tree
Showing 7 changed files with 334 additions and 6 deletions.
1 change: 0 additions & 1 deletion build/COPY_ROOT/opt/serverless/.move_complete

This file was deleted.

2 changes: 1 addition & 1 deletion build/COPY_ROOT/opt/serverless/handlers/basehandler.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def get_value(self, key, default = None):
return self.payload[key]

def get_input_dir(self):
return f"{self.INPUT_DIR}{self.request_id}/"
return f"{self.INPUT_DIR}"

def get_output_dir(self):
return f"{self.OUTPUT_DIR}"
Expand Down
70 changes: 70 additions & 0 deletions build/COPY_ROOT/opt/serverless/handlers/image2image.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
from handlers.basehandler import BaseHandler
import random
import time


"""
Handler classes are generally bound to a specific workflow file.
To modify values we have to be confident in the json structure.
One exception - RawWorkflow will send payload['workflow_json'] to the ComfyUI API - TODO
"""

class Image2Image(BaseHandler):

WORKFLOW_JSON = "/opt/serverless/workflows/image2image.json"

def __init__(self, payload):
super().__init__(payload, self.WORKFLOW_JSON)
self.apply_modifiers()


def apply_modifiers(self):
timestr = time.strftime("%Y%m%d-%H%M%S")
self.prompt["prompt"]["3"]["inputs"]["seed"] = self.get_value(
"seed",
random.randint(0,2**32))
self.prompt["prompt"]["3"]["inputs"]["steps"] = self.get_value(
"steps",
20)
self.prompt["prompt"]["3"]["inputs"]["sampler_name"] = self.get_value(
"sampler_name",
"dpmpp_2m")
self.prompt["prompt"]["3"]["inputs"]["scheduler"] = self.get_value(
"scheduler",
"normal")
self.prompt["prompt"]["3"]["inputs"]["denoise"] = self.get_value(
"denoise",
0.8700000000000001)
self.prompt["prompt"]["6"]["inputs"]["text"] = self.get_value(
"include_text",
"")
self.prompt["prompt"]["7"]["inputs"]["text"] = self.get_value(
"exclude_text",
"")
self.prompt["prompt"]["9"]["inputs"]["filename_prefix"] = f"{self.request_id}/img-{timestr}"
self.prompt["prompt"]["10"]["inputs"]["image"] = self.get_value(
"input_image",
"")
self.prompt["prompt"]["14"]["inputs"]["ckpt_name"] = self.get_value(
"ckpt_name",
"v1-5-pruned-emaonly.ckpt")


"""
Example Request Body:
{
"input": {
"handler": "Image2Image",
"aws_bucket_name": "ai-dock",
"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",
"denoise": 0.87,
"input_image": "https://raw.githubusercontent.com/comfyanonymous/ComfyUI/master/input/example.png"
}
}
"""

124 changes: 124 additions & 0 deletions build/COPY_ROOT/opt/serverless/handlers/rawworkflow.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
from handlers.basehandler import BaseHandler
import random


"""
Handler classes are generally bound to a specific workflow file.
To modify values we have to be confident in the json structure.
One exception - RawWorkflow will send payload['workflow_json'] to the ComfyUI API - TODO
"""

class RawWorkflow(BaseHandler):

WORKFLOW_JSON = None

def __init__(self, payload):
super().__init__(payload, self.WORKFLOW_JSON)
self.apply_modifiers()


def apply_modifiers(self):
# TODO - Parse the workflow for image URLs, download and update the prompt
pass



"""
Example Request Body:
{
"input": {
"handler": "RawWorkflow",
"aws_bucket_name": "ai-dock",
"workflow_json": {
"3": {
"inputs": {
"seed": 12345,
"steps": 20,
"cfg": 8,
"sampler_name": "euler",
"scheduler": "normal",
"denoise": 1,
"model": [
"4",
0
],
"positive": [
"6",
0
],
"negative": [
"7",
0
],
"latent_image": [
"5",
0
]
},
"class_type": "KSampler"
},
"4": {
"inputs": {
"ckpt_name": "v1-5-pruned-emaonly.ckpt"
},
"class_type": "CheckpointLoaderSimple"
},
"5": {
"inputs": {
"width": 1024,
"height": 1024,
"batch_size": 1
},
"class_type": "EmptyLatentImage"
},
"6": {
"inputs": {
"text": "a penguin chasing a polar bear",
"clip": [
"4",
1
]
},
"class_type": "CLIPTextEncode"
},
"7": {
"inputs": {
"text": "text, watermark",
"clip": [
"4",
1
]
},
"class_type": "CLIPTextEncode"
},
"8": {
"inputs": {
"samples": [
"3",
0
],
"vae": [
"4",
2
]
},
"class_type": "VAEDecode"
},
"9": {
"inputs": {
"filename_prefix": "image",
"images": [
"8",
0
]
},
"class_type": "SaveImage"
}
}
}
}
"""

43 changes: 41 additions & 2 deletions build/COPY_ROOT/opt/serverless/handlers/text2image.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from handlers.basehandler import BaseHandler
import random
import time


"""
Expand All @@ -19,21 +20,59 @@ def __init__(self, payload):


def apply_modifiers(self):
timestr = time.strftime("%Y%m%d-%H%M%S")
self.prompt["prompt"]["3"]["inputs"]["seed"] = self.get_value(
"seed",
random.randint(0,2**32))
self.prompt["prompt"]["3"]["inputs"]["steps"] = self.get_value(
"steps",
20)
self.prompt["prompt"]["3"]["inputs"]["sampler_name"] = self.get_value(
"sampler_name",
"euler")
self.prompt["prompt"]["3"]["inputs"]["scheduler"] = self.get_value(
"scheduler",
"normal")
self.prompt["prompt"]["4"]["inputs"]["ckpt_name"] = self.get_value(
"ckpt_name",
"v1-5-pruned-emaonly.ckpt")
self.prompt["prompt"]["5"]["inputs"]["width"] = self.get_value(
"width",
512)
self.prompt["prompt"]["5"]["inputs"]["height"] = self.get_value(
"height",
512)
self.prompt["prompt"]["5"]["inputs"]["batch_size"] = self.get_value(
"batch_size",
1)
self.prompt["prompt"]["6"]["inputs"]["text"] = self.get_value(
"include_text",
"")
self.prompt["prompt"]["7"]["inputs"]["text"] = self.get_value(
"exclude_text",
"")
self.prompt["prompt"]["9"]["inputs"]["filename_prefix"] = f"{self.request_id}/image"
self.prompt["prompt"]["9"]["inputs"]["filename_prefix"] = f"{self.request_id}/img-{timestr}"




"""
Example Request Body:
{
"input": {
"handler": "Text2Image",
"aws_bucket_name": "ai-dock",
"steps": 20,
"ckpt_name": "v1-5-pruned-emaonly.ckpt",
"sampler_name": "euler",
"scheduler": "normal",
"include_text": "A scuba diver exploring an ancient shipwreck",
"exclude_text": "steel, modern, cartoon",
"width": 512,
"height": 512,
"batch_size": 1
}
}
"""

2 changes: 0 additions & 2 deletions build/COPY_ROOT/opt/serverless/providers/runpod/worker.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import sys
sys.path.append('/opt/serverless')
from pydoc import locate
import handlers
import runpod

def get_handler(payload):
Expand All @@ -11,7 +10,6 @@ def get_handler(payload):
handler_class = locate(f"handlers.{m_name}.{c_name}")
handler = handler_class(payload)
except:
raise
raise IndexError(f"Handler ({c_name}) not found")

return handler
Expand Down
98 changes: 98 additions & 0 deletions build/COPY_ROOT/opt/serverless/workflows/image2image.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
{
"3": {
"inputs": {
"seed": 12345,
"steps": 20,
"cfg": 8,
"sampler_name": "dpmpp_2m",
"scheduler": "normal",
"denoise": 0.8700000000000001,
"model": [
"14",
0
],
"positive": [
"6",
0
],
"negative": [
"7",
0
],
"latent_image": [
"12",
0
]
},
"class_type": "KSampler"
},
"6": {
"inputs": {
"text": "pos prompt",
"clip": [
"14",
1
]
},
"class_type": "CLIPTextEncode"
},
"7": {
"inputs": {
"text": "neg prompt",
"clip": [
"14",
1
]
},
"class_type": "CLIPTextEncode"
},
"8": {
"inputs": {
"samples": [
"3",
0
],
"vae": [
"14",
2
]
},
"class_type": "VAEDecode"
},
"9": {
"inputs": {
"filename_prefix": "image",
"images": [
"8",
0
]
},
"class_type": "SaveImage"
},
"10": {
"inputs": {
"image": "example.png",
"choose file to upload": "image"
},
"class_type": "LoadImage"
},
"12": {
"inputs": {
"pixels": [
"10",
0
],
"vae": [
"14",
2
]
},
"class_type": "VAEEncode"
},
"14": {
"inputs": {
"ckpt_name": "v1-5-pruned-emaonly.ckpt"
},
"class_type": "CheckpointLoaderSimple"
}
}

0 comments on commit 7c9b02a

Please sign in to comment.