Update app.py

app.py

@@ -1,27 +1,26 @@
 import os
-# os.system("pip uninstall -y gradio") 
-# #os.system('pip install gradio==3.43.1')
-
+os.system("pip uninstall -y gradio") 
+os.system('pip install gradio==3.43.1')
 import torch
 import torchvision
 import torchvision.transforms as transforms
 from torch.utils.data import Dataset, DataLoader
 import gradio as gr
 import sys
-import os 
 import tqdm
 sys.path.append(os.path.abspath(os.path.join("", "..")))
-import torch
 import gc
 import warnings
 warnings.filterwarnings("ignore")
 from PIL import Image
-from utils import load_models, save_model_w2w, save_model_for_diffusers
+import numpy as np
+from utils import load_models
 from editing import get_direction, debias
 from sampling import sample_weights
 from lora_w2w import LoRAw2w
 from huggingface_hub import snapshot_download
-import numpy as np
+import spaces
+
 
 
 global device
@@ -32,11 +31,9 @@ global text_encoder
 global tokenizer
 global noise_scheduler
 global network
-global original_image
 device = "cuda:0"
 generator = torch.Generator(device=device)
-from gradio_imageslider import ImageSlider
-import spaces
+
 
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
@@ -125,12 +122,9 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     global pointy
     global wavy
     global large
-    global original_image
-   
     
     original_weights = network.proj.clone()
     
-    
     #pad to same number of PCs
     pcs_original = original_weights.shape[1]
     pcs_edits = young.shape[1]
