89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
import comfy.model_management as mm
|
|
from comfy.utils import ProgressBar
|
|
import torchvision.transforms.functional as F
|
|
|
|
depth_model_name_list = [
|
|
"Intel/dpt-hybrid-midas",
|
|
"LiheYoung/depth-anything-small-hf",
|
|
"facebook/dpt-dinov2-small-nyu",
|
|
]
|
|
|
|
class LoadDepthModel:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": (depth_model_name_list, {"default": depth_model_name_list[0]}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("DEPTH_MODEL", "IMAGE_PROCESSOR")
|
|
FUNCTION = "load_depth_model"
|
|
|
|
CATEGORY = "ComputerVision/Transformers/DepthEstimation"
|
|
|
|
def load_depth_model(self, model_name):
|
|
model = AutoModelForDepthEstimation.from_pretrained(model_name, trust_remote_code=True)
|
|
image_processor = AutoImageProcessor.from_pretrained(model_name, trust_remote_code=True)
|
|
|
|
return (model, image_processor)
|
|
|
|
class DepthEstimationInference:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": ("DEPTH_MODEL",),
|
|
"processor": ("IMAGE_PROCESSOR",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "depth_inference"
|
|
|
|
CATEGORY = "Transformers/ComputerVision/DepthEstimation"
|
|
|
|
def depth_inference(self, image, model, processor):
|
|
# Convert ComfyUI image to PIL image
|
|
img = 255. * image[0].cpu().numpy()
|
|
pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
|
|
|
# Prepare the image for the model
|
|
inputs = processor(images=pil_image, return_tensors="pt")
|
|
pbar = ProgressBar(len(inputs["pixel_values"]))
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
predicted_depth = outputs.predicted_depth
|
|
|
|
# Interpolate to original size
|
|
prediction = torch.nn.functional.interpolate(
|
|
predicted_depth.unsqueeze(1),
|
|
size=pil_image.size[::-1],
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
)
|
|
|
|
# Convert depth map to tensor in ComfyUI format
|
|
tensor_image = prediction.squeeze().cpu().numpy()
|
|
tensor_image = tensor_image.astype(np.float32) / tensor_image.max()
|
|
tensor_image = torch.from_numpy(tensor_image)[None,]
|
|
|
|
return (tensor_image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadDepthModel": LoadDepthModel,
|
|
"DepthEstimationInference": DepthEstimationInference,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadDepthModel": "Load Depth Model",
|
|
"DepthEstimationInference": "Depth Estimation Inference",
|
|
}
|