Update app.py

app.py

@@ -1,8 +1,10 @@
 import os
 from tempfile import TemporaryDirectory
 
+import gradio as gr
 import numpy as np
 import pandas as pd
+import spaces
 import torch
 from huggingface_hub import Repository
 from rlgym_tools.rocket_league.misc.serialize import serialize_game_state, serialize_scoreboard, \
@@ -23,6 +25,7 @@ MODEL = torch.jit.load("vortex-ngp/vortex-ngp-avid-frog.pt", map_location=DEVICE
 MODEL.eval()
 
 
+@spaces.GPU
 @torch.inference_mode()
 def infer(model, replay_file,
           nullify_goal_difference=False,
@@ -263,60 +266,56 @@ RADIO_INFO = """
 - **Ignore ties**: Makes the model pretend every situation is an overtime (e.g. ties are impossible).
 """.strip()
 
+with TemporaryDirectory() as temp_dir:
+    with gr.Blocks() as demo:
+        gr.Markdown(DESCRIPTION)
 
-def gradio_app():
-    import gradio as gr
+        # Use gr.Column to stack components vertically
+        with gr.Column():
+            file_input = gr.File(label="Upload Replay File", type="filepath", file_types=[".replay"])
+            checkboxes = gr.Radio(label="Options", choices=RADIO_OPTIONS, type="index", value=RADIO_OPTIONS[0],
+                                  info=RADIO_INFO)
+            submit_button = gr.Button("Generate Predictions")
+            plot_output = gr.Plot(label="Predictions")
+            download_button = gr.DownloadButton("Download Predictions", visible=False)
 
-    with TemporaryDirectory() as temp_dir:
-        with gr.Blocks() as demo:
-            gr.Markdown(DESCRIPTION)
-
-            # Use gr.Column to stack components vertically
-            with gr.Column():
-                file_input = gr.File(label="Upload Replay File", type="filepath", file_types=[".replay"])
-                checkboxes = gr.Radio(label="Options", choices=RADIO_OPTIONS, type="index", value=RADIO_OPTIONS[0],
-                                      info=RADIO_INFO)
-                submit_button = gr.Button("Generate Predictions")
-                plot_output = gr.Plot(label="Predictions")
-                download_button = gr.DownloadButton("Download Predictions", visible=False)
 
+        def make_plot(replay_file, radio_option, progress=gr.Progress(track_tqdm=True)):
             # Make plot on button click
-            def make_plot(replay_file, radio_option, progress=gr.Progress(track_tqdm=True)):
-                nullify_goal_difference = radio_option == 1
-                ignore_ties = radio_option == 2
-                print(f"Processing file: {replay_file}")
-
-                replay_stem = os.path.splitext(os.path.basename(replay_file))[0]
-                postfix = ""
-                if nullify_goal_difference:
-                    postfix += "_nullify_goal_difference"
-                elif ignore_ties:
-                    postfix += "_ignore_ties"
-                preds_file = os.path.join(temp_dir, f"predictions_{replay_stem}{postfix}.csv")
-                if os.path.exists(preds_file):
-                    print(f"Predictions file already exists: {preds_file}")
-                    preds = pd.read_csv(preds_file, dtype={"Touch": str})
-                else:
-                    preds = infer(MODEL, replay_file,
-                                  nullify_goal_difference=nullify_goal_difference,
-                                  ignore_ties=ignore_ties)
-                plt = plot_plotly(preds)
-                print(f"Plot generated for file: {replay_file}")
-                preds.to_csv(preds_file)
-                if len(os.listdir(temp_dir)) > 100:
-                    # Delete least recent file
-                    oldest_file = min(os.listdir(temp_dir), key=lambda f: os.path.getctime(os.path.join(temp_dir, f)))
-                    os.remove(os.path.join(temp_dir, oldest_file))
-                return plt, gr.DownloadButton(value=preds_file, visible=True)
-
-            submit_button.click(
-                fn=make_plot,
-                inputs=[file_input, checkboxes],
-                outputs=[plot_output, download_button],
-                show_progress="full",
-            )
-        demo.launch()
-
-
-if __name__ == '__main__':
-    gradio_app()
+            nullify_goal_difference = radio_option == 1
+            ignore_ties = radio_option == 2
+            print(f"Processing file: {replay_file}")
+
+            replay_stem = os.path.splitext(os.path.basename(replay_file))[0]
+            postfix = ""
+            if nullify_goal_difference:
+                postfix += "_nullify_goal_difference"
+            elif ignore_ties:
+                postfix += "_ignore_ties"
+            preds_file = os.path.join(temp_dir, f"predictions_{replay_stem}{postfix}.csv")
+            if os.path.exists(preds_file):
+                print(f"Predictions file already exists: {preds_file}")
+                preds = pd.read_csv(preds_file, dtype={"Touch": str})
+            else:
+                preds = infer(MODEL, replay_file,
+                              nullify_goal_difference=nullify_goal_difference,
+                              ignore_ties=ignore_ties)
+            plt = plot_plotly(preds)
+            print(f"Plot generated for file: {replay_file}")
+            preds.to_csv(preds_file)
+            if len(os.listdir(temp_dir)) > 100:
+                # Delete least recent file
+                oldest_file = min(os.listdir(temp_dir), key=lambda f: os.path.getctime(os.path.join(temp_dir, f)))
+                os.remove(os.path.join(temp_dir, oldest_file))
+            return plt, gr.DownloadButton(value=preds_file, visible=True)
+
+
+        submit_button.click(
+            fn=make_plot,
+            inputs=[file_input, checkboxes],
+            outputs=[plot_output, download_button],
+            show_progress="full",
+        )
+
+demo.queue(default_concurrency_limit=None)
+demo.launch()