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: Move the manual image resizing out of the depth anything pipeline

Browse Source
This commit is contained in:
blessedcoolant 2024-07-31 23:35:52 +05:30
parent 13fb2d1f49
commit daf899f9c4
2 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
invokeai/app/invocations/controlnet_image_processors.py
View File

@ -594,6 +594,7 @@ class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small", "small_v2"] DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small", "small_v2"]
# DepthAnything V2 Small model is licensed under Apache 2.0 but not the base and large models.
DEPTH_ANYTHING_MODELS = { DEPTH_ANYTHING_MODELS = {
"large": "LiheYoung/depth-anything-large-hf", "large": "LiheYoung/depth-anything-large-hf",
"base": "LiheYoung/depth-anything-base-hf", "base": "LiheYoung/depth-anything-base-hf",
@ -627,7 +628,13 @@ class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
source=DEPTH_ANYTHING_MODELS[self.model_size], loader=load_depth_anything source=DEPTH_ANYTHING_MODELS[self.model_size], loader=load_depth_anything
) as depth_anything_detector: ) as depth_anything_detector:
assert isinstance(depth_anything_detector, DepthAnythingPipeline) assert isinstance(depth_anything_detector, DepthAnythingPipeline)
return depth_anything_detector.generate_depth(image, self.resolution) depth_map = depth_anything_detector.generate_depth(image)
# Resizing to user target specified size
new_height = int(image.size[1] * (self.resolution / image.size[0]))
depth_map = depth_map.resize((self.resolution, new_height))
return depth_map
@invocation( @invocation(

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

@ -1,4 +1,4 @@
from typing import Optional, cast from typing import Optional
import torch import torch
from PIL import Image from PIL import Image
@ -14,13 +14,9 @@ class DepthAnythingPipeline(RawModel):
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, resolution: int = 512): def generate_depth(self, image: Image.Image) -> Image.Image:
image_width, image_height = image.size
depth_map = self.pipeline(image)["depth"] depth_map = self.pipeline(image)["depth"]
assert isinstance(depth_map, Image.Image) assert isinstance(depth_map, Image.Image)
new_height = int(image_height * (resolution / image_width))
depth_map = depth_map.resize((resolution, new_height))
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):