app.py

import os
import gradio as gr
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 = os.environ.get('token')


import utils.shared_utils as st


import torch, logging
logging.disable(logging.WARNING)
torch.cuda.empty_cache()
torch.manual_seed(3407)
from torch import autocast
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,revision="fp16", torch_dtype=torch.float16,use_auth_token=token).to(device)


def infer(prompt,samples):
    images=[]
    with context(device):
        for _ in range(samples):
            image = pipe([prompt],num_inference_steps=50, guidance_scale=7.5, height=400, width=400).images[0]
            images.append(image)
    return images



# 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= []
#     with context(device):
#         for _ in range(samples):
#             back_img = st.stableDiffusionAPICall(prompt)
#             images.append(back_img)
#     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":
        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 = "Green-Screen Image 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 - Abstract]</a>|<a href="https://docs.google.com/presentation/d/1baE6sF1vvb-mRC0rIvLfCk4MlHwtXfodBqHEmr6wnwk/edit?usp=sharing">[POSTER]</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. 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 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>
"""

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 transfers the style from your original image to re-map your new background realistically. Uses Nvidia FastPhotoStyle
* **Smoothing** - Given that the 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, a good choice for the style is input foreground image.

If you have a different image in mind, you can remove the default and upload it here.
Ideally transfer works better if your input foreground is also superimposed on the style image, so you may want to create it using the same steps
 ."""



demo = gr.Blocks()

with demo:
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>" + title + "</h1>")

    with gr.Box():
        with gr.Row():
            with gr.Column():
                gr.Markdown(description)
            with gr.Column():
                gr.Image("capture.png")

    # 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, label= "Prompt",
                                              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")

                    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.Box():
            with gr.Column(scale=1):
                supimp_btn = gr.Button("SuperImpose")
                overlay_img = gr.Image(shape=(800, 800), label="Overlay/Content Image", 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 Image", 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",variant="primary")
        #
        # ref_btn.click(fn=styleimpose, inputs=[final_input_img, ref_img], outputs=[ref_img])

        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,ref_img])
        style_btn.click(fn=st.style_transfer, inputs=[overlay_img,ref_img], outputs=[style_img])
        submit_btn.click(fn=st.smoother, inputs=[style_img,overlay_img], outputs=[output_img])

    gr.Examples(examples=[["profile_new.png","back_img.png"]], label="AlphaMatting- Remove BG",
                 inputs=[final_input_img, final_back_img], outputs=[overlay_img])
    gr.Examples(examples=[["profile_new.png",
                  "bedroom with a bookshelf in the background and a small stool to sit on the right side, photorealistic",
                  3]],inputs= [final_input_img,text,samples], label="Text2Img - Stable Diffisuon")
    gr.Examples(examples=[["cont_img.png","ref_img.png"]],inputs=[overlay_img, ref_img], label = "Nvidia - FastPhotoStyle")


demo.queue(concurrency_count=40, max_size=20)
demo.launch(enable_queue=True,max_threads=150)