File size: 1,138 Bytes
ac984b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
import runpod
import base64
from io import BytesIO
from PIL import Image
from pipeline import generate
def decode_to_image_obj(base64_string):
return Image.open(BytesIO(base64.b64decode(base64_string)))
def process_input(input):
"""
Execute the application code
"""
max_new_tokens = input['max_new_tokens']
category = input['category']
base64_string = input['image']
image = decode_to_image_obj(base64_string)
result = generate(decode_to_image_obj(image), category, max_new_tokens)
result = "This is a placeholder result."
return {
"result": result
}
# ---------------------------------------------------------------------------- #
# RunPod Handler #
# ---------------------------------------------------------------------------- #
def handler(event):
"""
This is the handler function that will be called by RunPod serverless.
"""
return process_input(event['input'])
if __name__ == '__main__':
print("Starting RunPod serverless worker.")
runpod.serverless.start({'handler': handler})
|