Update app.py

app.py

@@ -10,21 +10,21 @@ from pipeline_semantic_stable_diffusion_xl_img2img_ddpm import *
 from torch import inference_mode
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, AutoencoderKL
 from diffusers import DDIMScheduler
-# from share_btn import community_icon_html, loading_icon_html, share_js
+from share_btn import community_icon_html, loading_icon_html, share_js
 import torch
 from huggingface_hub import hf_hub_download
 from diffusers import DiffusionPipeline
-from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
+from cog_sdxl.dataset_and_utils import TokenEmbeddingsHandler
 import json
 from safetensors.torch import load_file
-# import lora
+import lora
 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 = [
         {
@@ -79,10 +79,9 @@ last_fused = False
 
 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)
-    
+    print(random_lora_index)
+    #print(sdxl_loras)
     repo_name = sdxl_loras[random_lora_index]["repo"]
     weight_name = sdxl_loras[random_lora_index]["weights"]
     
@@ -130,9 +129,13 @@ def load_lora(sdxl_loras, random_lora_index, lora_scale = 1.0, progress=gr.Progr
 
 
 ## SEGA ##
-def shuffle_lora(sdxl_loras):
-    #random_lora_index = load_lora(sdxl_loras)
-    random_lora_index = random.randrange(0, len(sdxl_loras), 1)
+def shuffle_lora(sdxl_loras, selected_lora=None):
+    print("selected_lora in shuffle_lora", selected_lora)
+    if(selected_lora is not None):
+        random_lora_index = selected_lora
+    else:
+        random_lora_index = random.randrange(0, len(sdxl_loras), 1)
+    print("random_lora_index in shuffle_lora: ", random_lora_index)
     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: 
@@ -140,18 +143,22 @@ def shuffle_lora(sdxl_loras):
                     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)
+    return gr.update(), 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)
+        return gr.Row(visible=False), gr.Column(elem_classes="output_column"), gr.Image(value=None), gr.State(value=None), gr.Button(visible=False)
     else:
-        return gr.Row(), gr.Column(), gr.Image()
+        return gr.Row(), gr.Column(), gr.Image(), None, gr.Button()
 
 def block_if_removed(input_image):
     if(input_image is None):
         raise gr.Warning("Photo removed. Upload a new one!")
     
+def select_lora(selected_state: gr.SelectData):
+    return selected_state.index
+    
+    
 def edit(sdxl_loras,
          input_image,
             wts, zs,
@@ -160,7 +167,7 @@ def edit(sdxl_loras,
          progress=gr.Progress(track_tqdm=True)
         ):
     show_share_button = gr.update(visible=True)
-    
+    print("random_lora_index in edit: ", random_lora_index)
     load_lora(sdxl_loras, random_lora_index)
     
     src_prompt = ""
@@ -169,7 +176,7 @@ def edit(sdxl_loras,
     tar_cfg_scale = 15
     src_cfg_scale = 3.5
     tar_prompt = ""
-    
+    print("Is do_inversion?", do_inversion)
     if do_inversion:
         image = load_image(input_image, device=device).to(torch.float16)
         with inference_mode():
@@ -209,7 +216,7 @@ def edit(sdxl_loras,
       edit_momentum_scale=0.3,
       edit_mom_beta=0.6,
       eta=1,)
-     
+    torch.manual_seed(torch.seed())
     sega_out = sd_pipe(prompt=tar_prompt, latents=latnets, guidance_scale = tar_cfg_scale,
                           # num_images_per_prompt=1,
                           # num_inference_steps=steps,
@@ -219,7 +226,7 @@ def edit(sdxl_loras,
     #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, gr.update(height=512), gr.Column(elem_classes="output_column_reverse")
+    return sega_out.images[0], wts, zs, do_inversion, gr.update(height=512), gr.Column(elem_classes="output_column_reverse"), gr.Button(visible=True)
     
     
         
@@ -252,7 +259,7 @@ def crop_image(image):
     
 
 
-########
+########r
 # demo #
 ########
 
@@ -273,6 +280,7 @@ with gr.Blocks(css="style.css") as demo:
         do_inversion = gr.State(value=True)
         gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
         gr_lora_index = gr.State()
+        gr_picked_lora = 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:
@@ -285,13 +293,26 @@ with gr.Blocks(css="style.css") as demo:
 
 
         with gr.Row():
-            run_button = gr.Button("Rerun with the same picture", visible=True, elem_id="run_again")
-
-
+            run_button = gr.Button("Rerun with the same picture", elem_id="run_again", visible=False)
+        with gr.Accordion("Tired of randomizing? Pick your LoRA", open=False, elem_id="pick"):
+            choose_gallery = gr.Gallery(
+                value=[(item["image"], item["title"]) for item in sdxl_loras_raw],
+                allow_preview=False,
+                columns=6,
+                elem_id="gallery",
+                show_share_button=False
+            )
+            
+        choose_gallery.select(
+            fn=select_lora,
+            outputs=[gr_picked_lora],
+            queue=False
+        )
+        
         run_button.click(
                 fn=shuffle_lora,
-                inputs=[gr_sdxl_loras],
-                outputs=[gr_lora_index, lora_image, lora_desc, loaded_lora, sega_edited_image],
+                inputs=[gr_sdxl_loras, gr_picked_lora],
+                outputs=[sega_edited_image, gr_lora_index, lora_image, lora_desc, loaded_lora, sega_edited_image],
                 queue=False
             ).then(
             fn=edit,
@@ -300,14 +321,13 @@ with gr.Blocks(css="style.css") as demo:
                     wts, zs,
                     do_inversion,
                     gr_lora_index
-
             ],
-            outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column])
+            outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column, run_button])
 
         input_image.change(
             fn = check_if_removed,
             inputs = [input_image],
-            outputs = [loaded_lora, output_column, sega_edited_image],
+            outputs = [loaded_lora, output_column, sega_edited_image, gr_picked_lora, run_button],
             queue=False,
             show_progress=False
         ).then(
@@ -320,8 +340,8 @@ with gr.Blocks(css="style.css") as demo:
             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],
+                inputs=[gr_sdxl_loras, gr_picked_lora],
+                outputs=[sega_edited_image, gr_lora_index, lora_image, lora_desc, loaded_lora, sega_edited_image],
                 queue=False
             ).then(
                 fn=edit,
@@ -331,11 +351,12 @@ with gr.Blocks(css="style.css") as demo:
                         do_inversion,
                         gr_lora_index
                 ],
-                outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column]
+                outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column, run_button]
             )
     gr.HTML('''
-    <img src="https://iccv2023.thecvf.com/img/LogoICCV23V04.svg" width="400" style="margin: 0 auto; display: none" id='iccv_logo' />
-    ''')
+    <h1>Hugging Face at</h1>
+    <img src="https://iccv2023.thecvf.com/img/LogoICCV23V04.svg" width="400" style="margin: 0 auto;" />
+    ''', elem_id="iccv_logo")
     
     
 demo.queue()