Additional changes for style transfer and examples

app.py

@@ -1,15 +1,41 @@
-from pathlib import Path
-import numpy as np
+import os
 import gradio as gr
-import utils.shared_utils as st
+from pathlib import Path
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+from huggingface_hub import notebook_login
+
+from huggingface_hub import notebook_login
+#if not (Path.home()/'.huggingface'/'token').exists():
+#token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
+token = "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"
 
+import src.utils.shared_utils as st
 
-import torch
+
+import torch, logging
+logging.disable(logging.WARNING)
+torch.cuda.empty_cache()
+torch.manual_seed(3407)
 from torch import autocast
-import torchvision.transforms as T
 from contextlib import nullcontext
+torch.backends.cudnn.benchmark = True
+
+model_id = "CompVis/stable-diffusion-v1-4"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 context = autocast if device == "cuda" else nullcontext
+
+# pipe = StableDiffusionPipeline.from_pretrained(model_id,use_auth_token=token).to(device)
+#
+#
+# def infer_original(prompt,samples):
+#     with context(device):
+#         images = pipe(samples*[prompt], guidance_scale=7.5).images
+#     return images
+
+
+
+
 # Apply the transformations needed
 
 
@@ -32,8 +58,14 @@ def infer(prompt,samples):
     return images
 
 
+# def newstyleimage(choice):
+#     print(choice)
+#     if choice == "yes":
+#         return gr.Image.update(visible=True,interactive=True)
+#     return
 
-
+def styleimpose(final_input_img, ref_img):
+    return st.superimpose(final_input_img, ref_img)[0]
 
 def change_bg_option(choice):
     if choice == "I have an Image":
@@ -49,16 +81,16 @@ def change_bg_option(choice):
 title = "FSDL- One-Shot, Green-Screen,   Composition-Transfer"
 DEFAULT_TEXT = "Photorealistic scenery of bookshelf in a room"
 description = """
-<center><a href="https://docs.google.com/document/d/1fde8XKIMT1nNU72859ytd2c58LFBxepS3od9KFBrJbM/edit?usp=sharing">[PAPER]</a> <a href="https://github.com/snknitin/FSDL-Project/blob/main/src/utils/shared_utils.py">[CODE]</a></center>
+<center><a href="https://docs.google.com/document/d/1fde8XKIMT1nNU72859ytd2c58LFBxepS3od9KFBrJbM/edit?usp=sharing">[PAPER - Documentation]</a> </center>
 <details>
 <summary><b>Instructions</b></summary>
 <p style="margin-top: -3px;">With this app, you can generate a suitable background image to overlay your portrait!<br />You have several ways to set how your final auto-edited image will look like:<br /></p>
  <ul style="margin-top: -20px;margin-bottom: -15px;">
-  <li style="margin-bottom: -10px;margin-left: 20px;">Use the "<i>Inputs</i>" tab to either upload an image from your device or allow the use of your webcam to capture</li>
-  <li style="margin-left: 20px;">Use the "<i>Background Image Inputs</i>" to upload your own background</li>
-  <li style="margin-left: 20px;">Use the "<i>Text prompt</i>" tab to generate a satisfactory bacground image.</li>
+  <li style="margin-bottom: -10px;margin-left: 20px;">Use the "<i>Inputs</i>" tab to either upload an image from your device OR allow the use of your webcam to capture</li>
+  <li style="margin-left: 20px;">Use the "<i>Background Image Inputs</i>" to upload your own background. OR</li>
+  <li style="margin-left: 20px;">Use the "<i>Text prompt</i>" tab to generate a satisfactory background image using Stable Diffusion.</li>
 </ul> 
-<p>After customization, just hit "<i>Edit</i>" and wait a few seconds.<br />The final image will be available for download <br /> <b>Enjoy!<b><p>
+<p>After deciding, just hit "<i>Select</i>" to ensure those images are processed.<br />The final image will be available for download <br /> <b>Enjoy!<b><p>
 </details>
 """
 
@@ -71,6 +103,13 @@ running = """
 * **Smoothing** - Given than image resolutions and clarity can be an issue, this smoothing button makes your final image crisp after the stylization transfer. Fair warning - this last process can take 5-10 mins
 """
 
+style_message = """ 
+This image above will be the content image. By default, the style will be copied from the input foreground image.
+
+If you have a different image in mind,  would you like to upload a different image?  
+Click yes to add a new style reference image"""
+
+
 
 demo = gr.Blocks()
 
@@ -101,10 +140,7 @@ with demo:
                                               placeholder="Enter your prompt to generate a background image... something like - Photorealistic scenery of bookshelf in a room")
 
                             samples = gr.Slider(label="Number of Images", minimum=1, maximum=5, value=2, step=1)
-                            btn = gr.Button("Generate images",variant="primary").style(
-                                margin=False,
-                                rounded=(False, True, True, False),
-                            )
+                            btn = gr.Button("Generate images",variant="primary")
 
                     gallery = gr.Gallery(label="Generated images", show_label=True).style(grid=(1, 3), height="auto")
                     # image_options = gr.Radio(label="Pick", interactive=True, choices=None, type="value")
@@ -130,10 +166,17 @@ with demo:
 
     with gr.Row(scale=1):
 
-        with gr.Column(scale=1):
-            supimp_btn = gr.Button("SuperImpose")
-            overlay_img = gr.Image(shape=(800, 800), label="Overlay", type="pil")
+        with gr.Box():
+            with gr.Column(scale=1):
+                supimp_btn = gr.Button("SuperImpose")
+                overlay_img = gr.Image(shape=(800, 800), label="Overlay", type="pil")
+                gr.Markdown(style_message)
+                #img_choice = gr.Radio(choices= ["yes"],interactive=True,type='value')
+                ref_img = gr.Image(shape=(800, 800),label="Style Reference", type="pil",interactive=True)
+                ref_img2 = gr.Image(shape=(800, 800), label="Style Reference", type="pil", interactive=True, visible=False)
+                ref_btn = gr.Button("Use this style")
 
+        ref_btn.click(fn=styleimpose, inputs=[final_input_img, ref_img], outputs=[ref_img2])
 
         with gr.Column(scale=1):
             style_btn = gr.Button("Composition-Transfer",variant="primary")
@@ -143,10 +186,14 @@ with demo:
             submit_btn = gr.Button("Smoothen",variant="primary")
             output_img = gr.Image(shape=(800, 800),label="FinalSmoothened Image",type="pil")
 
-        supimp_btn.click(fn=st.superimpose, inputs=[final_input_img, final_back_img], outputs=[overlay_img])
-        style_btn.click(fn=st.style_transfer, inputs=[overlay_img,final_input_img], outputs=[style_img])
+        supimp_btn.click(fn=st.superimpose, inputs=[final_input_img, final_back_img], outputs=[overlay_img,ref_img])
+        style_btn.click(fn=st.style_transfer, inputs=[overlay_img,ref_img2], outputs=[style_img])
         submit_btn.click(fn=st.smoother, inputs=[style_img,overlay_img], outputs=[output_img])
 
+
+    gr.Examples([["profile_new.png","back_img.png"]],[final_input_img, final_back_img])
+    gr.Examples([["profile_new.png","bedroom with a bookshelf in the background and a small stool to sit on the right side, photorealistic",3]], [final_input_img,text,samples])
+
 demo.queue()
 demo.launch()
 

utils/shared_utils.py

@@ -80,7 +80,7 @@ def memory_limit_image_resize(cont_img):
 def superimpose(input_img,back_img):
     matte_img = remove(input_img)
     back_img.paste(matte_img, (0, 0), matte_img)
-    return back_img
+    return back_img,input_img