Update app.py

app.py

@@ -21,12 +21,12 @@ def load_img(filename):
     img_tensor = pil_to_tensor(img)
     return img_tensor
 
-def process_img(image, filename):
+def process_img(image, UHD_LL_model):
     # Select the correct weight file based on the image filename
 
     # filename = image.name.split("/")[-1]
     # if filename in image_to_weights:
-    model_path = image_to_weights[1]
+    model_path = image_to_weights[0] if UHD_LL_model else image_to_weights[1]
     checkpoints = torch.load(model_path, map_location=device)
     model.load_state_dict(checkpoints['params'])
     model.to(device)
@@ -59,7 +59,9 @@ Due to the GPU memory limitations, the app might crash if you feed a high-resolu
 examples = [
             ['images/low00772.png'], 
             ['images/low00723.png'], 
-            ['images/low00748.png'], 
+            ['images/425_UHD_LL.JPG'],
+            ['images/1778_UHD_LL.JPG'],
+            ['images/1791_UHD_LL.JPG']
            ]
 
 css = """
@@ -74,9 +76,7 @@ css = """
 
 demo = gr.Interface(
     fn = process_img,
-    inputs = [
-            gr.Image(type = 'pil', label = 'input')
-    ],
+    inputs = [gr.Image(type = 'pil', label = 'input'), 'checkbox'],
     outputs = [gr.Image(type='pil', label = 'output')],
     title = title,
     description = description,