NikeZoldyck's picture
adding the gradio app code
4158574
from pathlib import Path
import numpy as np
import gradio as gr
import utils.shared_utils as st
import torch
from torch import autocast
import torchvision.transforms as T
from contextlib import nullcontext
device = "cuda" if torch.cuda.is_available() else "cpu"
context = autocast if device == "cuda" else nullcontext
# Apply the transformations needed
def select_input(input_img,webcm_img):
if input_img is None:
img= webcm_img
else:
img=input_img
return img
def infer(prompt,samples):
images= []
selections = ["Img_{}".format(str(i+1).zfill(2)) for i in range(samples)]
with context(device):
for _ in range(samples):
back_img = st.stableDiffusionAPICall(prompt)
images.append(back_img)
return images
def change_bg_option(choice):
if choice == "I have an Image":
return gr.Image(shape=(800, 800))
elif choice == "Generate one for me":
return gr.update(lines=8, visible=True, value="Please enter a text prompt")
else:
return gr.update(visible=False)
# TEXT
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>
<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>
</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>
</details>
"""
running = """
### Instructions for running the 3 S's in sequence
* **Superimpose** - This button allows you to isolate the foreground from your image and overlay it on the background. Remove background using alpha matting
* **Style-Transfer** - This button transfer the style from your original image to re-map your new background realistically. Uses Nvidia FastPhotoStyle
* **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
"""
demo = gr.Blocks()
with demo:
gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>" + title + "</h1>")
with gr.Box():
gr.Markdown(description)
# First row - Inputs
with gr.Row(scale=1):
with gr.Column():
with gr.Tabs():
with gr.TabItem("Upload "):
input_img = gr.Image(shape=(800, 800), interactive=True, label="You")
with gr.TabItem("Webcam Capture"):
webcm_img = gr.Image(source="webcam", streaming=True, shape=(800, 800), interactive=True)
inp_select_btn = gr.Button("Select")
with gr.Column():
with gr.Tabs():
with gr.TabItem("Upload"):
bgm_img = gr.Image(shape=(800, 800), type="pil", interactive=True, label="The Background")
bgm_select_btn = gr.Button("Select")
with gr.TabItem("Generate via Text Prompt"):
with gr.Box():
with gr.Row().style(mobile_collapse=False, equal_height=True):
text = gr.Textbox(lines=7,
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),
)
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")
text.submit(infer, inputs=[text, samples], outputs=gallery)
btn.click(infer, inputs=[text, samples], outputs=gallery, show_progress=True, status_tracker=None)
# Second Row - Backgrounds
with gr.Row(scale=1):
with gr.Column():
final_input_img = gr.Image(shape=(800, 800), type="pil", label="Foreground")
with gr.Column():
final_back_img = gr.Image(shape=(800, 800), type="pil", label="Background", interactive=True)
bgm_select_btn.click(fn=lambda x: x, inputs=bgm_img, outputs=final_back_img)
inp_select_btn.click(select_input, [input_img, webcm_img], final_input_img)
with gr.Row(scale=1):
with gr.Box():
gr.Markdown(running)
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.Column(scale=1):
style_btn = gr.Button("Composition-Transfer",variant="primary")
style_img = gr.Image(shape=(800, 800),label="Style-Transfer Image",type="pil")
with gr.Column(scale=1):
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])
submit_btn.click(fn=st.smoother, inputs=[style_img,overlay_img], outputs=[output_img])
demo.queue()
demo.launch()