added remove bg

app.py

@@ -3,6 +3,8 @@ import json
 import logging
 import torch
 import base64
+import rembg
+import numpy as np
 from io import BytesIO
 from PIL import Image
 from diffusers import (
@@ -35,8 +37,13 @@ def image_to_base64(image: Image) -> str:
     img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
     return img_base64
 
-
-
+def remove_bg(image: Image):
+    input_array_bg = np.array(image)
+    # Apply background removal using rembg
+    output_array_bg = rembg.remove(input_array_bg)
+    # Create a PIL Image from the output array
+    return Image.fromarray(output_array_bg)
+    
 def update_selection(evt: gr.SelectData):
     selected_lora = loras[evt.index]
     new_placeholder = f"Type a prompt for {selected_lora['title']}"
@@ -115,8 +122,8 @@ def run_lora(prompt, negative_prompt, cfg_scale, steps, scheduler, seed, width,
 
     # Unload LoRA weights
     pipe.unload_lora_weights()
-
-    return image_to_base64(image)
+    image_without_bg = remove_bg(image)
+    return image_to_base64(image_without_bg)
 
 with gr.Blocks(theme=gr.themes.Soft()) as app:
 

requirements.txt

@@ -2,4 +2,5 @@ torch
 diffusers
 spaces
 transformers
-peft
+peft
+rembg