update sdk_version to 6.6.0 and refactor Gradio app launch configuration

README.md

@@ -4,7 +4,7 @@ emoji: 🚗
 colorFrom: indigo
 colorTo: green
 sdk: gradio
-sdk_version: 5.22.0
+sdk_version: 6.6.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -5,10 +5,12 @@ import gradio as gr
 import matplotlib.cm as cm
 import numpy as np
 import plotly.graph_objects as go
+import spaces
 import torch
 import torch.nn.functional as F
 import torchdiffeq
 
+
 DESCRIPTION = """
 <div class="head">
 <div class="title">Fast LiDAR Data Generation with Rectified Flows</div>
@@ -92,6 +94,7 @@ def model_verbose(model, nfe, progress):
     return _model
 
 
+@spaces.GPU
 def generate(nfe: int, solver: str, phase: str, progress=gr.Progress()):
     model, lidar_utils, _ = torch.hub.load(config=model_dict[phase], **torch_hub_kwargs)
 
@@ -170,26 +173,7 @@ def setup_dropdown(value):
     return dropdown_solver, dropdown_nfe
 
 
-with gr.Blocks(
-    css="""
-.head {
-  text-align: center;
-  display: block;
-  font-size: var(--text-xl);
-}
-
-.title {
-  font-size: var(--text-xxl);
-  font-weight: bold;
-  margin-top: 2rem;
-}
-
-.description {
-  font-size: var(--text-lg);
-}
-    """,
-    theme=gr.themes.Ocean(),
-) as demo:
+with gr.Blocks() as demo:
     gr.HTML(DESCRIPTION)
 
     with gr.Row(variant="panel"):
@@ -222,4 +206,23 @@ with gr.Blocks(
 
 
 demo.queue()
-demo.launch()
+demo.launch(
+    css="""
+.head {
+  text-align: center;
+  display: block;
+  font-size: var(--text-xl);
+}
+
+.title {
+  font-size: var(--text-xxl);
+  font-weight: bold;
+  margin-top: 2rem;
+}
+
+.description {
+  font-size: var(--text-lg);
+}
+    """,
+    theme=gr.themes.Ocean(),
+)