SS

app.py

@@ -1,206 +1,3 @@
-<<<<<<< HEAD
-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.float32, #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:
-=======
 import streamlit as st
 import os
 import subprocess
@@ -404,4 +201,4 @@ if run_model_button == True and input_prompt is not None :
     torch.cuda.empty_cache()
 else:
 >>>>>>> 76f5454639852d715b029dd10f1e299517131c28
-    st.write("Please provide prompt and click the 'Run Model' button to proceed.")     
+    st.write("Please provide prompt and click the 'Run Model' button to proceed.")