Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: implement model to func for depth anything

Browse Source
This commit is contained in:
blessedcoolant 2024-08-03 00:37:17 +05:30
parent 332bc9da5b
commit a743f3c9b5
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py
View File

@ -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)