Update app.py

app.py

@@ -31,9 +31,9 @@ parser.add_argument("--temperature", type=float, default=1.0, help="temperature
 args = parser.parse_args()
 
 
-@spaces.GPU
+# @spaces.GPU
 def load_model():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cuda" #if torch.cuda.is_available() else "cpu"
     config = demo_util.get_config("configs/titok_l32.yaml")
     print(config)
     titok_tokenizer = demo_util.get_titok_tokenizer(config)
@@ -51,9 +51,9 @@ titok_tokenizer, titok_generator = load_model()
 def demo_infer(
                guidance_scale, randomize_temperature, num_sample_steps,
                class_label, seed):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    tokenizer = titok_tokenizer.to(device)
-    generator = titok_generator.to(device)
+    # device = "cuda" if torch.cuda.is_available() else "cpu"
+    tokenizer = titok_tokenizer #.to(device)
+    generator = titok_generator #.to(device)
     n = 4
     class_labels = [class_label for _ in range(n)]
     torch.manual_seed(seed)
@@ -83,7 +83,7 @@ with gr.Blocks() as demo:
                     with gr.Row():
                         i1k_class = gr.Dropdown(
                             list(imagenet_idx2classname.values()),
-                            value='macaw',
+                            value='Chihuahua',
                             type="index", label='ImageNet-1K Class'
                         )
                     guidance_scale = gr.Slider(minimum=1, maximum=25, step=0.1, value=3.5, label='Classifier-free Guidance Scale')