Update app.py

app.py

@@ -19,9 +19,13 @@ global vae
 global text_encoder
 global tokenizer
 global noise_scheduler
+global young_val
+global pointy_val
+global bags_val
 device = "cuda:0"
 generator = torch.Generator(device=device)
 
+
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
 mean = torch.load(f"{models_path}/mean.pt").bfloat16().to(device)
@@ -30,6 +34,7 @@ v = torch.load(f"{models_path}/V.pt").bfloat16().to(device)
 proj = torch.load(f"{models_path}/proj_1000pc.pt").bfloat16().to(device)
 df = torch.load(f"{models_path}/identity_df.pt")
 weight_dimensions = torch.load(f"{models_path}/weight_dimensions.pt")
+pinverse = torch.load(f"{models_path}/pinverse_1000pc.pt").bfloat16().to(device)
 
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
 global network
@@ -38,16 +43,64 @@ def sample_model():
     global unet
     del unet
     global network
+
     unet, _, _, _, _ = load_models(device)
     network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
  
-### start off with an initial model
-sample_model()
+ 
+@torch.no_grad()
+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 = generator.manual_seed(seed)
+    latents = torch.randn(
+        (1, unet.in_channels, 512 // 8, 512 // 8),
+        generator = generator,
+        device = device
+    ).bfloat16()
+   
+
+    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+
+    text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
+
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = tokenizer(
+                            [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
+                        )
+    uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]
+    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+    noise_scheduler.set_timesteps(ddim_steps) 
+    latents = latents * noise_scheduler.init_noise_sigma
+    
+    for i,t in enumerate(tqdm.tqdm(noise_scheduler.timesteps)):
+        latent_model_input = torch.cat([latents] * 2)
+        latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep=t)
+        with network:
+            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
+        #guidance
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
+    
+    latents = 1 / 0.18215 * latents
+    image = vae.decode(latents).sample
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
 
+    image = Image.fromarray((image * 255).round().astype("uint8"))
+
+    return [image] 
 
 
 @torch.no_grad()
-def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3):
+
     global device
     global generator 
     global unet
@@ -55,6 +108,15 @@ def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     global text_encoder
     global tokenizer
     global noise_scheduler
+    global young
+    global pointy
+    global bags
+    
+    original_weights = network.proj.clone()
+    
+
+    edited_weights = original_weights+a1*young+a2*pointy+a3*bags
+
     generator = generator.manual_seed(seed)
     latents = torch.randn(
         (1, unet.in_channels, 512 // 8, 512 // 8),
@@ -76,11 +138,23 @@ def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     noise_scheduler.set_timesteps(ddim_steps) 
     latents = latents * noise_scheduler.init_noise_sigma
     
+
+ 
     for i,t in enumerate(tqdm.tqdm(noise_scheduler.timesteps)):
         latent_model_input = torch.cat([latents] * 2)
         latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep=t)
+        
+        if t>start_noise:
+            pass
+        elif t<=start_noise:
+            network.proj = torch.nn.Parameter(edited_weights)
+            network.reset()
+
+
         with network:
             noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
+            
+        
         #guidance
         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
         noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
@@ -89,13 +163,61 @@ def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     latents = 1 / 0.18215 * latents
     image = vae.decode(latents).sample
     image = (image / 2 + 0.5).clamp(0, 1)
+
     image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
 
     image = Image.fromarray((image * 255).round().astype("uint8"))
 
+    #reset weights back to original 
+    network.proj = torch.nn.Parameter(original_weights)
+    network.reset()
+
     return [image] 
 
+
+
+
+def sample_then_run():
+    global young_val
+    global pointy_val
+    global bags_val
+    global young
+    global pointy
+    global bags
+
+    sample_model()
+        
+    young_val = network.proj@young[0]/(torch.norm(young)**2).item()
+    pointy_val = network.proj@pointy[0]/(torch.norm(pointy)**2).item()
+    bags_val = network.proj@bags[0]/(torch.norm(bags)**2).item()
     
+    prompt = "sks person"
+    negative_prompt = "low quality, blurry, unfinished, cartoon"
+    seed = 5
+    cfg = 3.0
+    steps = 50
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    return image
+
+
+#directions
+global young
+global pointy
+global bags
+young = get_direction(df, "Young", pinverse, 1000, device)
+young = debias(young, "Male", df, pinverse, device)
+young_max = torch.max(proj@young[0]/(torch.norm(young))**2).item()
+young_min = torch.min(proj@young[0]/(torch.norm(young))**2).item()
+
+pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
+pointy_max = torch.max(proj@pointy[0]/(torch.norm(pointy))**2).item()
+pointy_min = torch.min(proj@pointy[0]/(torch.norm(pointy))**2).item()
+
+bags = get_direction(df, "Bags_Under_Eyes", pinverse, 1000, device)
+bags_max = torch.max(proj@bags[0]/(torch.norm(bags))**2).item()
+bags_min = torch.min(proj@bags[0]/(torch.norm(bags))**2).item()
+
+
 
 css = ''
 with gr.Blocks(css=css) as demo:
@@ -103,36 +225,41 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown("Demo for the [h94/IP-Adapter-FaceID model](https://huggingface.co/h94/IP-Adapter-FaceID) - Generate AI images with your own face - Non-commercial license")
     with gr.Row():
         with gr.Column():
-            files = gr.Files(
-                        label="Upload a photo of your face to invert, or sample a new model",
-                        file_types=["image"]
-                    )
-            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=125)
-
             sample = gr.Button("Sample New Model")
