fix: Fixed Gradio app

app.py

@@ -9,46 +9,33 @@ from torchvision.transforms.functional import InterpolationMode
 from holocron import models
 
 
-def main(args):
-
-    model = models.rexnet1_3x(pretrained=True).eval()
-
-    preprocessor = Compose([
-        Resize(model.default_cfg['input_shape'][1:], interpolation=InterpolationMode.BILINEAR),
-        PILToTensor(),
-        ConvertImageDtype(torch.float32),
-        Normalize(model.default_cfg['mean'], model.default_cfg['std'])
-    ])
-
-    def predict(input):
-        input = Image.fromarray(input.astype('uint8'), 'RGB')
-        input = preprocessor(input)
-        with torch.inference_mode():
-            prediction = torch.nn.functional.softmax(model(input.unsqueeze(0))[0], dim=0)
-        return {class_name: float(conf) for class_name, conf in zip(model.default_cfg['classes'], prediction)}
-
-    image = gr.inputs.Image()
-    outputs = gr.outputs.Label(num_top_classes=3)
-
-    interface = gr.Interface(
-        fn=predict,
-        inputs=[image],
-        outputs=outputs,
-        title="Holocron: image classification demo",
-        article=("<p style='text-align: center'><a href='https://github.com/frgfm/Holocron'>" "Github Repo</a> | "
-                 "<a href='https://frgfm.github.io/Holocron/'>Documentation</a></p>"),
-        live=True,
-        theme="huggingface",
-        layout="horizontal",
-    )
-
-    interface.launch(server_port=args.port, show_error=True)
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Holocron image classification demo',
-                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    parser.add_argument("--port", type=int, default=8001, help="Port on which the webserver will be run")
-    args = parser.parse_args()
-
-    main(args)
+model = models.rexnet1_3x(pretrained=True).eval()
+
+preprocessor = Compose([
+    Resize(model.default_cfg['input_shape'][1:], interpolation=InterpolationMode.BILINEAR),
+    PILToTensor(),
+    ConvertImageDtype(torch.float32),
+    Normalize(model.default_cfg['mean'], model.default_cfg['std'])
+])
+
+def predict(img):
+    img = Image.fromarray(img.astype('uint8'), 'RGB')
+    img = preprocessor(img)
+    with torch.inference_mode():
+        prediction = torch.nn.functional.softmax(model(img.unsqueeze(0))[0], dim=0)
+    return {class_name: float(conf) for class_name, conf in zip(model.default_cfg['classes'], prediction)}
+
+image = gr.inputs.Image()
+outputs = gr.outputs.Label(num_top_classes=3)
+
+gr.Interface(
+    fn=predict,
+    inputs=[image],
+    outputs=outputs,
+    title="Holocron: image classification demo",
+    article=("<p style='text-align: center'><a href='https://github.com/frgfm/Holocron'>" "Github Repo</a> | "
+             "<a href='https://frgfm.github.io/Holocron/'>Documentation</a></p>"),
+    live=True,
+    theme="huggingface",
+    layout="horizontal",
+).launch(enable_queue=True,cache_examples=True)