Update rgb2x/gradio_demo_rgb2x.py

rgb2x/gradio_demo_rgb2x.py

@@ -38,16 +38,21 @@ def generate(
     generator = torch.Generator(device="cuda").manual_seed(seed)
     photo_name = photo.name
     if photo_name.endswith(".exr"):
-        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]
+        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
     radio = old_height / old_width
     max_side = 1000
-
     if old_height > old_width:
         new_height = max_side
         new_width = int(new_height / radio)
@@ -55,10 +60,11 @@ def generate(
         new_width = max_side
         new_height = int(new_width * radio)
 
-    new_width = new_width // 8 * 8
-    new_height = new_height // 8 * 8
+    if new_width % 8 != 0 or new_height % 8 != 0:
+        new_width = new_width // 8 * 8
+        new_height = new_height // 8 * 8
 
-    photo_resized = torchvision.transforms.Resize((new_height, new_width))(photo_tensor)
+    photo = torchvision.transforms.Resize((new_height, new_width))(photo)
 
     required_aovs = ["albedo", "normal", "roughness", "metallic", "irradiance"]
     prompts = {
@@ -69,37 +75,36 @@ def generate(
         "irradiance": "Irradiance (diffuse lighting)",
     }
 
-    return_list: list[Image.Image] = []
-
+    return_list = []
     for i in range(num_samples):
         for aov_name in required_aovs:
             prompt = prompts[aov_name]
-            result = pipe(
+            generated_image = pipe(
                 prompt=prompt,
-                photo=photo_resized,
+                photo=photo,
                 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)
 
-    # Also return the input image at the end
-    input_image_pil = torchvision.transforms.ToPILImage()(photo_tensor.cpu())
-    return_list.append(input_image_pil)
+            generated_image = (generated_image, f"Generated {aov_name} {i}")
+            return_list.append(generated_image)
 
+    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"])
@@ -129,6 +134,7 @@ with gr.Blocks() as demo:
                     value=1,
                 )
 
+        # Output side
         with gr.Column():
             gr.Markdown("### Output Gallery")
             result_gallery = gr.Gallery(
@@ -162,4 +168,4 @@ with gr.Blocks() as demo:
 
 
 if __name__ == "__main__":
-    demo.launch(debug=False, share=False, show_api=False)
+    demo.launch(debug=False, share=False, show_api=False)