change layout + adjust logic a bit

app.py

@@ -1,6 +1,6 @@
 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
@@ -21,7 +21,7 @@ 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
 global device
 global generator 
 global unet
@@ -32,7 +32,7 @@ global noise_scheduler
 global network
 device = "cuda:0"
 generator = torch.Generator(device=device)
-
+from gradio_imageslider import ImageSlider
 
 
 
@@ -61,7 +61,7 @@ def sample_model():
 
  
 @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
@@ -113,7 +113,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):
+def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
 
     global device
     global generator 
@@ -196,7 +196,7 @@ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, st
     network.proj = torch.nn.Parameter(original_weights)
     network.reset()
 
-    return image
+    return (image, input_image["background"])
 
 
 def sample_then_run():
@@ -344,10 +344,11 @@ def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
 
 
 
-def run_inversion(dict, pcs, epochs, weight_decay,lr):
+def run_inversion(input_image, pcs, epochs, weight_decay,lr):
     global network
-    init_image = dict["image"].convert("RGB").resize((512, 512))
-    mask = dict["mask"].convert("RGB").resize((512, 512))
+    print(len(input_image["layers"]))
+    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)
 
 
@@ -359,7 +360,7 @@ def run_inversion(dict, pcs, epochs, weight_decay,lr):
     steps = 50
     image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
-    return image, "model.pt"
+    return (image,init_image), "model.pt"
     
 
 
@@ -430,71 +431,83 @@ intro = """
 
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(intro)
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("""<div style="text-align: justify;"> Click below to sample an identity-encoding model.""")
-            sample = gr.Button("Sample New Model")
-            gr.Markdown("""<div style="text-align: justify;"> Or upload an image below and click \"invert\". You can also optionally draw over the face to define a mask.""")
-            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)
-            pcs = gr.Slider(label="# Principal Components", value=10000, step=1, minimum=1, maximum=10000, interactive=True)
-            with gr.Accordion("Advanced Options", open=False):
+    
+    gr.Markdown("""
+    Click sample (to sample an identity) *or* upload an image & click `invert` to get started ✨
+    > 💡 When inverting, draw a mask over the face for improved results. 
+    > To use a model previously downloaded from this demo see `Uplaoding a model` in the `Advanced options`
+    """)
+    with gr.Column():
+            with gr.Row():
                 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")
-            
-            gr.Markdown("""<div style="text-align: justify;"> Or you can upload a model below downloaded from this demo.""")
+                    # 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)
+                    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) 
+                    with gr.Row():
+                        sample = gr.Button("Sample New Model")
+                        invert_button = gr.Button("Invert")
+                with gr.Column():
+                    image_slider = ImageSlider(position=1., type="pil", height=512, width=512)
+                    # gallery1 = gr.Image(label="Identity from Original Model",height=512, width=512, interactive=False)
+                    
+            prompt1 = gr.Textbox(label="Prompt",
+                                            info="Make sure to include 'sks person'" ,
+                                            placeholder="sks person", 
+                                            value="sks person")
+                
+            # Editing 
+            with gr.Column():
+                #gallery2 = gr.Image(label="Identity from Edited Model", interactive=False, visible=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)
+#                 prompt2 = gr.Textbox(label="Prompt",
+#                                         info="Make sure to include 'sks person'" ,
+#                                         placeholder="sks person", 
+#                                         value="sks person", visible=False)
+#                 seed2 = gr.Number(value=5, label="Seed", precision=0, interactive=True, visible=False)
 
-            file_input = gr.File(label="Upload Model", container=True)
 
 
-            
-                
+#                 submit2 = gr.Button("Generate", visible=False)
 
-        with gr.Column():
-            gallery1 = gr.Image(label="Identity from Original Model", interactive=False)
-            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.Accordion("Advanced Options", open=False):
-                with gr.Column():     
-                        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")
+                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=400, 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():     
+                            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)
+                            seed1 = gr.Number(value=5, label="Seed", precision=0, interactive=True)
+                    with gr.Row():
+                            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.Tab("Editing"):
+                #     with gr.Column():     
+                        # 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)
+                        # injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
+                        # negative_prompt2 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
+                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)
 
             submit1 = gr.Button("Generate")
-
-            
-                
-        with gr.Column():
-            gallery2 = gr.Image(label="Identity from Edited Model", 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)
-            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.Accordion("Advanced Options", open=False):
-                with gr.Column():     
-                        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)
-                        injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
-                        negative_prompt2 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
-            
-            submit2 = gr.Button("Generate")
-            
+     
 
 
     gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
@@ -508,18 +521,21 @@ with gr.Blocks(css="style.css") as demo:
 
         invert_button.click(fn=run_inversion,
                     inputs=[input_image, pcs, epochs, weight_decay,lr], 
-                    outputs = [gallery1, file_output])
+                    outputs = [image_slider, file_output])
         
             
-        sample.click(fn=sample_then_run, outputs=[gallery1, file_output])
-
-        submit1.click(fn=inference,
-                    inputs=[prompt1, negative_prompt1, cfg1, steps1, seed1],
-                    outputs=gallery1)
-        submit2.click(fn=edit_inference,
-                    inputs=[prompt2, negative_prompt2, cfg2, steps2, seed2, injection_step, a1, a2, a3, a4],
-                    outputs=gallery2)
-        file_input.change(fn=file_upload, inputs=file_input, outputs = gallery1)
+        sample.click(fn=sample_then_run, outputs=[input_image, file_output])
+
+        # submit1.click(fn=inference,
+        #             inputs=[prompt1, negative_prompt1, cfg1, steps1, seed1],
+        #             outputs=gallery1)
+        # submit1.click(fn=edit_inference,
+        #             inputs=[input_image, prompt1, negative_prompt1, cfg1, steps1, seed1, injection_step, a1, a2, a3, a4],
+        #             outputs=image_slider)
+        submit1.click(
+        fn=edit_inference, inputs=[input_image, 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 = input_image)
 
 
     

requirements.txt

@@ -70,4 +70,5 @@ urllib3==2.2.1
 wandb==0.17.0
 xxhash==3.4.1
 yarl==1.9.4
-zipp==3.19.0
+zipp==3.19.0
+gradio_imageslider==0.0.20