File tree Expand file tree Collapse file tree 1 file changed +7
-4
lines changed
invokeai/backend/image_util/depth_anything Expand file tree Collapse file tree 1 file changed +7
-4
lines changed Original file line number Diff line number Diff line change @@ -12,17 +12,20 @@ class DepthAnythingPipeline(RawModel):
12
12
for Invoke's Model Management System"""
13
13
14
14
def __init__ (self , pipeline : DepthEstimationPipeline ) -> None :
15
- self .pipeline = pipeline
15
+ self ._pipeline = pipeline
16
16
17
17
def generate_depth (self , image : Image .Image ) -> Image .Image :
18
- depth_map = self .pipeline (image )["depth" ]
18
+ depth_map = self ._pipeline (image )["depth" ]
19
19
assert isinstance (depth_map , Image .Image )
20
20
return depth_map
21
21
22
22
def to (self , device : Optional [torch .device ] = None , dtype : Optional [torch .dtype ] = None ):
23
- pass
23
+ if device is not None and device .type not in {"cpu" , "cuda" }:
24
+ device = None
25
+ self ._pipeline .model .to (device = device , dtype = dtype )
26
+ self ._pipeline .device = self ._pipeline .model .device
24
27
25
28
def calc_size (self ) -> int :
26
29
from invokeai .backend .model_manager .load .model_util import calc_module_size
27
30
28
- return calc_module_size (self .pipeline .model )
31
+ return calc_module_size (self ._pipeline .model )
You can’t perform that action at this time.
0 commit comments