Update app.py

app.py

@@ -20,11 +20,10 @@ class dotdict(dict):
 
 
 class InferRunner:
-    @spaces.GPU()
     def __init__(self, device):
         vae_config = json.load(open("ckpts/ldm/vae_config.json"))
         self.vae = AutoencoderKL(**vae_config).to(device)
-        vae_weights = torch.load("ckpts/ldm/pytorch_model_vae.bin", map_location=device)
+        vae_weights = torch.load("ckpts/ldm/pytorch_model_vae.bin")
         self.vae.load_state_dict(vae_weights)
 
         train_args = dotdict(json.loads(open("ckpts/pico_model/summary.jsonl").readlines()[0]))
@@ -37,7 +36,8 @@ class InferRunner:
         ).eval().to(device)
         self.scheduler = DDPMScheduler.from_pretrained(train_args.scheduler_name, subfolder="scheduler")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+#device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda"
 runner = InferRunner(device)
 event_list = get_event()