Make fit for ZeroGPU

app.py

@@ -197,7 +197,7 @@ with TheraApp(
         examples_per_page=12,
         squeeze_canvas=True,
         advanced_settings_can_be_half_width=False,
-        #spaces_zero_gpu_enabled=True,
+        spaces_zero_gpu_enabled=True,
 ) as demo:
     demo.queue(
         api_open=False,

requirements.txt

@@ -1,6 +1,4 @@
 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-# torch-cpu is sufficient since we only use it for data loading
---extra-index-url https://download.pytorch.org/whl/cpu
 
 chex==0.1.7
 ConfigArgParse==1.7
@@ -25,10 +23,6 @@ optax==0.2.0
 orbax-checkpoint==0.2.4
 scipy==1.10.1
 timm==0.9.6
-torch==2.0.1+cpu
-torchaudio==2.0.2+cpu
-torchmetrics==1.2.0
-torchvision==0.15.2+cpu
 tqdm==4.65.0
 transformers==4.46.3
 Pillow==10.0.0