Update app.py

app.py

@@ -112,23 +112,43 @@ thick = debias(thick, "Black_Hair", df, pinverse, device)
 thick = debias(thick, "Brown_Hair", df, pinverse, device)
 thick = debias(thick, "Pale_Skin", df, pinverse, device)
 thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
-        
+
+
 
 @torch.no_grad()
-@spaces.GPU(duration=120)
-def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
-    device = self.device
-    self.unet.to(device)
-    self.text_encoder.to(device)
-    self.vae.to(device)
-    self.mean.to(device)
-    self.std.to(device)
-    self.v.to(device)
-    self.proj.to(device)
-    self.weights.to(device)
+@spaces.GPU
+def sample_then_run(net):
+    print(net)
+    # get mean and standard deviation for each principal component
+    m = torch.mean(proj, 0)
+    standev = torch.std(proj, 0)
+    
+    # sample
+    sample = torch.zeros([1, 1000]).to(device)
+    for i in range(1000):
+        sample[0, i] = torch.normal(m[i], standev[i], (1,1))
 
-    network = LoRAw2w( self.weights.bfloat16(), self.mean.bfloat16(), self.std.bfloat16(), self.v[:, :1000].bfloat16(), 
-                    self.unet,
+    net = "model_"+str(uuid.uuid4())[:4]+".pt"
+    torch.save(sample, net)
+    
+    image = prompt = "sks person"
+    negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
+    seed = 5
+    cfg = 3.0
+    steps = 25
+    image = inference(net, prompt, negative_prompt, cfg, steps, seed)
+    return net, image
+
+@torch.no_grad()
+@spaces.GPU(duration=120)
+def inference(net, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+    mean.to(device)
+    std.to(device)
+    v.to(device)
+    
+    weights = torch.load(net).to(device)
+    network = LoRAw2w(weights, mean, std, v[:, :1000], 
+                    unet,
                     rank=1,
                     multiplier=1.0,
                     alpha=27.0,
@@ -176,6 +196,8 @@ def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
         
     image = Image.fromarray((image * 255).round().astype("uint8"))
+
+    del network
         
     return image
 
@@ -270,23 +292,23 @@ def edit_inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, se
 
     return image
 
-@torch.no_grad()
-@spaces.GPU(duration=120)
-def sample_then_run(self):
-    self.unet = UNet2DConditionModel.from_pretrained(
-            "stablediffusionapi/realistic-vision-v51" , subfolder="unet", revision=None
-        )
-    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.weights.cpu().detach(), "model.pt" )
-    return image, "model.pt"
+# @torch.no_grad()
+# @spaces.GPU(duration=120)
+# def sample_then_run(self):
+#     self.unet = UNet2DConditionModel.from_pretrained(
+#             "stablediffusionapi/realistic-vision-v51" , subfolder="unet", revision=None
+#         )
+#     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.weights.cpu().detach(), "model.pt" )
+#     return image, "model.pt"
 
 
 
@@ -434,7 +456,7 @@ intro = """
     
     
 with gr.Blocks(css="style.css") as demo:
-      
+    net = gr.State()  
     gr.HTML(intro)
         
     gr.Markdown("""<div style="text-align: justify;"> In this demo, you can get an identity-encoding model by sampling or inverting. To use a model previously downloaded from this demo see \"Uploading a model\" in the Advanced Options. Next, you can generate new images from it, or edit the identity encoded in the model and generate images from the edited model. We provide detailed instructions and tips at the bottom of the page.""")
@@ -511,7 +533,7 @@ with gr.Blocks(css="style.css") as demo:
         #                 outputs = [input_image, file_output])
             
                 
-        # sample.click(fn=model.sample_then_run, outputs=[input_image, file_output])
+        sample.click(fn=sample_then_run,inputs = [net], outputs=[net, input_image])
             
         # submit.click(
         #         fn=model.edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]