Skip to content

Commit a743f3c

Browse files
fix: implement model to func for depth anything
1 parent 332bc9d commit a743f3c

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,20 @@ class DepthAnythingPipeline(RawModel):
1212
for Invoke's Model Management System"""
1313

1414
def __init__(self, pipeline: DepthEstimationPipeline) -> None:
15-
self.pipeline = pipeline
15+
self._pipeline = pipeline
1616

1717
def generate_depth(self, image: Image.Image) -> Image.Image:
18-
depth_map = self.pipeline(image)["depth"]
18+
depth_map = self._pipeline(image)["depth"]
1919
assert isinstance(depth_map, Image.Image)
2020
return depth_map
2121

2222
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
2427

2528
def calc_size(self) -> int:
2629
from invokeai.backend.model_manager.load.model_util import calc_module_size
2730

28-
return calc_module_size(self.pipeline.model)
31+
return calc_module_size(self._pipeline.model)

0 commit comments

Comments
 (0)