Update app.py

app.py

@@ -9,6 +9,7 @@ import time
 import argparse
 import demo_util
 import os
+import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -42,8 +43,12 @@ titok_tokenizer = titok_tokenizer.to(device)
 titok_generator = titok_generator.to(device)
 
 
+@spaces.GPU
 def demo_infer(guidance_scale, randomize_temperature, num_sample_steps,
                class_label, seed):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    titok_tokenizer = titok_tokenizer.to(device)
+    titok_generator = titok_generator.to(device)
     n = 4
     class_labels = [class_label for _ in range(n)]
     torch.manual_seed(seed)