|
21 | 21 | from controlnet_aux.util import HWC3, ade_palette
|
22 | 22 | from PIL import Image
|
23 | 23 | from pydantic import BaseModel, Field, field_validator, model_validator
|
| 24 | +from transformers import pipeline |
| 25 | +from transformers.pipelines import DepthEstimationPipeline |
24 | 26 |
|
25 | 27 | from invokeai.app.invocations.baseinvocation import (
|
26 | 28 | BaseInvocation,
|
|
44 | 46 | from invokeai.app.services.shared.invocation_context import InvocationContext
|
45 | 47 | from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, heuristic_resize
|
46 | 48 | from invokeai.backend.image_util.canny import get_canny_edges
|
47 |
| -from invokeai.backend.image_util.depth_anything import DEPTH_ANYTHING_MODELS, DepthAnythingDetector |
| 49 | +from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline |
48 | 50 | from invokeai.backend.image_util.dw_openpose import DWPOSE_MODELS, DWOpenposeDetector
|
49 | 51 | from invokeai.backend.image_util.hed import HEDProcessor
|
50 | 52 | from invokeai.backend.image_util.lineart import LineartProcessor
|
51 | 53 | from invokeai.backend.image_util.lineart_anime import LineartAnimeProcessor
|
52 | 54 | from invokeai.backend.image_util.util import np_to_pil, pil_to_np
|
53 |
| -from invokeai.backend.util.devices import TorchDevice |
54 | 55 |
|
55 | 56 |
|
56 | 57 | class ControlField(BaseModel):
|
@@ -592,36 +593,48 @@ def run_processor(self, image: Image.Image) -> Image.Image:
|
592 | 593 | return color_map
|
593 | 594 |
|
594 | 595 |
|
595 |
| -DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small"] |
| 596 | +DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small", "small_v2"] |
| 597 | +# DepthAnything V2 Small model is licensed under Apache 2.0 but not the base and large models. |
| 598 | +DEPTH_ANYTHING_MODELS = { |
| 599 | + "large": "LiheYoung/depth-anything-large-hf", |
| 600 | + "base": "LiheYoung/depth-anything-base-hf", |
| 601 | + "small": "LiheYoung/depth-anything-small-hf", |
| 602 | + "small_v2": "depth-anything/Depth-Anything-V2-Small-hf", |
| 603 | +} |
596 | 604 |
|
597 | 605 |
|
598 | 606 | @invocation(
|
599 | 607 | "depth_anything_image_processor",
|
600 | 608 | title="Depth Anything Processor",
|
601 | 609 | tags=["controlnet", "depth", "depth anything"],
|
602 | 610 | category="controlnet",
|
603 |
| - version="1.1.2", |
| 611 | + version="1.1.3", |
604 | 612 | )
|
605 | 613 | class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
|
606 | 614 | """Generates a depth map based on the Depth Anything algorithm"""
|
607 | 615 |
|
608 | 616 | model_size: DEPTH_ANYTHING_MODEL_SIZES = InputField(
|
609 |
| - default="small", description="The size of the depth model to use" |
| 617 | + default="small_v2", description="The size of the depth model to use" |
610 | 618 | )
|
611 | 619 | resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
|
612 | 620 |
|
613 | 621 | def run_processor(self, image: Image.Image) -> Image.Image:
|
614 |
| - def loader(model_path: Path): |
615 |
| - return DepthAnythingDetector.load_model( |
616 |
| - model_path, model_size=self.model_size, device=TorchDevice.choose_torch_device() |
617 |
| - ) |
| 622 | + def load_depth_anything(model_path: Path): |
| 623 | + depth_anything_pipeline = pipeline(model=str(model_path), task="depth-estimation", local_files_only=True) |
| 624 | + assert isinstance(depth_anything_pipeline, DepthEstimationPipeline) |
| 625 | + return DepthAnythingPipeline(depth_anything_pipeline) |
618 | 626 |
|
619 | 627 | with self._context.models.load_remote_model(
|
620 |
| - source=DEPTH_ANYTHING_MODELS[self.model_size], loader=loader |
621 |
| - ) as model: |
622 |
| - depth_anything_detector = DepthAnythingDetector(model, TorchDevice.choose_torch_device()) |
623 |
| - processed_image = depth_anything_detector(image=image, resolution=self.resolution) |
624 |
| - return processed_image |
| 628 | + source=DEPTH_ANYTHING_MODELS[self.model_size], loader=load_depth_anything |
| 629 | + ) as depth_anything_detector: |
| 630 | + assert isinstance(depth_anything_detector, DepthAnythingPipeline) |
| 631 | + depth_map = depth_anything_detector.generate_depth(image) |
| 632 | + |
| 633 | + # Resizing to user target specified size |
| 634 | + new_height = int(image.size[1] * (self.resolution / image.size[0])) |
| 635 | + depth_map = depth_map.resize((self.resolution, new_height)) |
| 636 | + |
| 637 | + return depth_map |
625 | 638 |
|
626 | 639 |
|
627 | 640 | @invocation(
|
|
0 commit comments