@@ -141,7 +135,7 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     large_pad = torch.cat((large, padding), 1)
     
 
-    edited_weights = original_weights+a1*1e6*young_pad+a2*1e6*pointy_pad+a3*1e6*wavy_pad+a4*8e5*large_pad
+    edited_weights = original_weights+a1*1e6*young_pad+a2*1e6*pointy_pad+a3*1e6*wavy_pad+a4*2e6*large_pad
    
     generator = generator.manual_seed(seed)
     latents = torch.randn(
@@ -197,22 +191,19 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     #reset weights back to original 
     network.proj = torch.nn.Parameter(original_weights)
     network.reset()
-   
-    return (original_image, image)
+
+    return image
     
 def sample_then_run():
-    global original_image
     sample_model()    
     prompt = "sks person"
     negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
     seed = 5
     cfg = 3.0
     steps = 50
-    original_image = inference( prompt, negative_prompt, cfg, steps, seed)
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
-
-    
-    return (original_image, original_image), "model.pt"
+    return image, "model.pt"
 
 
 global young
@@ -275,14 +266,10 @@ class CustomImageDataset(Dataset):
             image = self.transform(image)
         return image
     
-def invert(dict, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
+def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
     global unet
     del unet
     global network
-
-    image = dict["background"].convert("RGB").resize((512, 512))
-    mask = dict["layers"][0].convert("RGB").resize((512, 512))
-
     unet, _, _, _, _ = load_models(device)
     
     proj = torch.zeros(1,pcs).bfloat16().to(device)
@@ -294,18 +281,13 @@ def invert(dict, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
                         train_method="xattn-strict"
                     ).to(device, torch.bfloat16)
 
-
-    
-
-
     ### load mask
     mask = transforms.Resize((64,64), interpolation=transforms.InterpolationMode.BILINEAR)(mask)
     mask = torchvision.transforms.functional.pil_to_tensor(mask).unsqueeze(0).to(device).bfloat16()[:,0,:,:].unsqueeze(1)
     ### check if an actual mask was draw, otherwise mask is just all ones
     if torch.sum(mask) == 0:
         mask = torch.ones((1,1,64,64)).to(device).bfloat16()
-
-
+        
     ### single image dataset
     image_transforms = transforms.Compose([transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
                                                 transforms.RandomCrop(512),
@@ -313,11 +295,9 @@ def invert(dict, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
                                                 transforms.Normalize([0.5], [0.5])])
 
 
-    train_dataset = CustomImageDataset([image], transform=image_transforms)
+    train_dataset = CustomImageDataset(image, transform=image_transforms)
     train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True) 
 
-
-
     ### optimizer 
     optim = torch.optim.Adam(network.parameters(), lr=lr, weight_decay=weight_decay)    
 
@@ -347,40 +327,34 @@ def invert(dict, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
                 optim.step()
 
     ### return optimized network
-    
     return network
 
 
 
-
 def run_inversion(dict, pcs, epochs, weight_decay,lr):
     global network
-    global original_image
-    # init_image = dict["image"].convert("RGB").resize((512, 512))
-    # mask = dict["ma  print(dict)
-    network = invert( dict, pcs, epochs, weight_decay,lr)
+    init_image = dict["image"].convert("RGB").resize((512, 512))
+    mask = dict["mask"].convert("RGB").resize((512, 512))
+    network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
 
 
     #sample an image
     prompt = "sks person"
-    negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
+    negative_prompt = "low quality, blurry, unfinished, nudity"
     seed = 5
     cfg = 3.0
     steps = 50
-    original_image = inference( prompt, negative_prompt, cfg, steps, seed)
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
-    return (original_image, original_image), "model.pt"
+    return image, "model.pt"
+    
     
-
-
 
 def file_upload(file):
     global unet
     del unet
     global network
     global device
-    global original_image
-    
 
     
     
@@ -393,39 +367,38 @@ def file_upload(file):
 
     unet, _, _, _, _ = load_models(device)
 
-   
-    network = LoRAw2w( proj, mean, std, v[:, :10000], 
+
+    network = LoRAw2w( proj, mean, std, v[:, :pcs], 
                         unet,
                         rank=1,
                         multiplier=1.0,
                         alpha=27.0,
                         train_method="xattn-strict"
                     ).to(device, torch.bfloat16)
-   
+        
     
     prompt = "sks person"
-    negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
+    negative_prompt = "low quality, blurry, unfinished, nudity"
     seed = 5
     cfg = 3.0
     steps = 50
-    original_image = inference( prompt, negative_prompt, cfg, steps, seed)
-    return (original_image, original_image)
-    
-
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    return image
     
     
+
     
     
     
     
 intro = """
 <div style="display: flex;align-items: center;justify-content: center">
-    <h2 style="display: inline-block;margin-left: 10px;margin-top: 6px;font-weight: 500">Interpreting the Weight Space of Customized Diffusion Models (aka <b> <em>weights2weights</em></b>)</h2>
+    <h1 style="margin-left: 12px;text-align: center;margin-bottom: 7px;display: inline-block">weights2weights</h1>
+    <h3 style="display: inline-block;margin-left: 10px;margin-top: 6px;font-weight: 500">Interpreting the Weight Space of Customized Diffusion Models</h3>
 </div>
 <p style="font-size: 0.95rem;margin: 0rem;line-height: 1.2em;margin-top:1em;display: inline-block">
-    <a href="https://snap-research.github.io/weights2weights/" target="_blank">Project Page</a> | <a href="https://arxiv.org/abs/2406.09413" target="_blank">Paper</a>
+    <a href="https://snap-research.github.io/weights2weights/" target="_blank">project page</a> | <a href="https://arxiv.org/abs/2406.09413" target="_blank">paper</a>
      | 
-     <a href="https://github.com/snap-research/weights2weights" target="_blank">Code</a> |
     <a href="https://huggingface.co/spaces/Snapchat/w2w-demo?duplicate=true" target="_blank" style="
         display: inline-block;
     ">
@@ -437,115 +410,86 @@ intro = """
 
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(intro)
-    with gr.Tab("Model Editing"):
-        gr.Markdown("""
-        Click the `Sample New Model` to sample a new identity-encoding model or upload a model to get started ✨
-        """)
-        with gr.Column():
-                with gr.Row():
-                    with gr.Column():
+    
+    gr.Markdown("""<div style="text-align: justify;"> Click below to sample an identity-encoding model, or upload an image below and click \"invert\". You can also optionally draw over the face to define a mask. To use model previously downloaded from this demo see \"Uplaoding a model\" in the Advanced options""")
+    with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(source='upload', elem_id="image_upload", tool='sketch', type='pil', label="Upload image and draw to define mask",
+                                            height=512, width=512, brush_color='#00FFFF', mask_opacity=0.6)
+                  
+                    with gr.Row():
                         sample = gr.Button("🎲 Sample New Model")
-                        file_output1 = gr.File(label="Download Sampled Model", container=True, interactive=False)
-                        file_input = gr.File(label="Upload Model", container=True)
-
-
-                    with gr.Column():
-                        image_slider1 = ImageSlider(position=0.5, type="pil", height=512, width=512, label= "Reference Identity | Generated Samples by User")
-                        
-                        prompt1 = gr.Textbox(label="Prompt",
-                                                        info="Make sure to include 'sks person'" ,
-                                                        placeholder="sks person", 
-                                                        value="sks person")
-                        seed1 = gr.Number(value=5, label="Seed", precision=0, interactive=True)
-
+                        invert_button = gr.Button("⬆️ Invert")
+                with gr.Column():
+                    gallery = gr.Image(label="Image",height=512, width=512, interactive=False)
                     
+            prompt = gr.Textbox(label="Prompt",
+                                            info="Make sure to include 'sks person'" ,
+                                            placeholder="sks person", 
+                                            value="sks person")
+            
+            seed = gr.Number(value=5, label="Seed", precision=0, interactive=True)
                 
-
-                        with gr.Row():
-                            a1_1 = gr.Slider(label="- Young +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-                            a2_1 = gr.Slider(label="- Pointy Nose +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)       
-                        with gr.Row():        
-                            a3_1 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-                            a4_1 = gr.Slider(label="- Thick Eyebrows +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-
-                
-                
-                        with gr.Accordion("Advanced Options", open=False):
-                            cfg1= gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
-                            steps1 = gr.Slider(label="Inference Steps",  value=50, step=1, minimum=0, maximum=100, interactive=True)
-                            negative_prompt1 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
-                            injection_step1 = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
-                                
-                                
-                    
-
-                        submit1 = gr.Button("Generate")
-
-    with gr.Tab("Inversion"):
-        gr.Markdown("""
-        Upload an image and optionally define a mask by drawing over the face. Then click `invert` to get started ✨
-        """)
-        with gr.Row():
+            # Editing 
             with gr.Column():
-                input_image = gr.ImageEditor(elem_id="image_upload", type='pil', label="Upload image and draw to define mask", height=512, width=512, brush=gr.Brush(), layers=False) 
-                lr = gr.Number(value=1e-1, label="Learning Rate", interactive=True)
-                pcs = gr.Slider(label="# Principal Components", value=10000, step=1, minimum=1, maximum=10000, interactive=True)
-                epochs = gr.Slider(label="Epochs", value=400, step=1, minimum=1, maximum=2000, interactive=True)
-                weight_decay = gr.Number(value=1e-10, label="Weight Decay", interactive=True)
-                invert_button = gr.Button("Invert")
-                file_output2 = gr.File(label="Download Inverted Model", container=True, interactive=False)
+                with gr.Row():
+                    a1 = gr.Slider(label="- Young +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+                    a2 = gr.Slider(label="- Pointy Nose +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)       
+                with gr.Row():        
+                    a3 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+                    a4 = gr.Slider(label="- Thick Eyebrows +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
 
-            with gr.Column():
-                image_slider2 = ImageSlider(position=0.5, type="pil", height=512, width=512, label= "Reference Identity | Generated Samples by User")
+            
+            with gr.Accordion("Advanced Options", open=False):
+                with gr.Tab("Inversion"):
+                    with gr.Row():     
+                        lr = gr.Number(value=1e-1, label="Learning Rate", interactive=True)
+                        pcs = gr.Slider(label="# Principal Components", value=10000, step=1, minimum=1, maximum=10000, interactive=True)
+                    with gr.Row():
+                        epochs = gr.Slider(label="Epochs", value=800, step=1, minimum=1, maximum=2000, interactive=True)
+                        weight_decay = gr.Number(value=1e-10, label="Weight Decay", interactive=True)
+                with gr.Tab("Sampling"):
+                    with gr.Row():     
+                            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)
+                    with gr.Row():
+                            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
+                            injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
                             
-                prompt2 = gr.Textbox(label="Prompt",
-                                    info="Make sure to include 'sks person'" ,
-                                    placeholder="sks person", 
-                                    value="sks person")
-                seed2 = gr.Number(value=5, label="Seed", precision=0, interactive=True)
+                with gr.Tab("Uploading a model"):
+                    gr.Markdown("""<div style="text-align: justify;">Upload a model below downloaded from this demo.""")
 
-                        
-                    
+                    file_input = gr.File(label="Upload Model", container=True)
 
-                with gr.Row():
-                    a1_2 = gr.Slider(label="- Young +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-                    a2_2 = gr.Slider(label="- Pointy Nose +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)       
-                with gr.Row():        
-                    a3_2 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-                    a4_2 = gr.Slider(label="- Thick Eyebrows +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+            submit = gr.Button("Generate")
+     
 
-                    
-                    
-                with gr.Accordion("Advanced Options", open=False):
-                    cfg2= gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
-                    steps2 = gr.Slider(label="Inference Steps",  value=50, step=1, minimum=0, maximum=100, interactive=True)
-                    negative_prompt2 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
-                    injection_step2 = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
-                                    
-                                    
-                        
 
-                submit2 = gr.Button("Generate")
+    gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
 
+    with gr.Row():
+        file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
         
 
-        
-                    
-        
+
+
+
+    invert_button.click(fn=run_inversion,
+                    inputs=[input_image, pcs, epochs, weight_decay,lr], 
+                    outputs = [gallery, file_output])       
+    sample.click(fn=sample_then_run, outputs=[gallery, file_output])
+    submit.click(
+        fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]
+    )
+    file_input.change(fn=file_upload, inputs=file_input, outputs = input_image)
 
 
     
-                
-    sample.click(fn=sample_then_run, outputs=[image_slider1, file_output1])        
-    submit1.click(fn=edit_inference, inputs=[ prompt1, negative_prompt1, cfg1, steps1, seed1, injection_step1, a1_1, a2_1, a3_1, a4_1], outputs=image_slider1)
-    file_input.change(fn=file_upload, inputs=file_input, outputs = image_slider1)
 
 
-    invert_button.click(fn=run_inversion, inputs=[input_image, pcs, epochs, weight_decay,lr], outputs = [image_slider2, file_output2])
-    submit2.click(fn=edit_inference, inputs=[ prompt2, negative_prompt2, cfg2, steps2, seed2, injection_step2, a1_2, a2_2, a3_2, a4_2], outputs=image_slider2)
 
 
-