Update app.py

app.py

@@ -23,8 +23,8 @@ import spaces
 
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
-device = "cuda"
 
+device = "cuda"
 mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 std = torch.load(f"{models_path}/files/std.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bfloat16().to(device)
@@ -34,8 +34,7 @@ weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
 pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 
 
-
-
+unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
 
 
 young = get_direction(df, "Young", pinverse, 1000, device)
@@ -79,10 +78,9 @@ thick = debias(thick, "Brown_Hair", df, pinverse, device)
 thick = debias(thick, "Pale_Skin", df, pinverse, device)
 thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
 
-def sample_model():
-    global unet
+def sample_model(unet, network):
     del unet
-    global network
+    del network
     mean.to(device)
     std.to(device)
     v.to(device)
@@ -93,13 +91,6 @@ def sample_model():
 @torch.no_grad()
 @spaces.GPU
 def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
-    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),
@@ -396,29 +387,6 @@ intro = """
 
 
 with gr.Blocks(css="style.css") as demo:
-    generator = gr.State()
-    unet = gr.State()
-    vae = gr.State()
-    text_encoder = gr.State()
-    tokenizer = gr.State()
-    noise_scheduler = gr.State()
-    network = gr.State()
-    #device = gr.State(torch.device("cuda"))
-    device = "cuda"
-    #generator = torch.Generator(device=device)
-    young = gr.State()
-    pointy = gr.State()
-    wavy = gr.State()
-    thick = gr.State()
-    unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
-
-
-
-
-
-
-
-    
     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.""")
@@ -495,7 +463,7 @@ with gr.Blocks(css="style.css") as demo:
                     outputs = [input_image, file_output])
         
             
-        sample.click(fn=sample_then_run, outputs=[input_image, file_output])
+        sample.click(fn=sample_then_run, inputs=[unet, network] outputs=[input_image, file_output])
         
         submit.click(
             fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]