Spaces:
Running
on
Zero
Running
on
Zero
Commit
·
d967d62
1
Parent(s):
1352148
Upgrade the weights loading method to avoid duplicated loading.
Browse files
app.py
CHANGED
|
@@ -35,13 +35,6 @@ class ImagePreprocessor():
|
|
| 35 |
return image
|
| 36 |
|
| 37 |
|
| 38 |
-
from transformers import AutoModelForImageSegmentation
|
| 39 |
-
weights_path = 'BiRefNet'
|
| 40 |
-
birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', weights_path)), trust_remote_code=True)
|
| 41 |
-
birefnet.to(device)
|
| 42 |
-
birefnet.eval()
|
| 43 |
-
birefnet.weights_path = weights_path
|
| 44 |
-
|
| 45 |
usage_to_weights_file = {
|
| 46 |
'General': 'BiRefNet',
|
| 47 |
'Portrait': 'BiRefNet-portrait',
|
|
@@ -51,6 +44,13 @@ usage_to_weights_file = {
|
|
| 51 |
'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
|
| 52 |
}
|
| 53 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 54 |
|
| 55 |
@spaces.GPU
|
| 56 |
def predict(image, resolution, weights_file):
|
|
|
|
| 35 |
return image
|
| 36 |
|
| 37 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 38 |
usage_to_weights_file = {
|
| 39 |
'General': 'BiRefNet',
|
| 40 |
'Portrait': 'BiRefNet-portrait',
|
|
|
|
| 44 |
'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
|
| 45 |
}
|
| 46 |
|
| 47 |
+
from transformers import AutoModelForImageSegmentation
|
| 48 |
+
weights_path = 'General'
|
| 49 |
+
birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file[weights_path])), trust_remote_code=True)
|
| 50 |
+
birefnet.to(device)
|
| 51 |
+
birefnet.eval()
|
| 52 |
+
birefnet.weights_path = weights_path
|
| 53 |
+
|
| 54 |
|
| 55 |
@spaces.GPU
|
| 56 |
def predict(image, resolution, weights_file):
|