diff --git a/invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py b/invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py index 5663ccd7b0..b6ef529687 100644 --- a/invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py +++ b/invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py @@ -12,17 +12,20 @@ class DepthAnythingPipeline(RawModel): for Invoke's Model Management System""" def __init__(self, pipeline: DepthEstimationPipeline) -> None: - self.pipeline = pipeline + self._pipeline = pipeline def generate_depth(self, image: Image.Image) -> Image.Image: - depth_map = self.pipeline(image)["depth"] + depth_map = self._pipeline(image)["depth"] assert isinstance(depth_map, Image.Image) return depth_map def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None): - pass + if device is not None and device.type not in {"cpu", "cuda"}: + device = None + self._pipeline.model.to(device=device, dtype=dtype) + self._pipeline.device = self._pipeline.model.device def calc_size(self) -> int: from invokeai.backend.model_manager.load.model_util import calc_module_size - return calc_module_size(self.pipeline.model) + return calc_module_size(self._pipeline.model)