From 0a26654cccca8501bddf4e026d18cfee6c9891b2 Mon Sep 17 00:00:00 2001
From: Rick Staa <rick.staa@outlook.com>
Date: Thu, 13 Jun 2024 22:38:01 +0200
Subject: [PATCH] fix(api): fix incorrect I2I OpenAPI gateway spec (#105)

This commit ensures that the right format is set in the OpenAPI spec for
when the I2I response is coming from the gateway. This is a little bit
hacky since we don't yet generate the OpenAPI spec where it is
implemented (i.e. in go-livepeer).
---
 runner/app/routes/image_to_video.py |  8 ++++++--
 runner/gen_openapi.py               | 17 ++++++++++++++---
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/runner/app/routes/image_to_video.py b/runner/app/routes/image_to_video.py
index 5d66aa74..cf1d7ca5 100644
--- a/runner/app/routes/image_to_video.py
+++ b/runner/app/routes/image_to_video.py
@@ -82,7 +82,7 @@ async def image_to_video(
             motion_bucket_id=motion_bucket_id,
             noise_aug_strength=noise_aug_strength,
             safety_check=safety_check,
-            seed=seed
+            seed=seed,
         )
     except Exception as e:
         logger.error(f"ImageToVideoPipeline error: {e}")
@@ -95,7 +95,11 @@ async def image_to_video(
     for frames in batch_frames:
         output_frames.append(
             [
-                {"url": image_to_data_url(frame), "seed": seed, "nsfw": has_nsfw_concept[0]}
+                {
+                    "url": image_to_data_url(frame),
+                    "seed": seed,
+                    "nsfw": has_nsfw_concept[0],
+                }
                 for frame in frames
             ]
         )
diff --git a/runner/gen_openapi.py b/runner/gen_openapi.py
index c6dddec0..dcd21eea 100644
--- a/runner/gen_openapi.py
+++ b/runner/gen_openapi.py
@@ -1,6 +1,7 @@
 import argparse
 import json
 import os
+import copy
 
 import yaml
 from app.main import app, use_route_names_as_operation_ids
@@ -21,9 +22,11 @@ def translate_to_gateway(openapi):
     entrypoint created by the https://github.com/livepeer/go-livepeer package.
 
     .. note::
-       Differences between 'runner' and 'gateway' entrypoints:
-       - 'health' endpoint is removed.
-       - 'model_id' is enforced in all endpoints.
+        Differences between 'runner' and 'gateway' entrypoints:
+        - 'health' endpoint is removed.
+        - 'model_id' is enforced in all endpoints.
+        - 'VideoResponse' schema is updated to match the Gateway's transcoded mp4
+            response.
 
     Args:
         openapi (dict): The OpenAPI schema to be translated.
@@ -49,6 +52,14 @@ def translate_to_gateway(openapi):
                         if "model_id" in schema["properties"]:
                             schema["required"].append("model_id")
 
+    # Update the 'VideoResponse' schema to match the Gateway's response.
+    # NOTE: This is necessary because the Gateway transcodes the runner's response and
+    # returns an mp4 file.
+    openapi["components"]["schemas"]["VideoResponse"] = copy.deepcopy(
+        openapi["components"]["schemas"]["ImageResponse"]
+    )
+    openapi["components"]["schemas"]["VideoResponse"]["title"] = "VideoResponse"
+
     return openapi