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.")