import os import torch import numpy as np import matplotlib.cm as cm import cv2 import folder_paths import comfy.model_management as mm from comfy.utils import ProgressBar, load_torch_file from .video_depth_anything.video_depth import VideoDepthAnything def ensure_even(value): return value if value % 2 == 0 else value + 1 def preprocess(tensor, max_res=-1): # Convert from float [0,1] to uint8 [0,255] frames = (tensor.numpy() * 255).astype(np.uint8) original_height, original_width = frames.shape[1:3] # Resize if max_res specified if max_res > 0 and max(original_height, original_width) > max_res: scale = max_res / max(original_height, original_width) height = ensure_even(round(original_height * scale)) width = ensure_even(round(original_width * scale)) resized_frames = [] for frame in frames: frame = cv2.resize(frame, (width, height), interpolation=cv2.INTER_AREA) resized_frames.append(frame) frames = np.array(resized_frames) return frames def postprocess_inferno(depths): colormap = np.array(cm.get_cmap("inferno").colors) d_min, d_max = depths.min(), depths.max() depth_frames = [] for i in range(depths.shape[0]): depth = depths[i] depth_norm = ((depth - d_min) / (d_max - d_min) * 255).astype(np.uint8) depth_vis = (colormap[depth_norm] * 255).astype(np.uint8) depth_frames.append(depth_vis) return torch.from_numpy(np.array(depth_frames).astype(np.float32) / 255.0) def postprocess_gray(depths): d_min, d_max = depths.min(), depths.max() # global normalisation depths = (depths - d_min) / (d_max - d_min + 1e-6) # to Torch tensor = torch.from_numpy(depths.astype(np.float32)) # [T,H,W] tensor = tensor.unsqueeze(-1).repeat(1, 1, 1, 3) # [T,H,W,3] return tensor class LoadVideoDepthAnythingModel: @classmethod def INPUT_TYPES(s): return { "required": { "model": (['video_depth_anything_vits.pth', 'video_depth_anything_vitb.pth', 'video_depth_anything_vitl.pth', 'metric_video_depth_anything_vits.pth', 'metric_video_depth_anything_vitb.pth', 'metric_video_depth_anything_vitl.pth'], {"default": 'video_depth_anything_vits.safetensors'}) }, } RETURN_TYPES = ("VDAMODEL",) RETURN_NAMES = ("vda_model",) FUNCTION = "loadmodel" CATEGORY = "VideoDepthAnything" def loadmodel(self, model): device = mm.get_torch_device() download_path = os.path.join(folder_paths.models_dir, "videodepthanything") model_path = os.path.join(download_path, model) if not os.path.exists(model_path): print(f"[VideoDepthAnything] - Downloading model to: {model_path}") from huggingface_hub import snapshot_download match model: case "video_depth_anything_vits.pth": repo_id = "depth-anything/Video-Depth-Anything-Small" case "video_depth_anything_vitb.pth": repo_id = "depth-anything/Video-Depth-Anything-Base" case "video_depth_anything_vitl.pth": repo_id = "depth-anything/Video-Depth-Anything-Large" case "metric_video_depth_anything_vits.pth": repo_id = "depth-anything/Metric-Video-Depth-Anything-Small" case "metric_video_depth_anything_vitb.pth": repo_id = "depth-anything/Metric-Video-Depth-Anything-Base" case "metric_video_depth_anything_vitl.pth": repo_id = "depth-anything/Metric-Video-Depth-Anything-Large" snapshot_download(repo_id=repo_id, allow_patterns=[f"*{model}*"], local_dir=download_path) print(f"[VideoDepthAnything] - Loading model from: {model_path}") model_configs = { 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}, 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}, } if "vits" in model: encoder = "vits" elif "vitb" in model: encoder = "vitb" elif "vitl" in model: encoder = "vitl" self.model = VideoDepthAnything(**model_configs[encoder], metric="metric" in model) state_dict = load_torch_file(model_path) self.model.load_state_dict(state_dict, strict=True) self.model = self.model.to(device).eval() return (self.model,) class VideoDepthAnythingProcess: @classmethod def INPUT_TYPES(s): return {"required": { "vda_model": ("VDAMODEL", ), "images": ("IMAGE", ), "input_size": ("INT",{"default": 518}), "max_res": ("INT",{"default": 1280}), "precision": (['fp16', 'fp32'], {"default": 'fp16'}), }, } RETURN_TYPES = ("DEPTHS",) RETURN_NAMES =("depths", ) FUNCTION = "process" CATEGORY = "VideoDepthAnything" def process(self, vda_model, images, input_size, max_res, precision): device = mm.get_torch_device() offload_device = mm.unet_offload_device() vda_model.to(device) pbar = ProgressBar(images.shape[0]) images_np = preprocess(images, max_res) depths = vda_model.infer_video_depth(images_np, input_size=input_size, device=str(device), pbar=pbar, fp32= True if precision == 'fp32' else False) vda_model.to(offload_device) return (depths,) class VideoDepthAnythingOutput: @classmethod def INPUT_TYPES(s): return {"required": { "depths": ("DEPTHS", ), "colormap": (['inferno', 'gray'], {"default": 'gray'}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES =("images", ) FUNCTION = "output" CATEGORY = "VideoDepthAnything" def output(self, depths, colormap): if colormap == "inferno": output = postprocess_inferno(depths) else: output = postprocess_gray(depths) return (output, ) class VideoDepthAnythingSaveEXR: @classmethod def INPUT_TYPES(s): return {"required": { "depths": ("DEPTHS", ), "folder_name": ("STRING", {"default": 'exr-1', "required": True}) }, } RETURN_TYPES = ("STRING",) RETURN_NAMES =("exr_path", ) FUNCTION = "save_exr" CATEGORY = "VideoDepthAnything" OUTPUT_NODE = True def save_exr(self, depths, folder_name): depth_exr_dir = os.path.join(folder_paths.output_directory, folder_name) os.makedirs(depth_exr_dir, exist_ok=True) import OpenEXR import Imath for i, depth in enumerate(depths): output_exr = f"{depth_exr_dir}/frame_{i:05d}.exr" header = OpenEXR.Header(depth.shape[1], depth.shape[0]) header["channels"] = { "Z": Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT)) } exr_file = OpenEXR.OutputFile(output_exr, header) exr_file.writePixels({"Z": depth.tobytes()}) exr_file.close() print(f"[VideoDepthAnything] - Saved EXR to {depth_exr_dir}") return (depth_exr_dir,)