Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import os 
 os.system("pip uninstall -y gradio") 
 os.system('pip install gradio==3.43.1')
 
@@ -9,6 +8,7 @@ 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
@@ -112,6 +112,7 @@ def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
 
 
 
+
 @torch.no_grad()
 def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
 
@@ -129,8 +130,17 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     
     original_weights = network.proj.clone()
     
+    #pad to same number of PCs
+    pcs_original = original_weights.shape[1]
+    pcs_edits = young.shape[1]
+    padding =  torch.zeros((1,pcs_original-pcs_edits)).to(device)
+    young_pad = torch.cat((young, padding), 1)
+    pointy_pad = torch.cat((pointy, padding), 1)
+    wavy_pad = torch.cat((wavy, padding), 1)
+    large_pad = torch.cat((large, padding), 1)
+    
 
-    edited_weights = original_weights+a1*1e6*young+a2*1e6*pointy+a3*1e6*wavy+a4*2e6*large
+    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(
@@ -191,16 +201,16 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
 
 
 
-
 def sample_then_run():
     sample_model()    
     prompt = "sks person"
-    negative_prompt = "low quality, blurry, unfinished, cartoon"
+    negative_prompt = "low quality, blurry, unfinished, nudity"
     seed = 5
     cfg = 3.0
     steps = 50
     image = inference( prompt, negative_prompt, cfg, steps, seed)
-    return image
+    torch.save(network.proj, "sampled_model.pt" )
+    return image, "sampled_model.pt"
 
 
 
@@ -342,20 +352,55 @@ def run_inversion(dict, pcs, epochs, weight_decay,lr):
 
     #sample an image
     prompt = "sks person"
-    negative_prompt = "low quality, blurry, unfinished, cartoon"
+    negative_prompt = "low quality, blurry, unfinished, nudity"
     seed = 5
     cfg = 3.0
     steps = 50
     image = inference( prompt, negative_prompt, cfg, steps, seed)
-    torch.save(network.proj, "model.pt" )
-    return image, "model.pt"
+    torch.save(network.proj, "inverted_model.pt" )
+    return image, "inverted_model.pt"
 
 
 
+def file_upload(file):
+    global unet
+    del unet
+    global network
+    global device
 
+    
+    
+    proj = torch.load(file.name).to(device)
+
+    #pad to 10000 Principal components to keep everything consistent
+    pcs = proj.shape[1]
+    padding =  torch.zeros((1,10000-pcs)).to(device)
+    proj = torch.cat((proj, padding), 1)
+
+    unet, _, _, _, _ = load_models(device)
 
 
+    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"
+    seed = 5
+    cfg = 3.0
+    steps = 50
+    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    return image
+    
 
+    
+    
+    
 intro = """
 <div style="display: flex;align-items: center;justify-content: center">
     <h1 style="margin-left: 12px;text-align: center;margin-bottom: 7px;display: inline-block">weights2weights</h1>
@@ -378,39 +423,47 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Tab("Sampling Models + Editing"):
         with gr.Row():
             with gr.Column():
-                gallery1 = gr.Image(label="Identity from Sampled Model")
+                gallery1 = gr.Image(label="Identity from Sampled or Uploaded Model")
                 sample = gr.Button("Sample New Model")
-            gallery2 = gr.Image(label="Identity from Edited Model")
+                file_input = gr.File(label="Upload Model", container=True)
+                file_input.change(fn=file_upload, inputs=file_input, outputs = gallery1)
+
+
+            
                 
 
-        with gr.Row():
-            with gr.Column():          
+            with gr.Column():
+                    gallery2 = gr.Image(label="Identity from Edited Model")
                     prompt = gr.Textbox(label="Prompt",
-                                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")
-        with gr.Row():
+                                    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 = 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():                
+                        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="- placeholder for some fourth attribute +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+                        a4 = gr.Slider(label="- placeholder +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+
+                    
+                    with gr.Accordion("Advanced Options", open=False):
+                        with gr.Column():     
+                                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)
+                                injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
+                                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity", value="low quality, blurry, unfinished, nudity")
+
+                    submit = gr.Button("Generate")
+
+        file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
 
-        
-        with gr.Accordion("Advanced Options", open=False):
-            with gr.Column():     
-                    seed = gr.Number(value=5, label="Seed", precision=0, 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)
-                    injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
 
 
                     
-        submit = gr.Button("Generate")
+        
             
-        sample.click(fn=sample_then_run, outputs=gallery1)
+        sample.click(fn=sample_then_run, outputs=[gallery1, file_output])
 
         submit.click(fn=edit_inference,
                     inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4],
@@ -421,15 +474,18 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Tab("Inversion"):
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(sources='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) 
-                # input_image = gr.ImageEditor(sources='upload', elem_id="image_upload", type='pil', label="Upload image and draw to define mask",
-                #                         height=512, width=512) 
+                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) 
+            
                     
                 lr = gr.Number(value=1e-1, label="Learning Rate", interactive=True)
-                weight_decay = gr.Number(value=1e-10, label="Weight Decay", 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)
+                with gr.Accordion("Advanced Options", open=False):
+                        with gr.Column():
+                            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")
                 
@@ -439,13 +495,15 @@ with gr.Blocks(css="style.css") as demo:
                                 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", precision=0, 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)
+                with gr.Accordion("Advanced Options", open=False):
+                        with gr.Column():     
+                                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)
+                                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity", value="low quality, blurry, unfinished, nudity")
                 submit = gr.Button("Generate")
                 
-        file_output = gr.File(label="Download Model", container=False)
+        file_output = gr.File(label="Download Inverted Model", container=True, interactive=False)
 
                 
     
@@ -459,10 +517,4 @@ with gr.Blocks(css="style.css") as demo:
                     outputs=gallery)
             
         
-
-
-
-            
-            
-demo.queue().launch()  #share=True)
-
+demo.queue().launch()