fix: make DA Pipeline a subclass of RawModel

This commit is contained in:
blessedcoolant
2024-07-31 21:14:49 +05:30
parent 18f89ed5ed
commit b4cf78a95d
2 changed files with 9 additions and 6 deletions

View File

@ -1,10 +1,13 @@
from typing import cast
from typing import Optional, cast
import torch
from PIL import Image
from transformers.pipelines import DepthEstimationPipeline
from invokeai.backend.raw_model import RawModel
class DepthAnythingPipeline:
class DepthAnythingPipeline(RawModel):
"""Custom wrapper for the Depth Estimation pipeline from transformers adding compatibility
for Invoke's Model Management System"""
@ -20,6 +23,9 @@ class DepthAnythingPipeline:
depth_map = depth_map.resize((resolution, new_height))
return depth_map
def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
pass
def calc_size(self) -> int:
from invokeai.backend.model_manager.load.model_util import calc_module_size