kadirnar
e21711a049
- Updated model selection in image-to-image, mask generation, object detection, video classification, zero-shot image classification, zero-shot object detection, document question answering, feature extraction, image-to-text, image-text-to-text, visual question answering, conversational, fill mask, question answering, sentence similarity, table question answering, zero-shot classification, text classification, text generation, and token classification pipelines to use a default string instead of a model list. - Removed unused keypoint detection, summarization, text2text generation, and translation pipelines. - Added utility nodes for type conversion: IntToString, StringToInt, FloatToString, and StringToFloat. - Updated all pipeline calls to include trust_remote_code=True for enhanced security. - Improved requirements by adding timm and opencv-python for better functionality.
83 lines
2.5 KiB
Python
83 lines
2.5 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
|
|
|
|
class LoadDepthModel:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": ("STRING", {"default": "Intel/dpt-hybrid-midas"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("DEPTH_MODEL", "IMAGE_PROCESSOR")
|
|
FUNCTION = "load_depth_model"
|
|
|
|
CATEGORY = "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/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",
|
|
}
|