remove(image)

app.py

@@ -6,6 +6,7 @@ import random
 from diffusers import DiffusionPipeline
 import torch
 from sympy.core.random import choice
+from rembg import remove
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -55,6 +56,10 @@ def infer(
         generator=generator,
     ).images[0]
 
+    if del_back:
+        image = remove(image)
+        
+
     return image, seed