mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
fix: implement model to func for depth anything
This commit is contained in:
parent
332bc9da5b
commit
a743f3c9b5
@ -12,17 +12,20 @@ class DepthAnythingPipeline(RawModel):
|
|||||||
for Invoke's Model Management System"""
|
for Invoke's Model Management System"""
|
||||||
|
|
||||||
def __init__(self, pipeline: DepthEstimationPipeline) -> None:
|
def __init__(self, pipeline: DepthEstimationPipeline) -> None:
|
||||||
self.pipeline = pipeline
|
self._pipeline = pipeline
|
||||||
|
|
||||||
def generate_depth(self, image: Image.Image) -> Image.Image:
|
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)
|
assert isinstance(depth_map, Image.Image)
|
||||||
return depth_map
|
return depth_map
|
||||||
|
|
||||||
def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
|
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:
|
def calc_size(self) -> int:
|
||||||
from invokeai.backend.model_manager.load.model_util import calc_module_size
|
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)
|
||||||
|
Loading…
Reference in New Issue
Block a user