Update app.py

app.py

@@ -34,13 +34,13 @@ device = "cuda"
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
-mean = torch.load(f"{models_path}/files/mean.pt").bfloat16()#.to(device)
-std = torch.load(f"{models_path}/files/std.pt").bfloat16()#.to(device)
-v = torch.load(f"{models_path}/files/V.pt").bfloat16()#.to(device)
-proj = torch.load(f"{models_path}/files/proj_1000pc.pt").bfloat16()#.to(device)
+mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+std = torch.load(f"{models_path}/files/std.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bfloat16().to(device)
+proj = torch.load(f"{models_path}/files/proj_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 df = torch.load(f"{models_path}/files/identity_df.pt")
 weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
-pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt")#.bfloat16()#.to(device)
+pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu'))#.bfloat16()#.to(device)
 
 
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)