Spaces:
Running
on
Zero
Running
on
Zero
Show runtime counter in logs.
Browse files
app.py
CHANGED
@@ -20,7 +20,8 @@ from io import BytesIO
|
|
20 |
torch.set_float32_matmul_precision('high')
|
21 |
torch.jit.script = lambda f: f
|
22 |
|
23 |
-
device = "cuda" if torch.cuda.is_available() else "
|
|
|
24 |
|
25 |
|
26 |
def array_to_pil_image(image: np.ndarray, size: Tuple[int, int] = (1024, 1024)) -> Image.Image:
|
@@ -67,7 +68,8 @@ def predict(image, resolution, weights_file):
|
|
67 |
global birefnet
|
68 |
# Load BiRefNet with chosen weights
|
69 |
_weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
|
70 |
-
|
|
|
71 |
birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
|
72 |
birefnet.to(device)
|
73 |
birefnet.eval()
|
|
|
20 |
torch.set_float32_matmul_precision('high')
|
21 |
torch.jit.script = lambda f: f
|
22 |
|
23 |
+
device = "cuda" if torch.cuda.is_available() else "CPU"
|
24 |
+
counter_runtime = 0
|
25 |
|
26 |
|
27 |
def array_to_pil_image(image: np.ndarray, size: Tuple[int, int] = (1024, 1024)) -> Image.Image:
|
|
|
68 |
global birefnet
|
69 |
# Load BiRefNet with chosen weights
|
70 |
_weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
|
71 |
+
counter_runtime += 1
|
72 |
+
print('Using weights: {} in runtime {}'.format(_weights_file, counter_runtime))
|
73 |
birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
|
74 |
birefnet.to(device)
|
75 |
birefnet.eval()
|