Update app.py

app.py

@@ -125,7 +125,7 @@ class main():
         print("")
 
         
-        self.network = None
+        self.weights = None
     
         young = get_direction(df, "Young", pinverse, 1000, device)
         young = debias(young, "Male", df, pinverse, device)
@@ -170,11 +170,7 @@ class main():
         self.thick = thick
 
 
-    @torch.no_grad()
-    @spaces.GPU(duration=1000)
-    def sample_model(self):
-        self.unet, _, _, _, _ = load_models(self.device)
-        self.network = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)
+   
             
  
     @torch.no_grad()
@@ -184,8 +180,19 @@ class main():
         self.unet.to(device)
         self.text_encoder.to(device)
         self.vae.to(device)
-        self.network.to(device)
-        
+        self.mean.to(device)
+        self.std.to(device)
+        self.v.to(device)
+        self.proj.to(device)
+        self.weights.to(device)
+
+        network = LoRAw2w( self.weights, self.mean, self.std, self.v, 
+                    self.unet,
+                    rank=1,
+                    multiplier=1.0,
+                    alpha=27.0,
+                    train_method="xattn-strict"
+                ).to(device, torch.bfloat16)
         
 
         
@@ -213,18 +220,9 @@ class main():
         for i,t in enumerate(tqdm.tqdm(self.noise_scheduler.timesteps)):
             latent_model_input = torch.cat([latents] * 2)
             latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep=t)
-            with self.network:
-                print(latent_model_input.device)
-                print(self.unet.device)
-                print(self.text_encoder.device)
-                print(self.vae.device)
-                print(self.network.proj.device)
-                print(self.network.mean.device)
-                print(self.network.std.device)
-                print(self.network.v.device)
-                print(text_embeddings.device)
+            with network:
                 noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
-            print("after inference")
+            
             #guidance
             noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
             noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
@@ -315,16 +313,22 @@ class main():
         
         return image
 
-    @spaces.GPU
+    @torch.no_grad()
+    @spaces.GPU(duration=1000)
     def sample_then_run(self):
-        self.sample_model()    
+        self.unet = UNet2DConditionModel.from_pretrained(
+            pretrained_model_name_or_path, subfolder="unet", revision=revision
+        )
+        self.unet.to(self.device, dtype=torch.bfloat16)
+        self.weights = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)
+        
         prompt = "sks person"
         negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
         seed = 5
         cfg = 3.0
         steps = 25
-        image = self.inference( prompt, negative_prompt, cfg, steps, seed)
-        torch.save(self.network.proj, "model.pt" )
+        image = self.inference( weights, prompt, negative_prompt, cfg, steps, seed)
+        torch.save(self.weights.cpu().detach(), "model.pt" )
         return image, "model.pt"