Refactor image post-processing in gradio_demo_rgb2x.py

rgb2x/gradio_demo_rgb2x.py

@@ -1,4 +1,5 @@
 import spaces
+import numpy as np
 import os
 from typing import cast
 import gradio as gr
@@ -94,7 +95,11 @@ def generate(
 
             generated_image = (generated_image, f"Generated {aov_name} {i}")
             return_list.append(generated_image)
-    return_list.append((photo, "Input Image"))
+
+    def post_process_image(img, **kwargs):
+        return (img.cpu().numpy().transpose(1, 2, 0), kwargs.get("label", "Image"))
+
+    return_list.append(post_process_image(photo, label="Input Image"))
     return return_list