Update app.py

app.py

@@ -12,6 +12,7 @@ from torch import autocast, inference_mode
 from diffusers import StableDiffusionPipeline
 from diffusers import DDIMScheduler
 from transformers import AutoProcessor, BlipForConditionalGeneration
+from share_btn import community_icon_html, loading_icon_html, share_js
 
 # load pipelines
 sd_model_id = "stabilityai/stable-diffusion-2-1-base"
@@ -144,7 +145,7 @@ def edit(input_image,
          randomize_seed,
          src_prompt,
          src_cfg_scale):
-
+    show_share_button = gr.update(visible=True)
     if do_inversion or randomize_seed:
         x0 = load_512(input_image, device=device).to(torch.float16)
         # invert and retrieve noise maps and latent
@@ -173,7 +174,7 @@ def edit(input_image,
                           num_inference_steps=steps,
                           use_ddpm=True,  wts=wts.value, zs=zs.value[skip:], **editing_args)
       
-      return sega_out.images[0], reconstruct_button.update(visible=True), do_reconstruction, reconstruction, wts, zs, do_inversion
+      return sega_out.images[0], reconstruct_button.update(visible=True), do_reconstruction, reconstruction, wts, zs, do_inversion, show_share_button
     
     else: # if sega concepts were not added, performs regular ddpm sampling
       
@@ -181,9 +182,9 @@ def edit(input_image,
           pure_ddpm_img = sample(zs.value, wts.value, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
           reconstruction = gr.State(value=pure_ddpm_img)
           do_reconstruction = False
-          return pure_ddpm_img, reconstruct_button.update(visible=False), do_reconstruction, reconstruction, wts, zs, do_inversion
+          return pure_ddpm_img, reconstruct_button.update(visible=False), do_reconstruction, reconstruction, wts, zs, do_inversion, show_share_button
       
-      return reconstruction.value, reconstruct_button.update(visible=False), do_reconstruction, reconstruction, wts, zs, do_inversion
+      return reconstruction.value, reconstruct_button.update(visible=False), do_reconstruction, reconstruction, wts, zs, do_inversion, show_share_button
         
 
 def randomize_seed_fn(seed, randomize_seed):
@@ -309,7 +310,8 @@ def swap_visibilities(input_image,
         concept2_update = update_display_concept("Remove" if neg_guidance_2 else "Add", edit_concept_2, neg_guidance_2, sega_concepts_counter+1)
     else:
         concept2_update = gr.update(visible=False), gr.update(visible=False),gr.update(visible=False), gr.update(value=neg_guidance_2),gr.update(visible=True),gr.update(visible=False),sega_concepts_counter+1
-    return (*concept1_update[:-1], *concept2_update)
+    
+    return (gr.update(visible=True), *concept1_update[:-1], *concept2_update)
     
 
 
@@ -470,50 +472,53 @@ with gr.Blocks(css="style.css") as demo:
     sega_concepts_counter = gr.State(0)
     image_caption = gr.State(value="")
 
-
-
     with gr.Row():
-        input_image = gr.Image(label="Input Image", interactive=True)
+        input_image = gr.Image(label="Input Image", interactive=True, elem_id="input_image")
         ddpm_edited_image = gr.Image(label=f"Pure DDPM Inversion Image", interactive=False, visible=False)
-        sega_edited_image = gr.Image(label=f"LEDITS Edited Image", interactive=False)
+        sega_edited_image = gr.Image(label=f"LEDITS Edited Image", interactive=False, elem_id="output_image")
         input_image.style(height=365, width=365)
         ddpm_edited_image.style(height=365, width=365)
         sega_edited_image.style(height=365, width=365)
-
+    
+    with gr.Group(visible=False) as share_btn_container:
+        with gr.Group(elem_id="share-btn-container"):
+            community_icon = gr.HTML(community_icon_html, visible=True)
+            loading_icon = gr.HTML(loading_icon_html, visible=False)
+            share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)
+        
     with gr.Row():
-      with gr.Box(visible=False) as box1:
+      with gr.Box(visible=False, elem_id="box1") as box1:
         with gr.Row():
-          concept_1 = gr.Button(scale=3)
+          concept_1 = gr.Button(scale=3, value="")
           remove_concept1 = gr.Button("x", scale=1, min_width=10)
         with gr.Row():
             guidnace_scale_1 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
                             info="How strongly the concept should modify the image",
                                                   value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
                                                   step=0.5, interactive=True)
-      with gr.Box(visible=False) as box2:
+      with gr.Box(visible=False, elem_id="box2") as box2:
         with gr.Row():
-          concept_2 = gr.Button(scale=3)
+          concept_2 = gr.Button(scale=3, value="")
           remove_concept2 = gr.Button("x", scale=1, min_width=10)
         with gr.Row():
           guidnace_scale_2 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
                               info="How strongly the concept should modify the image",
                                                     value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
                                                     step=0.5, interactive=True)
-      with gr.Box(visible=False) as box3:
+      with gr.Box(visible=False, elem_id="box3") as box3:
         with gr.Row():
-          concept_3 = gr.Button(visible=False, scale=3)
+          concept_3 = gr.Button(scale=3, value="")
           remove_concept3 = gr.Button("x", scale=1, min_width=10)
         with gr.Row():
           guidnace_scale_3 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
                               info="How strongly the concept should modify the image",
                                                     value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
-                                                    step=0.5, interactive=True,visible=False)
+                                                    step=0.5, interactive=True)
 
 
     with gr.Row():
         inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
         
-
     with gr.Box():
         intro_segs = gr.Markdown("Add/Remove Concepts from your Image <span style=\"font-size: 12px; color: rgb(156, 163, 175)\">with Semantic Guidance</span>")
                   # 1st SEGA concept
@@ -601,6 +606,7 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Row().style(mobile_collapse=False, equal_height=True):
                 tar_prompt = gr.Textbox(
                                 label="Describe your edited image (optional)",
+                                elem_id="target_prompt",
                                 # show_label=False,
                                 max_lines=1, value="", scale=3,
                                 placeholder="Target prompt, DDPM Inversion", info = "DDPM Inversion Prompt. Can help with global changes, modify to what you would like to see"
@@ -716,7 +722,7 @@ with gr.Blocks(css="style.css") as demo:
 
 
         ],
-        outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs, do_inversion])
+        outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs, do_inversion, share_btn_container])
     # .success(fn=update_gallery_display, inputs= [prev_output_image, sega_edited_image], outputs = [gallery, gallery, prev_output_image])
 
 
@@ -829,7 +835,7 @@ with gr.Blocks(css="style.css") as demo:
         outputs = [seed],
         queue = False)
 
-
+    share_button.click(None, [], [], _js=share_js)
     
     gr.Examples(
         label='Examples',
@@ -852,7 +858,7 @@ with gr.Blocks(css="style.css") as demo:
                     tar_cfg_scale,
                     sega_concepts_counter
                ],
-        outputs=[box1, concept_1, guidnace_scale_1,neg_guidance_1, row1, row2,box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3,sega_concepts_counter],
+        outputs=[share_btn_container, box1, concept_1, guidnace_scale_1,neg_guidance_1, row1, row2,box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3,sega_concepts_counter],
         cache_examples=True
     )