Update app.py

app.py

@@ -22,6 +22,8 @@ from sampling import sample_weights
 from lora_w2w import LoRAw2w
 from huggingface_hub import snapshot_download
 import numpy as np
+
+
 global device
 global generator 
 global unet
@@ -59,9 +61,8 @@ def sample_model():
     unet, _, _, _, _ = load_models(device)
     network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
  
-
 @torch.no_grad()
-def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     global device
     global generator 
     global unet
@@ -110,7 +111,6 @@ def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
     return image
 
 
-
 @torch.no_grad()
 def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
 
@@ -199,7 +199,7 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     network.reset()
    
     return (original_image, image)
-
+    
 def sample_then_run():
     global original_image
     sample_model()    
@@ -210,6 +210,7 @@ def sample_then_run():
     steps = 50
     original_image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
+
     
     return (original_image, original_image), "model.pt"
 
@@ -273,11 +274,15 @@ class CustomImageDataset(Dataset):
         if self.transform:
             image = self.transform(image)
         return image
-
-def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
+    
+def invert(dict, 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)
@@ -308,7 +313,7 @@ def invert(image, mask, 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) 
 
 
@@ -346,12 +351,14 @@ def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
     return network
 
 
-# @spaces.GPU(duration=200)
-def run_inversion(input_image, pcs, epochs, weight_decay,lr):
+
+
+def run_inversion(dict, pcs, epochs, weight_decay,lr):
     global network
-    init_image = input_image["background"].convert("RGB").resize((512, 512))
-    mask = input_image["layers"][0].convert("RGB").resize((512, 512))
-    network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
+    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)
 
 
     #sample an image
@@ -360,20 +367,19 @@ def run_inversion(input_image, pcs, epochs, weight_decay,lr):
     seed = 5
     cfg = 3.0
     steps = 50
-    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    original_image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
-    return (image,init_image), "model.pt"
+    return (original_image, original_image), "model.pt"
     
 
 
 
-
-# @spaces.GPU()
 def file_upload(file):
     global unet
     del unet
     global network
     global device
+    global original_image
     
 
     
@@ -402,18 +408,20 @@ def file_upload(file):
     seed = 5
     cfg = 3.0
     steps = 50
-    image = inference( prompt, negative_prompt, cfg, steps, seed)
-    return image
+    original_image = inference( prompt, negative_prompt, cfg, steps, seed)
+    return (original_image, original_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>
@@ -443,13 +451,12 @@ with gr.Blocks(css="style.css") as demo:
                 with gr.Row():
                     with gr.Column():
                         sample = gr.Button("🎲 Sample New Model")
-                        file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
+                        file_output1 = gr.File(label="Download Sampled Model", container=True, interactive=False)
                         file_input = gr.File(label="Upload Model", container=True)
 
 
-                            # invert_button = gr.Button("⏪ Invert")
                     with gr.Column():
-                        image_slider = ImageSlider(position=0.5, type="pil", height=512, width=512)
+                        image_slider1 = ImageSlider(position=0.5, type="pil", height=512, width=512)
                         
                         prompt1 = gr.Textbox(label="Prompt",
                                                         info="Make sure to include 'sks person'" ,
@@ -461,11 +468,11 @@ with gr.Blocks(css="style.css") as demo:
                 
 
                         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)       
+                            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 = 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)
+                            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)
 
                 
                 
@@ -473,7 +480,7 @@ with gr.Blocks(css="style.css") as demo:
                             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_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
+                            injection_step1 = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
                                 
                                 
                     
@@ -481,44 +488,48 @@ with gr.Blocks(css="style.css") as demo:
                         submit1 = gr.Button("Generate")
 
     with gr.Tab("Inversion"):
-        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)
+        with gr.Row():
+            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")
+
+            with gr.Column():
+                image_slider2 = ImageSlider(position=0.5, type="pil", height=512, width=512)
+                            
+                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.Column():
-            image_slider = ImageSlider(position=0.5, type="pil", height=512, width=512)
                         
-            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)
-
                     
-                
 
-            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.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)
 
-                
-                
-            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_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
-                                
-                                
                     
+                    
+                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")
+                submit2 = gr.Button("Generate")
+
+        file_output2 = gr.File(label="Download Inverted Model", container=True, interactive=False)
 
         
                     
@@ -527,17 +538,14 @@ with gr.Blocks(css="style.css") as demo:
 
     
                 
-    sample.click(fn=sample_then_run, outputs=[image_slider, file_output])
-
-        
-    submit1.click(fn=edit_inference, inputs=[ prompt1, negative_prompt1, cfg1, steps1, seed1, injection_step, a1, a2, a3, a4], outputs=image_slider)
-    file_input.change(fn=file_upload, inputs=file_input, outputs = image_slider)
+    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)
 
 
-    
-    
-        
-            
+ 
 demo.queue().launch()