From a743f3c9b59e1116480934c634fc3b52945bc763 Mon Sep 17 00:00:00 2001 From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com> Date: Sat, 3 Aug 2024 00:37:17 +0530 Subject: [PATCH] fix: implement model to func for depth anything --- .../depth_anything/depth_anything_pipeline.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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)