Fix Gradio components

.gitignore

@@ -1 +1,2 @@
 __pycache__
+.gradio

ModelLoader.py

@@ -58,14 +58,17 @@ class ModelLoader:
     def load(self) -> None:
         self.model = create_model(self.opt)
         self.model.load_networks('latest')
-    def inference(self, src=''):
+    def inference(self, src='', image_pil=None):
         if self.model == None: self.load()
-        if not os.path.isfile(src):
-            raise Exception('The image %s is not found!' % src)
         # Loading
-        print('Loading the image %s' % src)
-        source = Image.open(src).convert('RGB')
-        img = self.transform(source).unsqueeze(0)
+        if isinstance(image_pil, Image.Image):
+            img = self.transform(image_pil.convert('RGB')).unsqueeze(0)
+        else:
+            if not os.path.isfile(src):
+                raise Exception('The image %s is not found!' % src)
+            print('Loading the image %s' % src)
+            source = Image.open(src).convert('RGB')
+            img = self.transform(source).unsqueeze(0)
         print(img.shape)
         # Inference
         self.model.set_input({

app.py

@@ -2,23 +2,24 @@ import gradio as gr
 from ModelLoader import ModelLoader
 import os
 
-max_img_wh = 1024  # Set the maximum image size in pixels
+max_img_wh = 4096  # Set the maximum image size in pixels
 
 def inference(image, use_gpu):
     gpu_ids = '0' if use_gpu else ''
     model = ModelLoader(gpu_ids=gpu_ids, max_img_wh=max_img_wh)
     model.load()
-    output_img = model.inference(src=image)
+    output_img = model.inference(image_pil=image)
     return output_img
 
 demo = gr.Interface(
     fn=inference,
     inputs=[
-        gr.inputs.Image(type="pil", label="Input Image"),
-        gr.inputs.Checkbox(label="Use GPU", value=True)  # Precheck the GPU checkbox
+        gr.components.Image(type='pil', label='Input Image'),
+        gr.components.Checkbox(label='Use GPU', value=True)  # Precheck the GPU checkbox
     ],
-    outputs=gr.outputs.Image(type="pil", label="Output Image"),
-    title="Model Demo",
-    description="Upload an image to see the model output."
+    outputs=gr.components.Image(type='pil', label='Output Image', format='jpeg'),
+    title='UnderWater Photo - published by Yarflam',
+    description='Upload an image to see the model output.',
+    flagging_mode='never'
 )
 demo.launch()