Update app.py

app.py

@@ -13,7 +13,29 @@ from enum import Enum
 
 api_key = os.getenv("FAI_API_KEY")
 api = os.getenv("FAI_API")
+rmbgkey = os.getenv("RMBGKEY")
 
+def rmbg(pil_image):
+    
+    # Convert PIL image to bytes
+    image_bytes = BytesIO()
+    pil_image.save(image_bytes, format='PNG')
+    image_bytes.seek(0)
+    
+    # Send the image to the remove.bg API
+    response = requests.post(
+        'https://api.remove.bg/v1.0/removebg',
+        files={'image_file': ('filename.png', image_bytes, 'image/png')},
+        data={'size': 'auto'},
+        headers={'X-Api-Key': rmbgkey}
+    )
+    
+    if response.status_code == 200:        
+        # Convert the bytes response to a PIL image
+        result_image = Image.open(BytesIO(response.content))
+        return result_image
+    else:
+        return None
 
 def image_to_base64(image):
     # Open the image file
@@ -113,8 +135,17 @@ def resize_to_fit(max_size, original_size):
     return new_width, new_height
 
     
-def process_generate(fore, prompt, intensity, mode, refprompt, bg): 
-
+def process_generate(fore, prompt, intensity, mode, refprompt, isrmbg): 
+    
+    if isrmbg:
+        try:
+            rmbgfore = rmbg(fore)
+            print(f"Background removed!")
+            if rmbgfore:
+                fore = rmbgfore
+        except:
+            pass
+         
     fore = create_square_image(fore)
     size = fore.size
     image_width = size[0]
@@ -127,7 +158,8 @@ def process_generate(fore, prompt, intensity, mode, refprompt, bg):
         gr.Warning("ℹ️ The input image size is too big, I will lower it!")
         image_width, image_height = resize_to_fit((1500,1500), (image_width, image_height))
 
-
+    
+        
     
     forestr = image_to_base64(fore.convert("RGBA"))    
     data = {
@@ -156,6 +188,8 @@ def process_generate(fore, prompt, intensity, mode, refprompt, bg):
 
     return image
 
+
+
 def update_value(val):
   return val
 
@@ -210,10 +244,7 @@ with block:
                             value=Stage.FULL.value,
                             label="Generation Mode", type='value')
         mode.change(fn=update_value, inputs=mode, outputs=mode)
-        with gr.Column():
-            bg = gr.Checkbox(info="Remove Backgroung")
-            bg.change(fn=update_value, inputs=bg, outputs=bg)
-            
+        
         with gr.Column():
             gr.HTML('''
                 <div>
@@ -226,12 +257,13 @@ with block:
     with gr.Row():
         intensity = gr.Slider(label="Refiner Strength", minimum=1.0, maximum=7.0, value=3.0, step=0.5)
         intensity.change(fn=update_value, inputs=intensity, outputs=intensity)
+        isrmbg = gr.Checkbox(label="Remove Background")
+        isrmbg.change(fn=update_value, inputs=isrmbg, outputs=isrmbg)
         generate_button = gr.Button(value="Generate")
         
-
-
-
-    ips = [fore, prompt, intensity, mode, refprompt, bg]
+    
+    
+    ips = [fore, prompt, intensity, mode, refprompt, isrmbg]
     generate_button.click(fn=process_generate, inputs=ips, outputs=[result_gallery])