Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
wangsen
MinerU
Commits
d58b24b5
Commit
d58b24b5
authored
Jun 17, 2025
by
myhloli
Browse files
fix: add conditional imports for torch and torch_npu in model_utils.py
parent
bd5252d9
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
6 additions
and
6 deletions
+6
-6
mineru/utils/model_utils.py
mineru/utils/model_utils.py
+6
-6
No files found.
mineru/utils/model_utils.py
View file @
d58b24b5
...
@@ -6,6 +6,12 @@ import numpy as np
...
@@ -6,6 +6,12 @@ import numpy as np
from
mineru.utils.boxbase
import
get_minbox_if_overlap_by_ratio
from
mineru.utils.boxbase
import
get_minbox_if_overlap_by_ratio
try
:
import
torch
import
torch_npu
except
ImportError
:
pass
def
crop_img
(
input_res
,
input_img
,
crop_paste_x
=
0
,
crop_paste_y
=
0
):
def
crop_img
(
input_res
,
input_img
,
crop_paste_x
=
0
,
crop_paste_y
=
0
):
...
@@ -297,14 +303,11 @@ def get_res_list_from_layout_res(layout_res, iou_threshold=0.7, overlap_threshol
...
@@ -297,14 +303,11 @@ def get_res_list_from_layout_res(layout_res, iou_threshold=0.7, overlap_threshol
def
clean_memory
(
device
=
'cuda'
):
def
clean_memory
(
device
=
'cuda'
):
import
torch
if
device
==
'cuda'
:
if
device
==
'cuda'
:
if
torch
.
cuda
.
is_available
():
if
torch
.
cuda
.
is_available
():
torch
.
cuda
.
empty_cache
()
torch
.
cuda
.
empty_cache
()
torch
.
cuda
.
ipc_collect
()
torch
.
cuda
.
ipc_collect
()
elif
str
(
device
).
startswith
(
"npu"
):
elif
str
(
device
).
startswith
(
"npu"
):
import
torch_npu
if
torch_npu
.
npu
.
is_available
():
if
torch_npu
.
npu
.
is_available
():
torch_npu
.
npu
.
empty_cache
()
torch_npu
.
npu
.
empty_cache
()
elif
str
(
device
).
startswith
(
"mps"
):
elif
str
(
device
).
startswith
(
"mps"
):
...
@@ -322,13 +325,10 @@ def clean_vram(device, vram_threshold=8):
...
@@ -322,13 +325,10 @@ def clean_vram(device, vram_threshold=8):
def
get_vram
(
device
):
def
get_vram
(
device
):
import
torch
if
torch
.
cuda
.
is_available
()
and
str
(
device
).
startswith
(
"cuda"
):
if
torch
.
cuda
.
is_available
()
and
str
(
device
).
startswith
(
"cuda"
):
total_memory
=
torch
.
cuda
.
get_device_properties
(
device
).
total_memory
/
(
1024
**
3
)
# 将字节转换为 GB
total_memory
=
torch
.
cuda
.
get_device_properties
(
device
).
total_memory
/
(
1024
**
3
)
# 将字节转换为 GB
return
total_memory
return
total_memory
elif
str
(
device
).
startswith
(
"npu"
):
elif
str
(
device
).
startswith
(
"npu"
):
import
torch_npu
if
torch_npu
.
npu
.
is_available
():
if
torch_npu
.
npu
.
is_available
():
total_memory
=
torch_npu
.
npu
.
get_device_properties
(
device
).
total_memory
/
(
1024
**
3
)
# 转为 GB
total_memory
=
torch_npu
.
npu
.
get_device_properties
(
device
).
total_memory
/
(
1024
**
3
)
# 转为 GB
return
total_memory
return
total_memory
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment