remove bg

handler.py

@@ -1,11 +1,9 @@
 from typing import Dict, List, Any
 import torch
-from torch import autocast
 from huggingface_hub import hf_hub_download
 from diffusers import DiffusionPipeline
-import base64
-from io import BytesIO
 from safetensors.torch import load_file
+from transformers import pipeline
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -43,6 +41,14 @@ class EndpointHandler:
             tokenizer=self.pipe.tokenizer_2,
         )
 
+        self.remove_bg = pipeline(
+            "image-segmentation",
+            model="briaai/RMBG-1.4",
+            device=device,
+            revision="22532afbdabdc36b2d30a334076720ac72a06f83",
+            trust_remote_code=True,
+        )
+
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
         Args:
@@ -57,7 +63,9 @@ class EndpointHandler:
         images = self.pipe(inputs, **data["parameters"]).images
         image = images[0]
 
-        return image
+        image_no_bg = self.remove_bg(image)
+
+        return image_no_bg
 
 
 if __name__ == "__main__":