Update app.py

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['composite'])
 
 
 def sample_then_run():
@@ -435,13 +435,16 @@ with gr.Blocks(css="style.css") as demo:
     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) 
+                    # 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():
-                    gallery1 = gr.Image(label="Identity from Original Model",height=512, width=512, interactive=False)
+                    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'" ,
@@ -471,18 +474,22 @@ with gr.Blocks(css="style.css") as demo:
             
             with gr.Accordion("Advanced Options", open=False):
                 with gr.Tab("Inversion"):
-                    with gr.Column():
+                    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.Column():     
+                    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)
-                            negative_prompt1 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
                             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)
@@ -498,10 +505,10 @@ with gr.Blocks(css="style.css") as demo:
      
 
 
-    # gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
+    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, info="After sampling a new model or inverting, you can download the model below")
+        file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
         
 
 
@@ -509,18 +516,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 = [input_image, file_output])
         
             
-        sample.click(fn=sample_then_run, outputs=[gallery1, file_output])
+        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=[prompt1, negative_prompt1, cfg1, steps1, seed1, injection_step, a1, a2, a3, a4],
-                    outputs=gallery1)
-        file_input.change(fn=file_upload, inputs=file_input, 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)