Update app.py

app.py

@@ -4,6 +4,8 @@ import numpy as np
 import requests
 import random
 from io import BytesIO
+# from utils import *
+# from constants import *
 from pipeline_semantic_stable_diffusion_xl_img2img_ddpm import *
 from torch import inference_mode
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, AutoencoderKL
@@ -20,8 +22,9 @@ import copy
 import json
 import gc
 import random
+from time import sleep
 
-with open("sdxl_loras.json", "r") as file:
+with open("sdxl_loras-Copy1.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
         {
@@ -74,12 +77,12 @@ last_merged = False
 last_fused = False
 
 
-def load_lora(sdxl_loras, lora_scale = 1.0, progress=gr.Progress(track_tqdm=True)):
+def load_lora(sdxl_loras, random_lora_index, lora_scale = 1.0, progress=gr.Progress(track_tqdm=True)):
     global last_lora, last_merged, last_fused, sd_pipe
     
     randomize()
-    random_lora_index = random.randrange(0, len(sdxl_loras), 1)
-
+    #random_lora_index = random.randrange(0, len(sdxl_loras), 1)
+    
     repo_name = sdxl_loras[random_lora_index]["repo"]
     weight_name = sdxl_loras[random_lora_index]["weights"]
     
@@ -121,22 +124,44 @@ def load_lora(sdxl_loras, lora_scale = 1.0, progress=gr.Progress(track_tqdm=True
                 merge_incompatible_lora(full_path_lora, lora_scale)
                 last_fused = False
             last_merged = True
-        print("DONE")
-        return random_lora_index
+        print("DONE MERGING")
+        #return random_lora_index
 
 
 
 ## SEGA ##
+def shuffle_lora(sdxl_loras):
+    #random_lora_index = load_lora(sdxl_loras)
+    random_lora_index = random.randrange(0, len(sdxl_loras), 1)
+    lora_repo = sdxl_loras[random_lora_index]["repo"]
+    lora_title = sdxl_loras[random_lora_index]["title"]
+    lora_desc = f"""#### LoRA used to edit this image: 
+                    ## {lora_title}
+                    by `{lora_repo.split('/')[0]}`
+                    """
+    lora_image = sdxl_loras[random_lora_index]["image"]
+    return random_lora_index, lora_image, lora_desc, gr.update(visible=True), gr.update(height=369)
 
+def check_if_removed(input_image):
+    if(input_image is None):
+        return gr.Row(visible=False), gr.Column(elem_classes="output_column"), gr.Image(value=None)
+    else:
+        return gr.Row(), gr.Column(), gr.Image()
+
+def block_if_removed(input_image):
+    if(input_image is None):
+        raise gr.Warning("Photo removed. Upload a new one!")
+    
 def edit(sdxl_loras,
          input_image,
             wts, zs,
          do_inversion,
-
+         random_lora_index,
+         progress=gr.Progress(track_tqdm=True)
         ):
     show_share_button = gr.update(visible=True)
     
-    random_lora_index = load_lora(sdxl_loras)
+    load_lora(sdxl_loras, random_lora_index)
     
     src_prompt = ""
     skip = 18
@@ -161,12 +186,12 @@ def edit(sdxl_loras,
                                                 # eta = 1.0,
                                               )
 
-        wts = gr.State(value=wts_tensor)
-        zs = gr.State(value=zs_tensor)
+        wts = wts_tensor
+        zs = zs_tensor
         do_inversion = False
     
 
-    latnets = wts.value[skip].expand(1, -1, -1, -1)
+    latnets = wts[skip].expand(1, -1, -1, -1)
     
     editing_prompt = [sdxl_loras[random_lora_index]["trigger_word"]]
     reverse_editing_direction = [False]
@@ -188,13 +213,13 @@ def edit(sdxl_loras,
     sega_out = sd_pipe(prompt=tar_prompt, latents=latnets, guidance_scale = tar_cfg_scale,
                           # num_images_per_prompt=1,
                           # num_inference_steps=steps,
-                         wts=wts.value, zs=zs.value[skip:], **editing_args)
+                         wts=wts, zs=zs[skip:], **editing_args)
     
-    lora_repo = sdxl_loras[random_lora_index]["repo"]
-    lora_desc = f"### LoRA Used To Edit this Image: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[random_lora_index]['is_nc'] else '' }"
-    lora_image = sdxl_loras[random_lora_index]["image"]
+    #lora_repo = sdxl_loras[random_lora_index]["repo"]
+    #lora_desc = f"### LoRA Used To Edit this Image: {lora_repo}' }"
+    #lora_image = sdxl_loras[random_lora_index]["image"]
       
-    return sega_out.images[0], wts, zs, do_inversion, lora_image, lora_desc, gr.update(visible=True), gr.update(visible=True)
+    return sega_out.images[0], wts, zs, do_inversion, gr.update(height=512), gr.Column(elem_classes="output_column_reverse")
     
     
         
@@ -232,64 +257,86 @@ def crop_image(image):
 ########
 
 with gr.Blocks(css="style.css") as demo:
-  
-
     def reset_do_inversion():
         return True
 
-
-    gr.HTML(
-        """<h1><img src="https://i.imgur.com/jpMRW5y.png" alt="LEDITS LoRA Photobooth"></h1>""",
-    )
-    wts = gr.State()
-    zs = gr.State()
-    reconstruction = gr.State()
-    do_inversion = gr.State(value=True)
-    gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
-
-    with gr.Row():
-        input_image = gr.Image(label="Input Image", interactive=True, source="webcam", height=512, width=512)
-        sega_edited_image = gr.Image(label=f"LEDITS Edited Image", interactive=False, elem_id="output_image", height=512, width=512)
-        # input_image.style(height=365, width=365)
-        # sega_edited_image.style(height=365, width=365)
-    
-    with gr.Row():
-        lora_image = gr.Image(interactive=False, height=128, width=128, visible=False)
-        lora_desc = gr.HTML(visible=False)
+    gr.HTML("""<img style="margin: 0 auto; width: 180px; margin-bottom: .5em" src="https://i.imgur.com/A4BP6Lx.png" alt="LEDITS SDXL LoRA Photobooth">""")
+    with gr.Box(elem_id="total_box"):
+        gr.HTML(
+            """<h1>LEDITS SDXL LoRA Photobooth</h1>
+            <h3>Smile, take a pic 📷✨ and <code>it'll be inverted on SDXL and a random SDXL LoRA will be applied</code></h3>
+            """,
+        )
+        wts = gr.State()
+        zs = gr.State()
+        reconstruction = gr.State()
+        do_inversion = gr.State(value=True)
+        gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
+        gr_lora_index = gr.State()
+        with gr.Row():
+            input_image = gr.Image(label="Input Image", interactive=True, source="webcam", height=512, width=512, elem_id="input_image")
+            with gr.Column(elem_classes="output_column") as output_column:
+                with gr.Row(visible=False) as loaded_lora:
+                    lora_image = gr.Image(interactive=False, height=128, width=128, elem_id="lora_image", show_label=False, show_download_button=False)
+                    lora_desc = gr.Markdown()
+                sega_edited_image = gr.Image(label=f"LEDITS Edited Image", interactive=False, elem_id="output_image", height=512, width=512)
+
+
+
+
+        with gr.Row():
+            run_button = gr.Button("Rerun with the same picture", visible=True, elem_id="run_again")
+
+
+        run_button.click(
+                fn=shuffle_lora,
+                inputs=[gr_sdxl_loras],
+                outputs=[gr_lora_index, lora_image, lora_desc, loaded_lora, sega_edited_image],
+                queue=False
+            ).then(
+            fn=edit,
+            inputs=[gr_sdxl_loras,
+                    input_image,
+                    wts, zs,
+                    do_inversion,
+                    gr_lora_index
+
+            ],
+            outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column])
+
+        input_image.change(
+            fn = check_if_removed,
+            inputs = [input_image],
+            outputs = [loaded_lora, output_column, sega_edited_image],
+            queue=False,
+            show_progress=False
+        ).then(
+            fn = block_if_removed,
+            inputs = [input_image],
+            queue=False,
+            show_progress=False
+        ).success(
+            fn = reset_do_inversion,
+            outputs = [do_inversion],
+            queue = False).then(
+                fn=shuffle_lora,
+                inputs=[gr_sdxl_loras],
+                outputs=[gr_lora_index, lora_image, lora_desc, loaded_lora, sega_edited_image],
+                queue=False
+            ).then(
+                fn=edit,
+                inputs=[gr_sdxl_loras,
+                        input_image,
+                        wts, zs,
+                        do_inversion,
+                        gr_lora_index
+                ],
+                outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column]
+            )
+    gr.HTML('''
+    <img src="https://iccv2023.thecvf.com/img/LogoICCV23V04.svg" width="400" style="margin: 0 auto; display: none" id='iccv_logo' />
+    ''')
     
-   
-    with gr.Row():
-        run_button = gr.Button("Run again!", visible=True)
-        
-
-    run_button.click(
-        fn=edit,
-        inputs=[gr_sdxl_loras,
-                input_image,
-                wts, zs,
-                do_inversion,
-
-        ],
-        outputs=[sega_edited_image, wts, zs, do_inversion,  lora_image, lora_desc, lora_image, lora_desc])
     
-    input_image.change(
-        fn = reset_do_inversion,
-        outputs = [do_inversion],
-        queue = False).then(
-        fn=edit,
-        inputs=[gr_sdxl_loras,
-                input_image,
-                wts, zs,
-                do_inversion,
-
-
-        ],
-        outputs=[sega_edited_image, wts, zs, do_inversion,  lora_image, lora_desc, lora_image, lora_desc]
-    )
-    
-
-    
-
-
 demo.queue()
-demo.launch(share=True)
+demo.launch()