Update app.py

app.py

@@ -24,7 +24,11 @@ import gc
 import random
 from time import sleep
 
-with open("sdxl_loras.json", "r") as file:
+IMAGE_DATASET_DIR = Path("image_dataset") / f"train-{uuid4()}"
+IMAGE_DATASET_DIR.mkdir(parents=True, exist_ok=True)
+IMAGE_JSONL_PATH = IMAGE_DATASET_DIR / "metadata.jsonl"
+
+with open("sdxl_loras-Copy1.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
         {
@@ -147,16 +151,24 @@ def shuffle_lora(sdxl_loras, selected_lora=None):
 
 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), gr.State(value=None), gr.Button(visible=False)
+        return gr.Row(visible=False), gr.Column(elem_classes="output_column"), gr.Image(value=None), gr.State(value=None), gr.Column(visible=False)
     else:
-        return gr.Row(), gr.Column(), gr.Image(), None, gr.Button()
+        return gr.Row(), gr.Column(), gr.Image(), None, gr.Column()
 
 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 select_lora(selected_state: gr.SelectData, sdxl_loras):
+    random_lora_index = selected_state.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: 
+                    ## {lora_title}
+                    by `{lora_repo.split('/')[0]}`
+                    """
+    lora_image = sdxl_loras[random_lora_index]["image"]
+    return random_lora_index, lora_image, lora_desc
     
     
 def edit(sdxl_loras,
@@ -186,7 +198,7 @@ def edit(sdxl_loras,
                                                 source_prompt= src_prompt,
                                                 # source_prompt_2 = None,
                                                 source_guidance_scale = src_cfg_scale,
-                                                negative_prompt = "blurry, ugly, bad quality",
+                                                negative_prompt = "blurry, bad quality",
                                                 # negative_prompt_2 = None,
                                                 num_inversion_steps = steps,
                                                 skip_steps = skip,
@@ -226,7 +238,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"), gr.Button(visible=True)
+    return sega_out.images[0], wts, zs, do_inversion, gr.update(height=512), gr.Column(elem_classes="output_column_reverse"), gr.Row(visible=True)
     
     
         
@@ -257,7 +269,16 @@ def crop_image(image):
     return image
 
     
-
+def save_preferences(sdxl_loras, selected_lora, input_image, result_image):
+    lora_id = sdxl_loras[selected_lora]["repo"]
+    uuid = uuid4()
+    input_image_path = IMAGE_DATASET_DIR / f"{uuid}-input.png"
+    output_image_path = IMAGE_DATASET_DIR / f"{uuid}-output.png"
+    Image.fromarray(input_image).save(input_image_path)
+    Image.fromarray(result_image).save(output_image_path)
+    with IMAGE_JSONL_PATH.open("a") as f:
+        json.dump({"selected_lora": lora_id, "input_image":input_image_path.name, "result_image":output_image_path.name}, f)
+        f.write("\n")
 
 ########r
 # demo #
@@ -288,24 +309,25 @@ with gr.Blocks(css="style.css") as demo:
                     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.Column(visible=False) as buttons_area:
+            with gr.Row(elem_id="buttons_area"):
+                print_button = gr.HTML('<button onclick="window.print()" class="lg secondary svelte-cmf5ev" id="print" style="width: 100%;margin-bottom:0">Print 🖨️</button>')
+                run_button = gr.Button("Regenerate with the same picture 🖼️🎲", elem_id="run_again")
 
 
-
-
-        with gr.Row():
-            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
-            )
+            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],
+            inputs=gr_sdxl_loras,
+            outputs=[gr_picked_lora, lora_image, lora_desc],
             queue=False
         )
         
@@ -322,12 +344,17 @@ 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, run_button])
+            outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column, buttons_area]
+        ).then(
+            fn=save_preferences,
+            inputs=[gr_sdxl_loras, gr_lora_index, input_image, sega_edited_image],
+            queue=False
+        )
 
         input_image.change(
             fn = check_if_removed,
             inputs = [input_image],
-            outputs = [loaded_lora, output_column, sega_edited_image, gr_picked_lora, run_button],
+            outputs = [loaded_lora, output_column, sega_edited_image, gr_picked_lora, buttons_area],
             queue=False,
             show_progress=False
         ).then(
@@ -351,7 +378,11 @@ 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, run_button]
+                outputs=[sega_edited_image, wts, zs, do_inversion, sega_edited_image, output_column, buttons_area]
+            ).then(
+                fn=save_preferences,
+                inputs=[gr_sdxl_loras, gr_lora_index, input_image, sega_edited_image],
+                queue=False
             )
     gr.HTML('''
     <h1>Hugging Face at</h1>