Update app.py

app.py

@@ -21,15 +21,8 @@ import spaces
 import uuid
 
 global device
-global generator 
-global unet
-global vae 
-global text_encoder
-global tokenizer
-global noise_scheduler
-global network
 device = "cuda"
-#generator = torch.Generator(device=device)
+
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
@@ -41,29 +34,76 @@ df = torch.load(f"{models_path}/files/identity_df.pt")
 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)
+global young
+global pointy
+global wavy 
+global thick
+young = get_direction(df, "Young", pinverse, 1000, device)
+young = debias(young, "Male", df, pinverse, device)
+young = debias(young, "Pointy_Nose", df, pinverse, device)
+young = debias(young, "Wavy_Hair", df, pinverse, device)
+young = debias(young, "Chubby", df, pinverse, device)
+young = debias(young, "No_Beard", df, pinverse, device)
+young = debias(young, "Mustache", df, pinverse, device)
+    
+pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
+pointy = debias(pointy, "Young", df, pinverse, device)
+pointy = debias(pointy, "Male", df, pinverse, device)
+pointy = debias(pointy, "Wavy_Hair", df, pinverse, device)
+pointy = debias(pointy, "Chubby", df, pinverse, device)
+pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
+
+wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
+wavy = debias(wavy, "Young", df, pinverse, device)
+wavy = debias(wavy, "Male", df, pinverse, device)
+wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
+wavy = debias(wavy, "Chubby", df, pinverse, device)
+wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
+
+thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
+thick = debias(thick, "Male", df, pinverse, device)
+thick = debias(thick, "Young", df, pinverse, device)
+thick = debias(thick, "Pointy_Nose", df, pinverse, device)
+thick = debias(thick, "Wavy_Hair", df, pinverse, device)
+thick = debias(thick, "Mustache", df, pinverse, device)
+thick = debias(thick, "No_Beard", df, pinverse, device)
+thick = debias(thick, "Sideburns", df, pinverse, device)
+thick = debias(thick, "Big_Nose", df, pinverse, device)
+thick = debias(thick, "Big_Lips", df, pinverse, device)
+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)
 
-def sample_model():
-    global unet
-    del unet
-    global network
+
+
+@torch.no_grad()
+@spaces.GPU
+def sample_then_run(network, unet):
+    #load models
     mean.to(device)
     std.to(device)
     v.to(device)
     proj.to(device)
     unet, _, _, _, _ = load_models(device)
     network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
- 
+    #inference
+    prompt = "sks person"
+    negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
+    seed = 5
+    cfg = 3.0
+    steps = 25
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    torch.save(network.proj, "model.pt" )
+    #return 
+    return image, "model.pt", unet, network
+
+
+    
 @torch.no_grad()
 @spaces.GPU
-def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+def inference(network, unet, 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),
@@ -192,61 +232,10 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
 
     return image
 
-@spaces.GPU
-def sample_then_run():
-    sample_model()    
-    prompt = "sks person"
-    negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
-    seed = 5
-    cfg = 3.0
-    steps = 25
-    image = inference( prompt, negative_prompt, cfg, steps, seed)
-    torch.save(network.proj, "model.pt" )
-    return image, "model.pt"
 
-#@spaces.GPU
-def start_items():
-    print("Starting items")
-    global young
-    global pointy
-    global wavy 
-    global thick
-    young = get_direction(df, "Young", pinverse, 1000, device)
-    young = debias(young, "Male", df, pinverse, device)
-    young = debias(young, "Pointy_Nose", df, pinverse, device)
-    young = debias(young, "Wavy_Hair", df, pinverse, device)
-    young = debias(young, "Chubby", df, pinverse, device)
-    young = debias(young, "No_Beard", df, pinverse, device)
-    young = debias(young, "Mustache", df, pinverse, device)
-    
-    pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
-    pointy = debias(pointy, "Young", df, pinverse, device)
-    pointy = debias(pointy, "Male", df, pinverse, device)
-    pointy = debias(pointy, "Wavy_Hair", df, pinverse, device)
-    pointy = debias(pointy, "Chubby", df, pinverse, device)
-    pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
-    
-    wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
-    wavy = debias(wavy, "Young", df, pinverse, device)
-    wavy = debias(wavy, "Male", df, pinverse, device)
-    wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
-    wavy = debias(wavy, "Chubby", df, pinverse, device)
-    wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
-
-    thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
-    thick = debias(thick, "Male", df, pinverse, device)
-    thick = debias(thick, "Young", df, pinverse, device)
-    thick = debias(thick, "Pointy_Nose", df, pinverse, device)
-    thick = debias(thick, "Wavy_Hair", df, pinverse, device)
-    thick = debias(thick, "Mustache", df, pinverse, device)
-    thick = debias(thick, "No_Beard", df, pinverse, device)
-    thick = debias(thick, "Sideburns", df, pinverse, device)
-    thick = debias(thick, "Big_Nose", df, pinverse, device)
-    thick = debias(thick, "Big_Lips", df, pinverse, device)
-    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)
+
+
+
 
 class CustomImageDataset(Dataset):
     def __init__(self, images, transform=None):
@@ -403,8 +392,15 @@ intro = """
 
 
 with gr.Blocks(css="style.css") as demo:
-    gr.HTML(intro)
+    network = gr.State()
+    unet = gr.State()
+    vae = gr.State()
+    text_encoder = gr.State()
+    tokenizer = gr.State()
+    noise_scheduler = gr.State()
+    _, 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.""")
     with gr.Column():
             with gr.Row():
@@ -479,7 +475,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, outputs=[network, unet, input_image, file_output])
         
         submit.click(
             fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]