Update rgb2x/gradio_demo_rgb2x.py

rgb2x/gradio_demo_rgb2x.py

@@ -38,21 +38,16 @@ def generate(
     generator = torch.Generator(device="cuda").manual_seed(seed)
     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 = 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]
-    old_width = photo.shape[2]
-    new_height = old_height
-    new_width = old_width
+        photo_tensor = load_exr_image(photo_name, tonemaping=True, clamp=True).to("cuda")
+    else:
+        photo_tensor = load_ldr_image(photo_name, from_srgb=True).to("cuda")
+
+    # Resize to multiple of 8
+    old_height = photo_tensor.shape[1]
+    old_width = photo_tensor.shape[2]
     radio = old_height / old_width
     max_side = 1000
+
     if old_height > old_width:
         new_height = max_side
         new_width = int(new_height / radio)
@@ -60,11 +55,10 @@ def generate(
         new_width = max_side
         new_height = int(new_width * radio)
 
-    if new_width % 8 != 0 or new_height % 8 != 0:
-        new_width = new_width // 8 * 8
-        new_height = new_height // 8 * 8
+    new_width = new_width // 8 * 8
+    new_height = new_height // 8 * 8
 
-    photo = torchvision.transforms.Resize((new_height, new_width))(photo)
+    photo_resized = torchvision.transforms.Resize((new_height, new_width))(photo_tensor)
 
     required_aovs = ["albedo", "normal", "roughness", "metallic", "irradiance"]
     prompts = {
@@ -75,36 +69,37 @@ def generate(
         "irradiance": "Irradiance (diffuse lighting)",
     }
 
-    return_list = []
+    return_list: list[Image.Image] = []
+
     for i in range(num_samples):
         for aov_name in required_aovs:
             prompt = prompts[aov_name]
-            generated_image = pipe(
+            result = pipe(
                 prompt=prompt,
-                photo=photo,
+                photo=photo_resized,
                 num_inference_steps=inference_step,
                 height=new_height,
                 width=new_width,
                 generator=generator,
                 required_aovs=[aov_name],
-            ).images[0][0]  # type: ignore
-
-            generated_image = torchvision.transforms.Resize((old_height, old_width))(
-                generated_image
             )
+            image_tensor = result.images[0][0]  # type: ignore
+            image_tensor = torchvision.transforms.Resize((old_height, old_width))(image_tensor)
+            image_pil = torchvision.transforms.ToPILImage()(image_tensor.cpu())
+            return_list.append(image_pil)
 
-            generated_image = (generated_image, f"Generated {aov_name} {i}")
-            return_list.append(generated_image)
+    # Also return the input image at the end
+    input_image_pil = torchvision.transforms.ToPILImage()(photo_tensor.cpu())
+    return_list.append(input_image_pil)
 
-    return_list.append((photo_name, "Input Image"))
     return return_list
 
 
 with gr.Blocks() as demo:
     with gr.Row():
         gr.Markdown("## Model RGB -> X (Realistic image -> Intrinsic channels)")
+
     with gr.Row():
-        # Input side
         with gr.Column():
             gr.Markdown("### Given Image")
             photo = gr.File(label="Photo", file_types=[".exr", ".png", ".jpg"])
@@ -134,7 +129,6 @@ with gr.Blocks() as demo:
                     value=1,
                 )
 
-        # Output side
         with gr.Column():
             gr.Markdown("### Output Gallery")
             result_gallery = gr.Gallery(