Update app.py

app.py

@@ -22,12 +22,6 @@ def load_stable_diffusion_model():
 # Initialize the model globally
 sd_model = load_stable_diffusion_model()
 
-def remove_background_stable_diffusion(image_path):
-    image = Image.open(image_path).convert("RGB")
-    prompt = "Remove background from the image"
-    output = sd_model(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
-    return output
-
 def remove_background_rembg(input_path):
     print(f"Removing background using rembg for image: {input_path}")
     with open(input_path, 'rb') as i:
@@ -437,9 +431,7 @@ def process_single_image(image_path, output_folder, bg_method, canvas_size_name,
     filename = os.path.basename(image_path)
     try:
         print(f"Processing image: {filename}")
-        if bg_method == 'stable_diffusion':
-            image_with_no_bg = remove_background_stable_diffusion(image_path)
-        elif bg_method == 'rembg':
+        if bg_method == 'rembg':
             image_with_no_bg = remove_background_rembg(image_path)  # Placeholder for existing function
         elif bg_method == 'bria':
             image_with_no_bg = remove_background_bria(image_path)  # Placeholder for existing function
@@ -471,8 +463,8 @@ def process_single_image(image_path, output_folder, bg_method, canvas_size_name,
         output_filename = f"{os.path.splitext(filename)[0]}.{output_ext}"
         output_path = os.path.join(output_folder, output_filename)
 
-        # Apply watermark only if the filename ends with "_01" and watermark_path is provided
-        if os.path.splitext(filename)[0].endswith("_01") and watermark_path:
+        # Apply watermark if watermark_path is provided
+        if watermark_path:
             watermark = Image.open(watermark_path).convert("RGBA")
             canvas.paste(watermark, (0, 0), watermark)
             log.append({"action": "add_watermark"})
@@ -689,7 +681,7 @@ with gr.Blocks(theme="NoCrypt/miku@1.2.2") as iface:
         num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=5)
 
     with gr.Row():
-        bg_method = gr.Radio(choices=["bria", "rembg", "stable_diffusion"], label="Background Removal Method", value="stable_diffusion")
+        bg_method = gr.Radio(choices=["bria", "rembg"], label="Background Removal Method", value="rembg")
         bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Background Choice", value="white")
         custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)