Use from_pretrained

app.py

@@ -2,7 +2,6 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
-from huggingface_hub import hf_hub_download
 import gradio as gr
 from gradio_imageslider import ImageSlider
 from briarmbg import BriaRMBG
@@ -10,15 +9,10 @@ import PIL
 from PIL import Image
 from typing import Tuple
 
-net=BriaRMBG()
-# model_path = "./model1.pth"
-model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
-if torch.cuda.is_available():
-    net.load_state_dict(torch.load(model_path))
-    net=net.cuda()
-else:
-    net.load_state_dict(torch.load(model_path,map_location="cpu"))
-net.eval() 
+net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net.to(device)
 
     
 def resize_image(image):