app.py

import streamlit as st
import os
import subprocess
from PIL import Image, ImageOps
import torch
from diffusers import StableDiffusionInpaintPipeline
import transformers
import cv2
import diffusers
import accelerate
import warnings
import numpy as np
import os 
import shutil
warnings.filterwarnings("ignore")
st.title('Background Generation')
st.write('This app generates new backgrounds for images.')
# set environment variable for dll
os.environ['KMP_DUPLICATE_LIB_OK']='True'

@st.cache_data
def mode(width, height):
    output_width = np.floor_divide(width, 8) * 8
    output_height = np.floor_divide(height, 8) * 8
    return output_width, output_height

def get_prompt():
    prompt = st.text_input('Enter your prompt here:', placeholder="Imagine our perfume bottle amidst a lush garden, surrounded by blooming flowers and vibrant colors.")
    return prompt

def get_negative_prompt():
    negative_prompt = st.text_input('Enter your negative prompt here:', placeholder="low quality, out of frame, watermark.. etc.")
    return negative_prompt

def get_user_input():
    st.subheader("Upload an image file, Press Clean Background Button.")
    uploaded_file = st.file_uploader("Upload a JPG image file", type=["jpg", "jpeg"])

    if uploaded_file is not None:
        user_file_path = os.path.join("data/custom_dataset/", uploaded_file.name)

        # Open the uploaded image
        uploaded_image = Image.open(uploaded_file)

        # Check if the width is larger than 640
        if uploaded_image.width > 640:
            # Calculate the proportional height based on the desired width of 640 pixels
            aspect_ratio = uploaded_image.width / uploaded_image.height
            resized_height = int(640 / aspect_ratio)
            # Resize the image to a width of 640 pixels and proportional height
            resized_image = uploaded_image.resize((640, resized_height))
        else:
            resized_image = uploaded_image

        return resized_image, user_file_path

    return None, None


def clean_files(directory):
    files = os.listdir(directory)
    for file in files:
        file_path = os.path.join(directory, file)
        if os.path.isfile(file_path):
            os.remove(file_path)

uploaded_file, user_file_path = get_user_input()
button_1 = st.button("Clean Background")

button_1_clicked = False  # Variable to track button state

def run_subprocess():
  mask_created = False
  command = "python main.py inference --dataset custom_dataset/ --arch 7 --img_size 640 --save_map True"
  subprocess.run(command, shell=True)
  mask_created = True


# Perform the necessary actions when the "Clean Background" button is clicked
st.write(button_1)

# Log data for analyzing the app later.
def log(copy = False):
    custom_dataset_directory = "data/custom_dataset/"
    processed_directory = "data/processed"
    for filename in os.listdir(custom_dataset_directory):
        file_path = os.path.join(custom_dataset_directory, filename)
        
        if copy == True:
            shutil.copy(file_path, processed_directory)  # Copy files
        else:
            shutil.move(file_path, processed_directory)  # Move files


def load_images():
    x = user_file_path.split('/')[-1]    
    uploaded_file_name = os.path.basename(user_file_path)
    image_path = os.path.join("data/custom_dataset/", x)
    dif_image = Image.open(image_path)    
    
    mask_path = os.path.join("mask/custom_dataset/", x.replace('.jpg', '.png'))
    png_image = Image.open(mask_path)
    inverted_image = ImageOps.invert(png_image) 
    return dif_image , inverted_image

if button_1:
    button_1_clicked = True    
    # Move items from data/custom_dataset/ to data/processed    
    log( copy= True)
    clean_files("data/custom_dataset/")    
    if uploaded_file is not None:   
        uploaded_file.save(user_file_path)        
        run_subprocess()
        st.success("Background cleaned.")
        log(copy = True)
        
        

st.subheader("Text your prompt and choose parameters, then press Run Model button")

# Create a two-column layout
col1, col2 = st.columns(2)

# Get user input for prompts
with col1:
    input_prompt = st.text_area('Enter Prompt', height=80)
with col2:
    input_negative_prompt = st.text_area('Enter Negative Prompt', height=80)

num_inference_steps = st.slider('Number of Inference Steps:', min_value=5, max_value=50, value=10)
num_images_per_prompt = st.slider('Image Count to be Produced:', min_value=1, max_value=2, value=1)

# use seed with torch generator
#torch.manual_seed(0)
# seed
#seed = st.slider('Seed:', min_value=0, max_value=100, value=1)
#generator = [torch.Generator(device="cuda").manual_seed(seed) for i in range(num_images_per_prompt)]

#generator = torch.Generator(device="cuda").manual_seed(0)
run_model_button = st.button("Run Model") 

@st.cache_resource
def initialize_pipe():
    pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting",
            revision="fp16",
            torch_dtype=torch.float16, #16 for gpu
            safety_checker = None,
            requires_safety_checker = False).to("cuda")

    pipe.safety_checker = None
    pipe.requires_safety_checker = False
    return pipe

def image_resize(dif_image):
    output_width, output_height = mode(dif_image.width, dif_image.height)
    while output_height > 800:
          output_height = output_height // 1.5
          output_width = output_width // 1.5
          output_width, output_height = mode(output_width, output_height)
    return output_width, output_height


def show_output(x5):
    if len(x5) == 1:
      col1, col2 = st.columns(2)
      with col1 :
        st.image(inverted_image, width=256, caption='Generated Mask', use_column_width=False)
      with col2:
        st.image(x5[0], width=256, caption='Generated Image', use_column_width=False)
        
    elif len(x5) == 2:
      col1, col2, col3 = st.columns(3)
      with col1 : 
        col1.image(inverted_image, width=256, caption='Generated Mask', use_column_width=False)
      with col2 :
        col2.image(x5[0], width=256, caption='Gener ted Image', use_column_width=False)
      with col3 :
        col3.image(x5[1], width=256, caption='Generated Image-2', use_column_width=False)

# Check if the button is clicked and all inputs are provided
if run_model_button == True and input_prompt is not None :
    st.write("Running the model...")
    dif_image , inverted_image = load_images()
    output_width, output_height = image_resize(dif_image)
    base_prompt = "high resolution, high quality, use mask. Do not distort the shape of the object. make the object stand out, show it clearly and vividly, preserving the shape of the object, use the mask"
    prompt = input_prompt + " " + base_prompt
    
    st.write("Pipe working with {0} inference steps and {1} image will be created for prompt".format(num_inference_steps, num_images_per_prompt))

    pipe = initialize_pipe()

    #output_height = 128
    #output_width = 128

    x5 = pipe(image=dif_image, mask_image=inverted_image, num_inference_steps=num_inference_steps,  # , generator= generator
                num_images_per_prompt=num_images_per_prompt, prompt=prompt, negative_prompt=input_negative_prompt,
                height=output_height, width=output_width).images
            
    show_output(x5)
    torch.cuda.empty_cache()
else:

    st.write("Please provide prompt and click the 'Run Model' button to proceed.")