fix: model output

Author: Justin Merrell

infer.py

@@ -32,4 +32,9 @@ def run(model_inputs):
     # Return Errors
     # return {"error": "Error Message"}
 
-    return {"images": ("/path/to/image.png"), "seed": "1234"}
+    return [
+        {
+            "image": "/path/to/image.png",
+            "seed": "1234"
+        }
+    ]

runpod/serverless/modules/inference.py

@@ -54,8 +54,10 @@ def run(self, job):
         '''
         input_errors = self.input_validation(job['input'])
         if input_errors:
-            return {
-                "error": input_errors
-            }
+            return [
+                {
+                    "error": input_errors
+                }
+            ]
 
         return infer.run(job)

runpod/serverless/modules/job.py

@@ -64,14 +64,15 @@ def run(job):
 
     job_output = model.run(job)
 
-    if "error" in job_output:
-        return {
-            "error": job_output["error"]
-        }
-
-    if "images" in job_output:
-        object_urls = upload.upload_image(job['id'], job_output["images"])
-        job_output["images"] = object_urls
+    for index, output in enumerate(job_output):
+        if "error" in output:
+            return {
+                "error": output["error"]
+            }
+
+        if "image" in job_output:
+            object_url = upload.upload_image(job['id'], output["image"], index)
+            output["image"] = object_url
 
     job_duration = time.time() - time_job_started
     job_duration_ms = int(job_duration * 1000)

runpod/serverless/modules/upload.py

@@ -36,50 +36,45 @@
 # ---------------------------------------------------------------------------- #
 #                                 Upload Image                                 #
 # ---------------------------------------------------------------------------- #
-def upload_image(job_id, job_results):
+def upload_image(job_id, job_result, result_index=0):
     '''
     Upload image to bucket storage.
     '''
     if boto_client is None:
         # Save the output to a file
-        for index, result in enumerate(job_results):
-            output = BytesIO()
-            img = Image.open(result)
-            img.save(output, format=img.format)
-
-            os.makedirs("uploaded", exist_ok=True)
-            with open(f"uploaded/{index}.png", "wb") as file_output:
-                file_output.write(output.getvalue())
-
-        return []
-
-    object_urls = []
-    for index, result in enumerate(job_results):
         output = BytesIO()
-        img = Image.open(result)
+        img = Image.open(job_result)
         img.save(output, format=img.format)
 
-        bucket = time.strftime('%m-%y')
+        os.makedirs("uploaded", exist_ok=True)
+        with open(f"uploaded/{result_index}.png", "wb") as file_output:
+            file_output.write(output.getvalue())
 
-        # Upload to S3
-        boto_client.put_object(
-            Bucket=f'{bucket}',
-            Key=f'{job_id}/{index}.png',
-            Body=output.getvalue(),
-            ContentType="image/png"
-        )
+        return None
 
-        output.close()
+    output = BytesIO()
+    img = Image.open(job_result)
+    img.save(output, format=img.format)
+
+    bucket = time.strftime('%m-%y')
+
+    # Upload to S3
+    boto_client.put_object(
+        Bucket=f'{bucket}',
+        Key=f'{job_id}/{result_index}.png',
+        Body=output.getvalue(),
+        ContentType="image/png"
+    )
 
-        presigned_url = boto_client.generate_presigned_url(
-            'get_object',
-            Params={
-                'Bucket': f'{bucket}',
-                'Key': f'{job_id}/{index}.png'
-            }, ExpiresIn=604800)
+    output.close()
 
-        object_urls.append(presigned_url)
+    presigned_url = boto_client.generate_presigned_url(
+        'get_object',
+        Params={
+            'Bucket': f'{bucket}',
+            'Key': f'{job_id}/{result_index}.png'
+        }, ExpiresIn=604800)
 
-        log(f"Presigned URL generated: {presigned_url}")
+    log(f"Presigned URL generated: {presigned_url}")
 
-    return object_urls
+    return presigned_url

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = runpod
-version = 0.4.2
+version = 0.4.3
 description = Official Python library for RunPod API & SDK.
 long_description = file: README.md
 long_description_content_type = text/markdown