+            gallery1 = gr.Gallery(label="Identity from Sampled Model")
 
-            with gr.Column(visible=False) as clear_button:
-                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+        with gr.Column():          
             prompt = gr.Textbox(label="Prompt",
-                       info="Make sure to include 'sks person'" ,
-                       placeholder="sks person", 
-                       value="sks person")
+                        info="Make sure to include 'sks person'" ,
+                        placeholder="sks person", 
+                        value="sks person")
             negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, cartoon", value="low quality, blurry, unfinished, cartoon")
-            seed = gr.Number(value=5, label="Seed", interactive=True)
+            seed = gr.Number(value=5, precision=0, label="Seed", interactive=True)
             cfg = gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
-            steps = gr.Slider(label="Inference Steps", value=50, step=1, minimum=0, maximum=100, interactive=True)
+            steps = gr.Slider(label="Inference Steps", precision=0, value=50, step=1, minimum=0, maximum=100, interactive=True)
+            injection_step = gr.Slider(label="Injection Step", precision=0, value=800, step=1, minimum=0, maximum=1000, interactive=True)
+
+
+            with gr.Row():
+                a1 = gr.Slider(label="Young", value=0, step=1, minimum=-1000000, maximum=1000000, interactive=True)
+                a2 = gr.Slider(label="Pointy Nose", value=0, step=1, minimum=-1000000, maximum=1000000, interactive=True)
+                a3 = gr.Slider(label="Undereye Bags", value=0, step=1, minimum=-1000000, maximum=1000000, interactive=True)
 
 
             submit = gr.Button("Submit")
 
         with gr.Column():
-            gallery = gr.Gallery(label="Generated Images")
+            gallery2 = gr.Gallery(label="Identity from Edited Model")
+
+
+
+        
+        sample.click(fn=sample_then_run, outputs=gallery1)
 
-        sample.click(fn=sample_model)
         
-        submit.click(fn=inference,
-                    inputs=[prompt, negative_prompt, cfg, steps, seed],
-                    outputs=gallery)
+        submit.click(fn=edit_inference,
+                    inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3],
+                    outputs=gallery2)
             
 
 
@@ -140,3 +267,8 @@ with gr.Blocks(css=css) as demo:
             
             
 demo.launch(share=True)
+
+
+
+            
+