Update app_3.py

app_3.py

@@ -75,11 +75,11 @@ transform_image = transforms.Compose(
 # Model paths
 model_path = './models/iclight_sd15_fc.safetensors'
 model_path2 = './checkpoints/depth_anything_v2_vits.pth'
-model_path3 = './checkpoints/sam2_hiera_large.pt'
+#model_path3 = './checkpoints/sam2_hiera_large.pt'
 model_path4 = './checkpoints/config.json'
 model_path5 = './checkpoints/preprocessor_config.json'
-model_path6 = './configs/sam2_hiera_l.yaml'
-model_path7 = './mvadapter_i2mv_sdxl.safetensors'
+#model_path6 = './configs/sam2_hiera_l.yaml'
+#model_path7 = './mvadapter_i2mv_sdxl.safetensors'
 
 # Base URL for the repository
 BASE_URL = 'https://huggingface.co/Ashoka74/Placement/resolve/main/'
@@ -88,11 +88,11 @@ BASE_URL = 'https://huggingface.co/Ashoka74/Placement/resolve/main/'
 model_urls = {
     model_path: 'iclight_sd15_fc.safetensors',
     model_path2: 'depth_anything_v2_vits.pth',
-    model_path3: 'sam2_hiera_large.pt',
+    #model_path3: 'sam2_hiera_large.pt',
     model_path4: 'config.json',
     model_path5: 'preprocessor_config.json',
-    model_path6: 'sam2_hiera_l.yaml',
-    model_path7: 'mvadapter_i2mv_sdxl.safetensors'
+    #model_path6: 'sam2_hiera_l.yaml',
+    #model_path7: 'mvadapter_i2mv_sdxl.safetensors'
 }
 
 # Ensure directories exist
@@ -149,10 +149,16 @@ tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")
 text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder")
 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-2.0", trust_remote_code=True)#, token=os.getenv('token'))
-rmbg = rmbg.to(device=device, dtype=torch.float32)  # Keep this as float32
+# rmbg = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True)#, token=os.getenv('token'))
+# rmbg = rmbg.to(device=device, dtype=torch.float32)  # Keep this as float32
+
+rmbg = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet", trust_remote_code=True
+        )
+rmbg = rmbg.to_device(device=device, dtype=torch.float32) 
 
 # remove bg
 # rmbg = AutoModelForImageSegmentation.from_pretrained(