Update app_3.py

app_3.py

@@ -150,8 +150,14 @@ vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae")
 unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")
 # Load model directly
 from transformers import AutoModelForImageSegmentation
-rmbg = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
-rmbg = rmbg.to(device=device, dtype=torch.float32)  # Keep this as float32
+# rmbg = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
+# rmbg = rmbg.to(device=device, dtype=torch.float32)  # Keep this as float32
+
+# remove bg
+rmbg = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+rmbg = rmbg.to(device)
 
 model = DepthAnythingV2(encoder='vits', features=64, out_channels=[48, 96, 192, 384])
 model.load_state_dict(torch.load('checkpoints/depth_anything_v2_vits.pth', map_location=device))