Refactor image post-processing in gradio_demo_rgb2x.py

rgb2x/gradio_demo_rgb2x.py

@@ -1,11 +1,11 @@
+from typing import cast
 import spaces
-import numpy as np
 import os
-from typing import cast
 import gradio as gr
 from PIL import Image
 import torch
 import torchvision
+
 from diffusers import DDIMScheduler
 from load_image import load_exr_image, load_ldr_image
 from pipeline_rgb2x import StableDiffusionAOVMatEstPipeline
@@ -36,15 +36,15 @@ def generate(
     num_samples: int,
 ) -> list[Image.Image]:
     generator = torch.Generator(device="cuda").manual_seed(seed)
-
-    if photo.name.endswith(".exr"):
-        photo = load_exr_image(photo.name, tonemaping=True, clamp=True).to("cuda")
+    photo_name = photo.name
+    if photo_name.endswith(".exr"):
+        photo = load_exr_image(photo_name, tonemaping=True, clamp=True).to("cuda")
     elif (
-        photo.name.endswith(".png")
-        or photo.name.endswith(".jpg")
-        or photo.name.endswith(".jpeg")
+        photo_name.endswith(".png")
+        or photo_name.endswith(".jpg")
+        or photo_name.endswith(".jpeg")
     ):
-        photo = load_ldr_image(photo.name, from_srgb=True).to("cuda")
+        photo = load_ldr_image(photo_name, from_srgb=True).to("cuda")
 
     # Check if the width and height are multiples of 8. If not, crop it using torchvision.transforms.CenterCrop
     old_height = photo.shape[1]
@@ -96,10 +96,7 @@ def generate(
             generated_image = (generated_image, f"Generated {aov_name} {i}")
             return_list.append(generated_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_list.append((photo_name, "Input Image"))
     return return_list
 
 
@@ -149,9 +146,9 @@ with gr.Blocks() as demo:
             examples = gr.Examples(
                 examples=[
                     [
-                        "rgb2x/example/Castlereagh_corridor_photo.png", # Photo
-                        0, # Seed
-                        50, # Inference Step
+                        "rgb2x/example/Castlereagh_corridor_photo.png",  # Photo
+                        0,  # Seed
+                        50,  # Inference Step
                         1,  # Samples
                     ]
                 ],