Update app.py

app.py

@@ -21,8 +21,6 @@ from lora_w2w import LoRAw2w
 from huggingface_hub import snapshot_download
 import spaces
 
-
-
 global device
 global generator 
 global unet
@@ -36,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").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)
 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")#.bfloat16()#.to(device)
 
 
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
@@ -51,7 +49,10 @@ def sample_model():
     global unet
     del unet
     global network
-
+    mean.to(device)
+    std.to(device)
+    v.to(device)
+    proj.to(device)
     unet, _, _, _, _ = load_models(device)
     network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)