Update app.py

app.py

@@ -12,15 +12,14 @@ from utils import load_models, save_model_w2w, save_model_for_diffusers
 from sampling import sample_weights
 from huggingface_hub import snapshot_download
 
-global device
-global generator 
-global unet
-global vae 
-global text_encoder
-global tokenizer
-global noise_scheduler
+#global device
+#global generator 
+#global unet
+#global vae 
+#global text_encoder
+#global tokenizer
+#global noise_scheduler
 device = "cuda:0"
-generator = torch.Generator(device=device)
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
@@ -32,24 +31,23 @@ df = torch.load(f"{models_path}/identity_df.pt")
 weight_dimensions = torch.load(f"{models_path}/weight_dimensions.pt")
 
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
-global network
+network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
+#global network
+
+#def sample_model():
+#    global unet
+#    del unet
+#    global network
+#    unet, _, _, _, _ = load_models(device)
 
-def sample_model():
-   global unet
-   del unet
-   global network
-   unet, _, _, _, _ = load_models(device)
-    
-@torch.no_grad()
 def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
-    global network
-    global device
-    global generator 
-    global unet
-    global vae 
-    global text_encoder
-    global tokenizer
-    global noise_scheduler
+    #global device
+    #global generator 
+    #global unet
+    #global vae 
+    #global text_encoder
+    #global tokenizer
+    #global noise_scheduler
     generator = torch.Generator(device=device).manual_seed(seed)
     latents = torch.randn(
         (1, unet.in_channels, 512 // 8, 512 // 8),
@@ -118,7 +116,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             gallery = gr.Gallery(label="Generated Images")
 
-        sample.click(fn=sample_model)
+        #sample.click(fn=sample_model)
         
         submit.click(fn=inference,
                     inputs=[prompt, negative_prompt, cfg, steps, seed],