Update app.py

app.py

@@ -14,7 +14,6 @@ from PIL import Image
 import numpy as np
 from utils import load_models
 from editing import get_direction, debias
-from sampling import sample_weights
 from lora_w2w import LoRAw2w
 from huggingface_hub import snapshot_download
 import spaces
@@ -83,6 +82,31 @@ thick.value = debias(thick.value, "Brown_Hair", df, pinverse, device.value)
 thick.value = debias(thick.value, "Pale_Skin", df, pinverse, device.value)
 thick.value = debias(thick.value, "Heavy_Makeup", df, pinverse, device.value)
 
+
+@torch.no_grad()
+@spaces.GPU
+def sample_weights(unet, proj, mean, std, v, device, factor = 1.0): 
+    # get mean and standard deviation for each principal component
+    m = torch.mean(proj, 0)
+    standev = torch.std(proj, 0)
+    del proj
+    torch.cuda.empty_cache()
+    # sample
+    sample = torch.zeros([1, 1000]).to(device)
+    for i in range(1000):
+        sample[0, i] = torch.normal(m[i], factor*standev[i], (1,1))
+
+    # load weights into network
+    net = LoRAw2w( sample, mean, std, v, 
+                    unet,
+                    rank=1,
+                    multiplier=1.0,
+                    alpha=27.0,
+                    train_method="xattn-strict"
+                ).to(device, torch.bfloat16)
+        
+    return net
+
 @torch.no_grad()
 @spaces.GPU
 def sample_model():