Update README.md

README.md

@@ -114,10 +114,10 @@ import torch
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
 
-birefnet = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
+model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
 torch.set_float32_matmul_precision(['high', 'highest'][0])
-birefnet.to('cuda')
-birefnet.eval()
+model.to('cuda')
+model.eval()
 
 # Data settings
 image_size = (1024, 1024)
@@ -132,7 +132,7 @@ input_images = transform_image(image).unsqueeze(0).to('cuda')
 
 # Prediction
 with torch.no_grad():
-    preds = birefnet(input_images)[-1].sigmoid().cpu()
+    preds = model(input_images)[-1].sigmoid().cpu()
 pred = preds[0].squeeze()
 pred_pil = transforms.ToPILImage()(pred)
 mask = pred_pil.resize(image.size)