nodes/Limbicnation-ComfyUIDepthEstimation
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
fca63accb1a37a90dbe43d56fba04773343c8acd
Limbicnation-ComfyUIDepthEs…/depth_estimation_node.py
T
limbicnation fca63accb1 fix: update Hugging Face model paths for V2 models 
- Updated model paths to use the correct organization 'depth-anything' instead of 'LiheYoung'
- Fixed direct URLs to point to the correct Hugging Face repositories
- Added more fallback model paths to try multiple organization/repo formats
- Updated error messages with correct download links
- Fixed tensor shape handling for output images
2025-05-04 00:07:37 +02:00

1742 lines
85 KiB
Python
Raw Blame History

import os
import numpy as np
import torch
import traceback
import time
import requests
import urllib.request
import wget
from pathlib import Path
from transformers import pipeline
from PIL import Image, ImageFilter, ImageOps, ImageDraw, ImageFont
import folder_paths
from comfy.model_management import get_torch_device, get_free_memory
import gc
import logging
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, List, Dict, Any, Optional, Union
# Try to import timm (for vision transformers)
try:
import timm
TIMM_AVAILABLE = True
except ImportError:
TIMM_AVAILABLE = False
print("Warning: timm not available. Direct loading of Depth Anything models may not work.")
# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("DepthEstimation")
# Depth Anything V2 Implementation
class DepthAnythingV2(nn.Module):
"""Direct implementation of Depth Anything V2 model"""
def __init__(self, encoder='vits', features=64, out_channels=[48, 96, 192, 384]):
super().__init__()
self.encoder = encoder
self.features = features
self.out_channels = out_channels
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Create encoder based on specification
if TIMM_AVAILABLE:
if encoder == 'vits':
self.backbone = timm.create_model('vit_small_patch16_224', pretrained=False)
self.embed_dim = 384
elif encoder == 'vitb':
self.backbone = timm.create_model('vit_base_patch16_224', pretrained=False)
self.embed_dim = 768
elif encoder == 'vitl':
self.backbone = timm.create_model('vit_large_patch16_224', pretrained=False)
self.embed_dim = 1024
else: # fallback to vits
self.backbone = timm.create_model('vit_small_patch16_224', pretrained=False)
self.embed_dim = 384
# Implement the rest of the model architecture
self.initialize_decoder()
else:
# Fallback if timm is not available
from torchvision.models import resnet50
self.backbone = resnet50(pretrained=False)
self.embed_dim = 2048
logger.warning("Using fallback ResNet50 model (timm not available)")
def initialize_decoder(self):
"""Initialize the decoder layers"""
self.neck = nn.Sequential(
nn.Conv2d(self.embed_dim, self.features, 1, 1, 0),
nn.Conv2d(self.features, self.features, 3, 1, 1),
)
# Create decoders for each level
self.decoders = nn.ModuleList([
self.create_decoder_level(self.features, self.out_channels[0]),
self.create_decoder_level(self.out_channels[0], self.out_channels[1]),
self.create_decoder_level(self.out_channels[1], self.out_channels[2]),
self.create_decoder_level(self.out_channels[2], self.out_channels[3])
])
# Final depth head
self.depth_head = nn.Sequential(
nn.Conv2d(self.out_channels[3], self.out_channels[3], 3, 1, 1),
nn.BatchNorm2d(self.out_channels[3]),
nn.ReLU(True),
nn.Conv2d(self.out_channels[3], 1, 1)
)
def create_decoder_level(self, in_channels, out_channels):
"""Create a decoder level"""
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, 3, 1, 1),
nn.BatchNorm2d(out_channels),
nn.ReLU(True),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
)
def forward(self, x):
"""Forward pass of the model"""
# For timm ViT models
if hasattr(self.backbone, 'forward_features'):
features = self.backbone.forward_features(x)
# Reshape features based on model type
if 'vit' in self.encoder:
# Reshape transformer output to spatial features
# Exact reshape depends on the model details
h = w = int(features.shape[1]**0.5)
features = features.reshape(-1, h, w, self.embed_dim).permute(0, 3, 1, 2)
# Process through decoder
x = self.neck(features)
# Apply decoder stages
for decoder in self.decoders:
x = decoder(x)
# Final depth prediction
depth = self.depth_head(x)
return depth
else:
# Fallback for ResNet
x = self.backbone.conv1(x)
x = self.backbone.bn1(x)
x = self.backbone.relu(x)
x = self.backbone.maxpool(x)
x = self.backbone.layer1(x)
x = self.backbone.layer2(x)
x = self.backbone.layer3(x)
x = self.backbone.layer4(x)
# Process through simple decoder
x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
x = self.depth_head(x)
return x
def infer_image(self, image):
"""Process an image and return the depth map
Args:
image: A numpy image in BGR format (OpenCV) or RGB PIL Image
Returns:
depth: A numpy array containing the depth map
"""
# Convert input to tensor
if isinstance(image, np.ndarray):
# Convert BGR to RGB
if image.shape[2] == 3:
image = image[:, :, ::-1]
# Normalize
image = image.astype(np.float32) / 255.0
# HWC to CHW
image = image.transpose(2, 0, 1)
# Add batch dimension
image = torch.from_numpy(image).unsqueeze(0)
elif isinstance(image, Image.Image):
# Convert PIL image to numpy
image = np.array(image).astype(np.float32) / 255.0
# HWC to CHW
image = image.transpose(2, 0, 1)
# Add batch dimension
image = torch.from_numpy(image).unsqueeze(0)
# Move to device
image = image.to(self.device)
# Set model to eval mode
self.eval()
# Get prediction
with torch.no_grad():
depth = self.forward(image)
# Convert to numpy
depth = depth.squeeze().cpu().numpy()
return depth
def __call__(self, image):
"""Compatible interface with the pipeline API"""
if isinstance(image, Image.Image):
# Convert to numpy for processing
depth = self.infer_image(image)
# Return in the format expected by the node
return {"predicted_depth": torch.from_numpy(depth).unsqueeze(0)}
else:
# Already a tensor, process directly
self.eval()
with torch.no_grad():
depth = self.forward(image)
return {"predicted_depth": depth}
# Configure model paths
if not hasattr(folder_paths, "models_dir"):
folder_paths.models_dir = os.path.join(folder_paths.base_path, "models")
# Register depth models path - support multiple possible directory structures
DEPTH_DIR = "depth_anything"
DEPTH_ANYTHING_DIR = "depthanything"
# Check which directory structure exists
possible_paths = [
os.path.join(folder_paths.models_dir, DEPTH_DIR),
os.path.join(folder_paths.models_dir, DEPTH_ANYTHING_DIR),
os.path.join(folder_paths.models_dir, DEPTH_ANYTHING_DIR, DEPTH_DIR),
os.path.join(folder_paths.models_dir, "checkpoints", DEPTH_DIR),
os.path.join(folder_paths.models_dir, "checkpoints", DEPTH_ANYTHING_DIR),
]
# Filter to only paths that exist
existing_paths = [p for p in possible_paths if os.path.exists(p)]
if not existing_paths:
# If none exists, create the default one
existing_paths = [os.path.join(folder_paths.models_dir, DEPTH_DIR)]
os.makedirs(existing_paths[0], exist_ok=True)
logger.info(f"Created model directory: {existing_paths[0]}")
# Log all found paths for debugging
logger.info(f"Found depth model directories: {existing_paths}")
# Register all possible paths for model loading
folder_paths.folder_names_and_paths[DEPTH_DIR] = (existing_paths, folder_paths.supported_pt_extensions)
# Set primary models directory to the first available path
MODELS_DIR = existing_paths[0]
logger.info(f"Using primary models directory: {MODELS_DIR}")
# Set Hugging Face cache to the models directory to ensure models are saved there
os.environ["TRANSFORMERS_CACHE"] = MODELS_DIR
os.environ["HF_HOME"] = MODELS_DIR
# Define model configurations for direct loading
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]},
'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
}
# Define all models mentioned in the README with memory requirements
DEPTH_MODELS = {
"Depth-Anything-Small": {
"path": "LiheYoung/depth-anything-small-hf", # Correct HF path for V1
"vram_mb": 1500,
"direct_url": "https://github.com/LiheYoung/Depth-Anything/releases/download/v1.0/depth_anything_vitb14.pt",
"model_type": "v1",
"encoder": "vitb"
},
"Depth-Anything-Base": {
"path": "LiheYoung/depth-anything-base-hf", # Correct HF path for V1
"vram_mb": 2500,
"direct_url": "https://github.com/LiheYoung/Depth-Anything/releases/download/v1.0/depth_anything_vitl14.pt",
"model_type": "v1",
"encoder": "vitl"
},
"Depth-Anything-Large": {
"path": "LiheYoung/depth-anything-large-hf", # Correct HF path for V1
"vram_mb": 4000,
"direct_url": "https://github.com/LiheYoung/Depth-Anything/releases/download/v1.0/depth_anything_vitl14.pt",
"model_type": "v1",
"encoder": "vitl"
},
"Depth-Anything-V2-Small": {
"path": "depth-anything/Depth-Anything-V2-Small-hf", # Updated corrected path as shown in example
"vram_mb": 1500,
"direct_url": "https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf/resolve/main/pytorch_model.bin",
"model_type": "v2",
"encoder": "vits",
"config": MODEL_CONFIGS["vits"]
},
"Depth-Anything-V2-Base": {
"path": "depth-anything/Depth-Anything-V2-Base-hf", # Updated corrected path
"vram_mb": 2500,
"direct_url": "https://huggingface.co/depth-anything/Depth-Anything-V2-Base-hf/resolve/main/pytorch_model.bin",
"model_type": "v2",
"encoder": "vitb",
"config": MODEL_CONFIGS["vitb"]
},
# Add MiDaS models as dedicated options with direct download URLs
"MiDaS-Small": {
"path": "Intel/dpt-hybrid-midas",
"vram_mb": 1000,
"midas_type": "MiDaS_small",
"direct_url": "https://github.com/intel-isl/MiDaS/releases/download/v2_1/midas_v21_small_256.pt"
},
"MiDaS-Base": {
"path": "Intel/dpt-hybrid-midas",
"vram_mb": 1200,
"midas_type": "DPT_Hybrid",
"direct_url": "https://github.com/intel-isl/MiDaS/releases/download/v3/dpt_hybrid-midas-501f0c75.pt"
}
}
class MiDaSWrapper:
def __init__(self, model_type, device):
self.device = device
try:
# Import required libraries
import torch.nn.functional as F
# Use a more reliable approach to loading MiDaS models
if model_type == "DPT_Hybrid" or model_type == "dpt_hybrid":
# Use direct URL download for MiDaS models
midas_url = "https://github.com/intel-isl/MiDaS/releases/download/v3/dpt_hybrid-midas-501f0c75.pt"
local_path = os.path.join(MODELS_DIR, "dpt_hybrid_midas.pt")
if not os.path.exists(local_path):
logger.info(f"Downloading MiDaS model from {midas_url}")
try:
response = requests.get(midas_url, stream=True)
if response.status_code == 200:
with open(local_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
logger.info(f"Downloaded MiDaS model to {local_path}")
else:
logger.error(f"Failed to download model: {response.status_code}")
except Exception as e:
logger.error(f"Error downloading MiDaS model: {e}")
# Load pretrained model
try:
# Create a simple model architecture
from torchvision.models import resnet50
self.model = resnet50()
self.model.fc = torch.nn.Linear(2048, 1)
# Load state dict if available
if os.path.exists(local_path):
logger.info(f"Loading MiDaS model from {local_path}")
state_dict = torch.load(local_path, map_location=device)
# Convert all parameters to float
floated_state_dict = {k: v.float() for k, v in state_dict.items()}
self.model.load_state_dict(floated_state_dict)
except Exception as e:
logger.error(f"Error loading MiDaS model state dict: {e}")
# Fallback to ResNet
self.model = resnet50(pretrained=True)
self.model.fc = torch.nn.Linear(2048, 1)
else: # Other model types or fallback
from torchvision.models import resnet50
self.model = resnet50(pretrained=True)
self.model.fc = torch.nn.Linear(2048, 1)
# Ensure model parameters are float
for param in self.model.parameters():
param.data = param.data.float()
# Explicitly convert model to FloatTensor
self.model = self.model.float()
# Move model to device and set to eval mode
self.model = self.model.to(device)
self.model.eval()
except Exception as e:
logger.error(f"Failed to load MiDaS model: {e}")
logger.error(traceback.format_exc())
# Create a minimal model as absolute fallback
from torchvision.models import resnet18
self.model = resnet18(pretrained=True).float().to(device)
self.model.fc = torch.nn.Linear(512, 1).float().to(device)
self.model.eval()
def __call__(self, image):
"""Process an image and return the depth map"""
try:
# Convert PIL image to tensor for processing
if isinstance(image, Image.Image):
# Get original dimensions
original_width, original_height = image.size
# Ensure dimensions are multiple of 32 (required for some models)
# This helps prevent tensor dimension mismatches
target_height = ((original_height + 31) // 32) * 32
target_width = ((original_width + 31) // 32) * 32
# Resize to dimensions that work well with the model
img_resized = image.resize((target_width, target_height), Image.LANCZOS)
# Log resize information
if (target_width != original_width) or (target_height != original_height):
logger.info(f"Resized input from {original_width}x{original_height} to {target_width}x{target_height} (multiples of 32)")
# Convert to numpy array
img_np = np.array(img_resized).astype(np.float32) / 255.0
# Convert to tensor with proper shape (B,C,H,W)
if len(img_np.shape) == 3:
# RGB image
img_np = img_np.transpose(2, 0, 1) # (H,W,C) -> (C,H,W)
else:
# Grayscale image - add channel dimension
img_np = np.expand_dims(img_np, axis=0)
# Add batch dimension and ensure float32
input_tensor = torch.from_numpy(img_np).unsqueeze(0).float()
else:
# Already a tensor - ensure float32 by explicitly converting
# This is the key fix for the "Input type (torch.cuda.DoubleTensor) and weight type (torch.cuda.FloatTensor)" error
if image.dtype == torch.float64 or image.dtype == torch.double:
logger.info(f"Converting input tensor from {image.dtype} to torch.float32")
input_tensor = image.float() # Convert DoubleTensor to FloatTensor
else:
# Still convert to ensure it's float32
input_tensor = image.float()
# Handle tensor shape issues
# Ensure we have batch and channel dimensions
if input_tensor.dim() == 2: # [H, W]
input_tensor = input_tensor.unsqueeze(0).unsqueeze(0) # Add batch and channel dims [1, 1, H, W]
elif input_tensor.dim() == 3:
# Could be [C, H, W] or [B, H, W]
if input_tensor.shape[0] <= 3: # Likely [C, H, W]
input_tensor = input_tensor.unsqueeze(0) # Add batch dim [1, C, H, W]
else: # Likely [B, H, W]
input_tensor = input_tensor.unsqueeze(1) # Add channel dim [B, 1, H, W]
# Move to device and ensure float type
input_tensor = input_tensor.to(self.device).float()
# Log tensor shape for debugging
logger.info(f"MiDaS input tensor shape: {input_tensor.shape}, dtype: {input_tensor.dtype}")
# Log tensor info for debugging
logger.info(f"Input tensor type before inference: {input_tensor.dtype}")
# Run inference
with torch.no_grad():
# Make sure input is float32 and model weights are float32
output = self.model(input_tensor)
# Reshape to expected format
if output.dim() == 2:
# Add channel dimension if missing
output = output.unsqueeze(1)
# Resize to match input resolution
if isinstance(image, Image.Image):
w, h = image.size
# Fix tensor dimensionality mismatch by ensuring output has proper dimensions
# This fixes the "Input and output must have the same number of spatial dimensions" error
if output.dim() == 3: # Add height/width dimension if missing
output = output.unsqueeze(2)
# Ensure output has at least 4 dimensions (B,C,H,W)
while output.dim() < 4:
output = output.unsqueeze(-1)
# Log the shape for debugging
logger.info(f"Resizing output tensor from shape {output.shape} to size ({h}, {w})")
# Now interpolate with proper dimensions
output = torch.nn.functional.interpolate(
output,
size=(h, w),
mode="bicubic",
align_corners=False
)
# Use same interface as the pipeline
return {"predicted_depth": output.float()} # Ensure output is float
except Exception as e:
logger.error(f"Error in MiDaS inference: {e}")
logger.error(traceback.format_exc())
# Return a placeholder depth map
if isinstance(image, Image.Image):
w, h = image.size
dummy_tensor = torch.ones((1, 1, h, w), device=self.device, dtype=torch.float32)
else:
# Try to get shape from tensor
shape = image.shape
if len(shape) >= 3:
if shape[0] == 3: # CHW format
h, w = shape[1], shape[2]
else: # HWC format
h, w = shape[0], shape[1]
else:
h, w = 512, 512
dummy_tensor = torch.ones((1, 1, h, w), device=self.device, dtype=torch.float32)
return {"predicted_depth": dummy_tensor}
class DepthEstimationNode:
"""
ComfyUI node for depth estimation using Depth Anything models.
This node provides depth map generation from images using various Depth Anything models
with configurable post-processing options like blur, median filtering, contrast enhancement,
and gamma correction.
"""
MEDIAN_SIZES = ["3", "5", "7", "9", "11"]
def __init__(self):
self.device = None
self.depth_estimator = None
self.current_model = None
logger.info("Initialized DepthEstimationNode")
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
"""Define the input types for the node."""
return {
"required": {
"image": ("IMAGE",),
"model_name": (list(DEPTH_MODELS.keys()),),
# Ensure minimum size is enforced by the UI
"input_size": ("INT", {"default": 518, "min": 256, "max": 1024, "step": 1}),
"blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}),
# Define median_size as a dropdown with specific string values
"median_size": (cls.MEDIAN_SIZES, {"default": "5"}),
"apply_auto_contrast": ("BOOLEAN", {"default": True}),
"apply_gamma": ("BOOLEAN", {"default": True})
},
"optional": {
"force_reload": ("BOOLEAN", {"default": False}),
"force_cpu": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "estimate_depth"
CATEGORY = "depth"
def cleanup(self) -> None:
"""Clean up resources and free VRAM."""
try:
if self.depth_estimator is not None:
# Save model name before deletion for logging
model_name = self.current_model
# Delete the estimator
del self.depth_estimator
self.depth_estimator = None
self.current_model = None
# Force CUDA cache clearing
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
logger.info(f"Cleaned up model resources for {model_name}")
# Log available memory after cleanup if CUDA is available
if torch.cuda.is_available():
try:
free_mem_info = get_free_memory(get_torch_device())
# Handle return value whether it's a tuple or a single value
if isinstance(free_mem_info, tuple):
free_mem, total_mem = free_mem_info
logger.info(f"Available VRAM after cleanup: {free_mem/1024:.2f}MB of {total_mem/1024:.2f}MB")
else:
logger.info(f"Available VRAM after cleanup: {free_mem_info/1024:.2f}MB")
except Exception as e:
logger.warning(f"Error getting memory info: {e}")
except Exception as e:
logger.warning(f"Error during cleanup: {e}")
logger.debug(traceback.format_exc())
def ensure_model_loaded(self, model_name: str, force_reload: bool = False, force_cpu: bool = False) -> None:
"""
Ensures the correct model is loaded with proper VRAM management and fallback options.
Args:
model_name: The name of the model to load
force_reload: If True, reload the model even if it's already loaded
force_cpu: If True, force loading on CPU regardless of GPU availability
Raises:
RuntimeError: If the model fails to load after all fallback attempts
"""
try:
if model_name not in DEPTH_MODELS:
available_models = list(DEPTH_MODELS.keys())
if len(available_models) > 0:
fallback_model = available_models[0]
logger.warning(f"Unknown model: {model_name}. Falling back to {fallback_model}")
model_name = fallback_model
else:
raise ValueError(f"No depth models available. Please check your installation.")
model_info = DEPTH_MODELS[model_name]
# Handle model_info as string or dict
if isinstance(model_info, dict):
model_path = model_info["path"]
required_vram = model_info.get("vram_mb", 2000) * 1024 # Convert to KB
else:
model_path = model_info
required_vram = 2000 * 1024 # Default 2GB
# Only reload if needed or forced
if force_reload or self.depth_estimator is None or self.current_model != model_path:
self.cleanup()
# Set up device
if self.device is None:
self.device = get_torch_device()
logger.info(f"Loading depth model: {model_name} on {'CPU' if force_cpu else self.device}")
# Check available memory if using CUDA
if torch.cuda.is_available() and not force_cpu:
try:
free_mem_info = get_free_memory(self.device)
# Handle different return types from get_free_memory
if isinstance(free_mem_info, tuple):
free_mem, total_mem = free_mem_info
logger.info(f"Available VRAM: {free_mem/1024:.2f}MB, Required: {required_vram/1024:.2f}MB")
else:
free_mem = free_mem_info
logger.info(f"Available VRAM: {free_mem/1024:.2f}MB, Required: {required_vram/1024:.2f}MB")
total_mem = free_mem * 2 # Estimate if not available
# If not enough memory, fall back to CPU
if free_mem < required_vram:
logger.warning(f"Insufficient VRAM for {model_name} ({required_vram/1024:.1f}MB required, {free_mem/1024:.1f}MB available). Falling back to CPU.")
force_cpu = True
except Exception as mem_error:
logger.warning(f"Error checking VRAM, using CPU to be safe: {str(mem_error)}")
force_cpu = True
# Determine device type for pipeline
device_type = 'cpu' if force_cpu else ('cuda' if torch.cuda.is_available() else 'cpu')
# Use FP16 for CUDA devices to save VRAM
dtype = torch.float16 if 'cuda' in str(self.device) and not force_cpu else torch.float32
# Create a dedicated cache directory for this model
cache_dir = os.path.join(MODELS_DIR, model_name.replace("-", "_").lower())
os.makedirs(cache_dir, exist_ok=True)
# Check if we should try direct model download
direct_url = model_info.get("direct_url", None)
if direct_url:
# Determine model filename from URL
model_filename = os.path.basename(direct_url)
model_path_local = os.path.join(cache_dir, model_filename)
# Check if model already exists locally
if not os.path.exists(model_path_local):
try:
logger.info(f"Attempting to download model directly from: {direct_url}")
logger.info(f"Saving to: {model_path_local}")
# Download with progress reporting
response = requests.get(direct_url, stream=True)
total_size = int(response.headers.get('content-length', 0))
block_size = 1024 # 1 Kibibyte
if response.status_code == 200:
with open(model_path_local, 'wb') as f:
if total_size > 0:
downloaded = 0
for data in response.iter_content(block_size):
f.write(data)
downloaded += len(data)
download_pct = (downloaded / total_size) * 100
if downloaded % (5 * 1024 * 1024) == 0: # Log every 5MB
logger.info(f"Downloaded: {downloaded/1024/1024:.1f}MB of {total_size/1024/1024:.1f}MB ({download_pct:.1f}%)")
else:
f.write(response.content)
logger.info(f"Model successfully downloaded to {model_path_local}")
else:
logger.warning(f"Failed to download model from {direct_url}, status code: {response.status_code}")
except Exception as download_error:
logger.warning(f"Error downloading model: {str(download_error)}")
# List of model paths to try (original and fallback)
model_paths_to_try = [
model_path, # Original path
model_path.replace("-hf", ""), # Remove -hf suffix if it exists
model_path if "-hf" in model_path else model_path + "-hf", # Add or keep -hf suffix
# Try correct organization name for V2 models
"depth-anything/Depth-Anything-V2-Small-hf" if "v2" in model_name.lower() and "small" in model_name.lower() else model_path,
"depth-anything/Depth-Anything-V2-Base-hf" if "v2" in model_name.lower() and "base" in model_name.lower() else model_path,
# Try alternative formats
model_path.replace("LiheYoung", "depth-anything"), # Try with depth-anything organization
model_path.replace("depth-anything", "LiheYoung"), # Try with LiheYoung organization
# Fallbacks
"Intel/dpt-hybrid-midas", # Midas model as fallback
"LiheYoung/depth-anything-small", # Fallback to regular Depth Anything model
"depth-anything/Depth-Anything-Small-hf" # One more fallback
]
# Log all paths we're going to try
logger.info(f"Will try loading from these paths: {model_paths_to_try}")
# Try each model path
success = False
last_error = None
logger.info(f"Loading model with device={device_type}, dtype={dtype}")
for path in model_paths_to_try:
try:
logger.info(f"Attempting to load from: {path}")
# Try with online mode first
try:
# Add more debugging information
logger.info(f"Loading with params: model={path}, device_map={device_type}, dtype={dtype}")
# Handle specific TypeError that might occur during unpacking
try:
self.depth_estimator = pipeline(
"depth-estimation",
model=path,
cache_dir=cache_dir,
local_files_only=False, # Try online first
device_map=device_type,
torch_dtype=dtype
)
# Verify that the estimator was properly initialized
if self.depth_estimator is None:
raise RuntimeError("Pipeline initialization returned None")
# Log more info for debugging
logger.info(f"Pipeline created: {type(self.depth_estimator)}")
# Test the model with a small image to ensure it works
test_img = Image.new("RGB", (64, 64), color=(128, 128, 128))
_ = self.depth_estimator(test_img)
logger.info("Model test successful")
success = True
logger.info(f"Successfully loaded model from {path}")
break
except TypeError as type_error:
# Handle unpacking errors by printing traceback
logger.error(f"Type error when loading model: {str(type_error)}")
logger.error(f"Traceback: {traceback.format_exc()}")
# Try alternative pipeline creation approach for older transformers versions
logger.info("Trying alternative pipeline creation method...")
from transformers import AutoModelForDepthEstimation, AutoImageProcessor
# Load model components separately to avoid unpacking issues
try:
processor = AutoImageProcessor.from_pretrained(path, cache_dir=cache_dir)
model = AutoModelForDepthEstimation.from_pretrained(path, cache_dir=cache_dir)
# Move model to correct device if needed
if not force_cpu and 'cuda' in device_type:
model = model.to(self.device)
# Create a custom pipeline class that wraps these components
class CustomDepthEstimator:
def __init__(self, model, processor):
self.model = model
self.processor = processor
def __call__(self, image):
# Process image and run model
inputs = self.processor(images=image, return_tensors="pt")
if not force_cpu and 'cuda' in device_type:
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = self.model(**inputs)
# Format output like the pipeline would
return {"predicted_depth": outputs.predicted_depth}
self.depth_estimator = CustomDepthEstimator(model, processor)
# Test the custom pipeline
test_img = Image.new("RGB", (64, 64), color=(128, 128, 128))
_ = self.depth_estimator(test_img)
success = True
logger.info(f"Successfully loaded model using custom pipeline")
break
except Exception as custom_error:
logger.error(f"Custom pipeline creation failed: {str(custom_error)}")
raise
except Exception as online_error:
logger.warning(f"Online loading failed for {path}: {str(online_error)}")
logger.debug(f"Error traceback: {traceback.format_exc()}")
# Try with local_files_only if online fails
try:
# Add more verbose logging
logger.info(f"Trying local cache with model={path}")
self.depth_estimator = pipeline(
"depth-estimation",
model=path,
cache_dir=cache_dir,
local_files_only=True, # Try local only as fallback
device_map=device_type,
torch_dtype=dtype
)
# Verify pipeline initialization success
if self.depth_estimator is None:
raise RuntimeError("Local pipeline initialization returned None")
# Test the model
test_img = Image.new("RGB", (64, 64), color=(128, 128, 128))
_ = self.depth_estimator(test_img)
success = True
logger.info(f"Successfully loaded model from local cache: {path}")
break
except Exception as local_error:
last_error = local_error
logger.warning(f"Local loading failed for {path}: {str(local_error)}")
logger.debug(f"Error traceback: {traceback.format_exc()}")
continue
except Exception as path_error:
last_error = path_error
logger.warning(f"Failed to load model from {path}: {str(path_error)}")
continue
# Prioritize direct model loading for V2 models and as fallback for other models
if not success:
logger.info("Transformers pipeline attempts failed, trying direct model loading with explicit configurations...")
# Try the direct loading approach
direct_model = self.load_model_direct(model_name, model_info, force_cpu)
if direct_model is not None:
self.depth_estimator = direct_model
success = True
logger.info(f"Successfully loaded model using direct loading approach")
else:
logger.error("Direct model loading also failed")
# If all attempts failed so far, try MiDaS as a final fallback
if not success:
try:
logger.info("Attempting to load MiDaS model as final fallback...")
midas_model = MiDaSWrapper("dpt_hybrid", self.device)
# Test the model
test_img = Image.new("RGB", (64, 64), color=(128, 128, 128))
_ = midas_model(test_img)
self.depth_estimator = midas_model
success = True
logger.info("Successfully loaded MiDaS fallback model")
except Exception as midas_error:
logger.error(f"MiDaS fallback also failed: {str(midas_error)}")
# If all attempts failed, try a different model
if model_name != "Depth-Anything-V2-Small" and "Depth-Anything-V2-Small" in DEPTH_MODELS:
logger.warning(f"Failed to load {model_name}, trying Depth-Anything-V2-Small as fallback")
try:
# Increase chances of success with CPU
return self.ensure_model_loaded("Depth-Anything-V2-Small", True, True)
except Exception as fallback_error:
logger.error(f"Fallback model also failed: {str(fallback_error)}")
# If still failing, show helpful message with instructions
if not success:
# Show all model directories for debugging
all_model_dirs = "\n".join(existing_paths)
# Check if the error is related to GPU issues
gpu_related = False
auth_related = False
tensor_related = False
error_str = str(last_error).lower()
if "cuda" in error_str or "gpu" in error_str or "vram" in error_str:
gpu_related = True
if "authentication" in error_str or "unauthorized" in error_str or "401" in error_str:
auth_related = True
if "tensor" in error_str or "dimension" in error_str or "shape" in error_str:
tensor_related = True
# Create a targeted error message based on the error type
if auth_related:
error_solution = """
AUTHENTICATION ERROR: The model couldn't be downloaded due to Hugging Face authentication requirements.
SOLUTION:
1. Use force_cpu=True in the node settings (this will use the MiDaS fallback model)
2. Download the model manually using one of these direct links that don't require authentication:
- https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf/resolve/main/pytorch_model.bin
- https://github.com/LiheYoung/Depth-Anything/releases/download/v2.0/depth_anything_v2_small.pt
- https://huggingface.co/ckpt/depth-anything-v2/resolve/main/depth_anything_v2_small.pt
Save the file to one of these directories:
{all_model_dirs}
"""
elif gpu_related:
error_solution = """
GPU ERROR: The model failed to load on your GPU.
SOLUTION:
1. Use force_cpu=True to use CPU processing instead
2. Reduce input_size parameter to 384 to reduce memory requirements
3. Try a smaller model like MiDaS-Small instead
4. Ensure you have the latest GPU drivers installed
"""
elif tensor_related:
error_solution = """
TENSOR DIMENSION ERROR: There was a problem with tensor shapes during model processing.
SOLUTION:
1. Use force_cpu=True to use CPU processing instead (more stable)
2. Set input_size to a multiple of 32 (e.g. 384, 512)
3. Try processing the image at a different resolution
4. Try a different model like MiDaS-Small
"""
else:
error_solution = f"""
Failed to load model {model_name} after trying multiple sources.
GENERAL SOLUTIONS:
1. Download the model manually using one of these direct URLs:
- Depth-Anything-V2-Small: https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf/resolve/main/pytorch_model.bin
- MiDaS Base: https://github.com/intel-isl/MiDaS/releases/download/v3/dpt_hybrid-midas-501f0c75.pt
2. Try using force_cpu=True in node settings
3. Try a different model version
4. Reduce input_size parameter to a smaller value like 384
"""
error_msg = f"""
MODEL LOADING ERROR: {str(last_error)}
{error_solution}
SEARCHED DIRECTORIES:
{all_model_dirs}
"""
logger.error(error_msg)
raise RuntimeError(error_msg)
# Ensure model is on the correct device
if not force_cpu and hasattr(self.depth_estimator, 'model'):
self.depth_estimator.model = self.depth_estimator.model.to(self.device)
self.current_model = model_path
except Exception as e:
self.cleanup()
error_msg = f"Failed to load model {model_name}: {str(e)}"
logger.error(error_msg)
logger.debug(traceback.format_exc())
raise RuntimeError(error_msg)
def load_model_direct(self, model_name, model_info, force_cpu=False):
"""
Directly loads a depth model without using transformers pipeline.
This is a fallback method when the normal pipeline loading fails.
Args:
model_name: Name of the model to load
model_info: Dictionary with model information
force_cpu: Whether to force CPU usage
Returns:
A depth estimation model that implements the __call__ interface
"""
try:
logger.info(f"Attempting direct model loading for {model_name}")
# Determine device
device_type = 'cpu' if force_cpu else ('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device(device_type)
# Look in all possible model directories
# This is important to support various directory structures
model_found = False
model_path_local = None
# Make a unique model cache directory for this specific model
model_subfolder = model_name.replace("-", "_").lower()
# Check all possible locations for the model file
for base_path in existing_paths:
# Try different possible locations and filename patterns
possible_model_locations = [
# Direct downloads in the model directory
os.path.join(base_path, model_subfolder),
# Using the full HF directory structure
os.path.join(base_path, model_info.get("path", "").replace("/", "_")),
# Directly in base directory
base_path,
]
# Add directory structure with model configs if V2
if model_info.get("model_type") == "v2":
v2_path = os.path.join(base_path, "v2")
possible_model_locations.append(v2_path)
possible_model_locations.append(os.path.join(v2_path, model_subfolder))
# Try all locations
logger.info(f"Searching for existing model in these directories: {possible_model_locations}")
for location in possible_model_locations:
# Check for model file with various naming patterns
if os.path.exists(location):
# Check for common filenames
for filename in ["pytorch_model.bin", "model.pt", "model.pth",
f"{model_subfolder}.pt", f"{model_subfolder}.bin"]:
file_path = os.path.join(location, filename)
if os.path.exists(file_path):
logger.info(f"Found existing model file: {file_path}")
model_path_local = file_path
model_found = True
break
if model_found:
break
if model_found:
break
# If model not found, use the first directory for downloading
cache_dir = os.path.join(existing_paths[0], model_subfolder)
os.makedirs(cache_dir, exist_ok=True)
# Get model configuration
model_type = model_info.get("model_type", "v1")
encoder = model_info.get("encoder", "vits")
config = model_info.get("config", MODEL_CONFIGS.get(encoder, MODEL_CONFIGS["vits"]))
# Step 1: If model not found locally, download it
# List of alternative URLs that don't require authentication
alternative_urls = {
"Depth-Anything-V2-Small": [
"https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf/resolve/main/pytorch_model.bin",
"https://github.com/LiheYoung/Depth-Anything/releases/download/v2.0/depth_anything_v2_small.pt",
"https://huggingface.co/ckpt/depth-anything-v2/resolve/main/depth_anything_v2_small.pt"
],
"Depth-Anything-V2-Base": [
"https://huggingface.co/depth-anything/Depth-Anything-V2-Base-hf/resolve/main/pytorch_model.bin",
"https://github.com/LiheYoung/Depth-Anything/releases/download/v2.0/depth_anything_v2_base.pt",
"https://huggingface.co/ckpt/depth-anything-v2/resolve/main/depth_anything_v2_base.pt"
],
"MiDaS-Base": [
"https://github.com/intel-isl/MiDaS/releases/download/v3/dpt_hybrid-midas-501f0c75.pt"
]
}
# Get primary URL from model_info
direct_url = model_info.get("direct_url")
# Add alternative URLs to try if the main one fails
urls_to_try = [direct_url] if direct_url else []
# Add alternative URLs for this model if available
if model_name in alternative_urls:
urls_to_try.extend(alternative_urls[model_name])
# Try downloading the model if not found locally
if not model_found and urls_to_try:
# Try each URL in sequence until one works
for url in urls_to_try:
if not url:
continue
try:
model_filename = os.path.basename(url)
model_path_local = os.path.join(cache_dir, model_filename)
if os.path.exists(model_path_local):
logger.info(f"Model already exists at {model_path_local}")
model_found = True
break
logger.info(f"Attempting to download model from {url} to {model_path_local}")
# Create parent directory if needed
os.makedirs(os.path.dirname(model_path_local), exist_ok=True)
# Try different download methods
download_success = False
# First try wget (more reliable for large files)
try:
logger.info(f"Downloading with wget: {url}")
wget.download(url, out=model_path_local)
logger.info(f"Downloaded model weights to {model_path_local}")
download_success = True
except Exception as wget_error:
logger.warning(f"wget download failed: {str(wget_error)}")
# Fallback to requests
try:
logger.info(f"Downloading with requests: {url}")
response = requests.get(url, stream=True)
if response.status_code == 200:
total_size = int(response.headers.get('content-length', 0))
logger.info(f"File size: {total_size/1024/1024:.1f} MB")
with open(model_path_local, 'wb') as f:
downloaded = 0
for data in response.iter_content(1024 * 1024): # 1MB chunks
f.write(data)
downloaded += len(data)
if total_size > 0 and downloaded % (10 * 1024 * 1024) == 0: # Log every 10MB
progress = (downloaded / total_size) * 100
logger.info(f"Downloaded {downloaded/1024/1024:.1f}MB of {total_size/1024/1024:.1f}MB ({progress:.1f}%)")
logger.info(f"Download complete: {model_path_local}")
download_success = True
else:
logger.warning(f"Failed to download from {url}: HTTP status {response.status_code}")
except Exception as req_error:
logger.warning(f"Requests download failed: {str(req_error)}")
# Try urllib as last resort
if not download_success:
try:
logger.info(f"Downloading with urllib: {url}")
urllib.request.urlretrieve(url, model_path_local)
logger.info(f"Downloaded model weights to {model_path_local}")
download_success = True
except Exception as urllib_error:
logger.warning(f"urllib download failed: {str(urllib_error)}")
# Check if download succeeded
if download_success and os.path.exists(model_path_local) and os.path.getsize(model_path_local) > 0:
logger.info(f"Successfully downloaded model to {model_path_local}")
model_found = True
break
else:
logger.warning(f"Download appeared to succeed but file is empty or missing")
# Try to remove the failed download
if os.path.exists(model_path_local):
try:
os.remove(model_path_local)
except:
pass
except Exception as dl_error:
logger.warning(f"Error downloading from {url}: {str(dl_error)}")
continue
if not model_found:
logger.error("All download attempts failed")
# Step 2: Create and load the appropriate model if found
if model_found and model_path_local and os.path.exists(model_path_local):
logger.info(f"Found model file at: {model_path_local}")
# Handle V2 models with DepthAnythingV2 implementation
if model_type == "v2" and TIMM_AVAILABLE:
try:
logger.info(f"Loading as DepthAnythingV2 model with config: {config}")
# Create model with the appropriate configuration
model = DepthAnythingV2(**config)
# Load weights from checkpoint
logger.info(f"Loading weights from {model_path_local}")
state_dict = torch.load(model_path_local, map_location=device)
# Convert state dict to float32 if needed
if any(v.dtype == torch.float64 for v in state_dict.values() if hasattr(v, 'dtype')):
logger.info("Converting state dict from float64 to float32")
state_dict = {k: v.float() if hasattr(v, 'dtype') else v for k, v in state_dict.items()}
# Attempt to load the state dict (handles different formats)
try:
if "model" in state_dict:
model.load_state_dict(state_dict["model"])
else:
model.load_state_dict(state_dict)
except Exception as e:
logger.warning(f"Error loading state dict: {str(e)}")
logger.warning("Trying to load with strict=False")
if "model" in state_dict:
model.load_state_dict(state_dict["model"], strict=False)
else:
model.load_state_dict(state_dict, strict=False)
# Move model to the correct device and ensure float32
model = model.float().to(device)
model.device = device
model.eval()
# Test the model
logger.info("Testing model with sample image")
test_img = Image.new("RGB", (64, 64), color=(128, 128, 128))
try:
_ = model(test_img)
logger.info("DepthAnythingV2 model loaded and tested successfully")
return model
except Exception as test_error:
logger.error(f"Error during model test: {str(test_error)}")
logger.debug(traceback.format_exc())
except Exception as e:
logger.error(f"Error loading DepthAnythingV2: {str(e)}")
logger.debug(traceback.format_exc())
# Fallback to MiDaS model if V2 loading failed or for V1 models
try:
logger.info("Falling back to MiDaS model")
# Determine the appropriate MiDaS model type
midas_model_type = "dpt_hybrid"
if "large" in model_name.lower():
midas_model_type = "dpt_large"
elif "small" in model_name.lower():
midas_model_type = "midas_v21_small"
# Create and test the MiDaS model
midas_model = MiDaSWrapper(midas_model_type, device)
# Test with a small image
test_img = Image.new("RGB", (64, 64), color=(128, 128, 128))
_ = midas_model(test_img)
logger.info("MiDaS model loaded and tested successfully")
return midas_model
except Exception as e:
logger.error(f"Error loading MiDaS: {str(e)}")
logger.debug(traceback.format_exc())
# If all else fails, return None
return None
except Exception as e:
logger.error(f"Direct model loading failed: {str(e)}")
logger.debug(traceback.format_exc())
return None
def process_image(self, image: Union[torch.Tensor, np.ndarray], input_size: int = 518) -> Image.Image:
"""
Converts input image to proper format for depth estimation and resizes it.
Args:
image: Input image as tensor or numpy array
input_size: Target size for the longest dimension of the image
Returns:
PIL Image ready for depth estimation
"""
try:
# Validate input_size
if input_size < 256:
logger.warning(f"Input size {input_size} is too small, using 256 instead")
input_size = 256
elif input_size > 1024:
logger.warning(f"Input size {input_size} is too large, using 1024 instead")
input_size = 1024
# Convert tensor to numpy array
if torch.is_tensor(image):
# Check tensor dtype and convert to float32 if needed
if image.dtype == torch.float64 or image.dtype == torch.double:
logger.info(f"Converting input tensor from {image.dtype} to torch.float32")
image = image.float() # Convert DoubleTensor to FloatTensor
# Check for NaN values in tensor
if torch.isnan(image).any():
logger.warning("Input tensor contains NaN values. Replacing with zeros.")
image = torch.nan_to_num(image, nan=0.0)
# Get first image from batch and convert to numpy
image_np = (image.cpu().numpy()[0] * 255).astype(np.uint8)
else:
# Check for NaN values in numpy array
if np.isnan(image).any():
logger.warning("Input array contains NaN values. Replacing with zeros.")
image = np.nan_to_num(image, nan=0.0)
# Convert float64 to float32 if needed
if image.dtype == np.float64:
logger.info("Converting numpy array from float64 to float32")
image = image.astype(np.float32)
image_np = (image * 255).astype(np.uint8)
# Handle different channel configurations
if len(image_np.shape) == 3:
if image_np.shape[-1] == 4: # Handle RGBA images
image_np = image_np[..., :3]
elif len(image_np.shape) == 2: # Handle grayscale images
image_np = np.stack([image_np] * 3, axis=-1)
# Convert to PIL image
pil_image = Image.fromarray(image_np)
# Resize the image while preserving aspect ratio
width, height = pil_image.size
# Determine which dimension to scale to input_size
if width > height:
new_width = input_size
new_height = int(height * (new_width / width))
else:
new_height = input_size
new_width = int(width * (new_height / height))
# Resize the image with antialiasing
resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
logger.info(f"Resized image from {width}x{height} to {new_width}x{new_height}")
return resized_image
except Exception as e:
logger.error(f"Error processing image: {str(e)}")
logger.debug(traceback.format_exc())
# Return a placeholder image on error
return Image.new('RGB', (512, 512), (128, 128, 128))
def _create_error_image(self, input_image=None):
"""Create an error image placeholder based on input image if possible."""
try:
if input_image is not None and isinstance(input_image, torch.Tensor) and input_image.shape[0] > 0:
# Check tensor type - if it's float64, log it for debugging
if input_image.dtype == torch.float64 or input_image.dtype == torch.double:
logger.info(f"Input tensor for error image is {input_image.dtype}, will create float32 error image")
# Create gray error image with same dimensions as input
# Ensure tensor has the right shape for error display (BHWC)
if input_image.ndim == 4:
if input_image.shape[-1] != 3: # if not BHWC format
if input_image.shape[1] == 3: # if BCHW format
# Extract height and width from BCHW
h, w = input_image.shape[2], input_image.shape[3]
else:
# Default to dimensions from input
h, w = input_image.shape[2], input_image.shape[3]
else:
# Already in BHWC format
h, w = input_image.shape[1], input_image.shape[2]
else:
# Unexpected shape, use default
return self._create_basic_error_image()
# Make sure dimensions aren't too small
if h <= 1 or w <= 1:
logger.warning(f"Input has invalid dimensions {h}x{w}, using default error image")
return self._create_basic_error_image()
# Gray background with slight red tint to indicate error - explicitly use float32
placeholder = torch.ones((1, h, w, 3), dtype=torch.float32) * torch.tensor([0.5, 0.4, 0.4], dtype=torch.float32)
if self.device is not None:
placeholder = placeholder.to(self.device)
# Verify the placeholder is float32
if placeholder.dtype != torch.float32:
logger.warning(f"Error image has unexpected dtype {placeholder.dtype}, converting to float32")
placeholder = placeholder.float()
return placeholder
else:
return self._create_basic_error_image()
except Exception as e:
logger.error(f"Error creating error image: {str(e)}")
return self._create_basic_error_image()
def _create_basic_error_image(self):
"""Create a basic error image when no input dimensions are available."""
# Standard size error image (512x512)
h, w = 512, 512
# Gray background with slight red tint to indicate error - explicitly use float32
placeholder = torch.ones((1, h, w, 3), dtype=torch.float32) * torch.tensor([0.5, 0.4, 0.4], dtype=torch.float32)
if self.device is not None:
placeholder = placeholder.to(self.device)
# Double-check that we're returning a float32 tensor
if placeholder.dtype != torch.float32:
placeholder = placeholder.float()
return placeholder
def _add_error_text_to_image(self, image_tensor, error_text):
"""Add error text to the image tensor for visual feedback."""
try:
# Convert tensor to PIL for text rendering
if image_tensor is None:
return
temp_img = self._tensor_to_pil(image_tensor)
# Draw error text
draw = ImageDraw.Draw(temp_img)
# Try to get a font, fall back to default if needed
try:
font = ImageFont.truetype("arial.ttf", 20)
except:
font = ImageFont.load_default()
# Split text into multiple lines if too long
lines = []
words = error_text.split()
current_line = words[0] if words else "Error"
for word in words[1:]:
if len(current_line + " " + word) < 50:
current_line += " " + word
else:
lines.append(current_line)
current_line = word
lines.append(current_line)
# Draw title
draw.text((10, 10), "Depth Estimation Error", fill=(255, 50, 50), font=font)
# Draw error message
y_position = 40
for line in lines:
draw.text((10, y_position), line, fill=(255, 255, 255), font=font)
y_position += 25
# Convert back to tensor
result = self._pil_to_tensor(temp_img)
# Copy to original tensor if shapes match
if image_tensor.shape == result.shape:
image_tensor.copy_(result)
return image_tensor
except Exception as e:
logger.error(f"Error adding text to error image: {e}")
return image_tensor
def _tensor_to_pil(self, tensor):
"""Convert a tensor to PIL Image."""
if tensor.shape[0] == 1: # Batch size 1
img_np = (tensor[0].cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(img_np)
return Image.new('RGB', (512, 512), color=(128, 100, 100))
def _pil_to_tensor(self, pil_img):
"""Convert PIL Image back to tensor."""
img_np = np.array(pil_img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np).unsqueeze(0)
if self.device is not None:
tensor = tensor.to(self.device)
return tensor
def estimate_depth(self,
image: torch.Tensor,
model_name: str,
input_size: int = 518,
blur_radius: float = 2.0,
median_size: str = "5",
apply_auto_contrast: bool = True,
apply_gamma: bool = True,
force_reload: bool = False,
force_cpu: bool = False) -> Tuple[torch.Tensor]:
"""
Estimates depth from input image with error handling and cleanup.
Args:
image: Input image tensor
model_name: Name of the depth model to use
input_size: Target size for the longest dimension of the image (between 256 and 1024)
blur_radius: Gaussian blur radius for smoothing
median_size: Size of median filter for noise reduction
apply_auto_contrast: Whether to enhance contrast automatically
apply_gamma: Whether to apply gamma correction
force_reload: Whether to force reload the model
force_cpu: Whether to force using CPU for inference
Returns:
Tuple containing depth map tensor
"""
error_image = None
start_time = time.time()
try:
# Validate inputs
if image is None or image.numel() == 0:
raise ValueError("Empty or null input image")
if image.ndim != 4:
raise ValueError(f"Expected 4D tensor for image, got {image.ndim}D.")
# Check for DoubleTensor and convert to FloatTensor if needed
if image.dtype == torch.float64 or image.dtype == torch.double:
logger.info(f"Converting input tensor from {image.dtype} to torch.float32 in estimate_depth")
image = image.float() # This is crucial for fixing the tensor type mismatch
# Create error image placeholder based on input dimensions
error_image = self._create_error_image(image)
if torch.isnan(image).any():
logger.warning("Input image contains NaN values. These will be replaced.")
image = torch.nan_to_num(image, nan=0.0)
# Handle case where median_size is passed as a boolean or other type
if isinstance(median_size, bool) or median_size is True or median_size == 'True':
logger.warning(f"median_size was passed as boolean: {median_size}. Defaulting to 5")
median_size = "5"
elif not isinstance(median_size, str) or median_size not in self.MEDIAN_SIZES:
logger.warning(f"Invalid median_size: {median_size}. Defaulting to 5")
median_size = "5"
# Make sure it's one of the allowed values before any processing
if median_size not in self.MEDIAN_SIZES:
logger.warning(f"median_size '{median_size}' not in allowed values {self.MEDIAN_SIZES}, defaulting to 5")
median_size = "5"
# Load model with fallback strategy - wrapped in try-except
try:
self.ensure_model_loaded(model_name, force_reload, force_cpu)
except Exception as model_error:
# Special handling for model loading errors - common issue
error_msg = f"Failed to load model '{model_name}': {str(model_error)}"
logger.error(error_msg)
# Add error text to error image
self._add_error_text_to_image(error_image, f"Model Error: {str(model_error)[:100]}...")
return (error_image,)
# Process input image with resizing
try:
# Ensure input_size is valid
# Add more strict validation to handle edge cases
if not isinstance(input_size, int):
logger.warning(f"Input size {input_size} is not an integer, using 518 instead")
input_size = 518
# Fix input_size if it's too small
if input_size < 256:
logger.warning(f"Input size {input_size} is too small, using 256 instead")
input_size = 256
elif input_size > 1024:
logger.warning(f"Input size {input_size} is too large, using 1024 instead")
input_size = 1024
# Log tensor type for debugging
logger.info(f"Input tensor type before processing: {image.dtype}")
pil_image = self.process_image(image, input_size)
except Exception as img_error:
logger.error(f"Image processing error: {str(img_error)}")
self._add_error_text_to_image(error_image, f"Image Error: {str(img_error)[:100]}...")
return (error_image,)
# Perform depth estimation with error catching
try:
with torch.inference_mode():
# Log tensor info before depth estimation
logger.info(f"Calling depth estimator with PIL image of size {pil_image.size}")
depth_result = self.depth_estimator(pil_image)
# Convert output to float32 if needed
predicted_depth = depth_result["predicted_depth"]
if predicted_depth.dtype != torch.float32:
logger.info(f"Converting output from {predicted_depth.dtype} to float32")
predicted_depth = predicted_depth.float()
depth_map = predicted_depth.squeeze().cpu().numpy()
except RuntimeError as rt_error:
# Check for tensor type mismatch errors
error_msg = str(rt_error)
if "Input type" in error_msg and "weight type" in error_msg:
# This is the specific error we're trying to fix
logger.error(f"Tensor type mismatch error: {error_msg}")
# Try to fall back to CPU with explicit float conversion
logger.info("Attempting to fall back to CPU with explicit float conversion")
try:
# Create a copy of the image tensor with explicit float32 type
float_image = image.float().cpu() # Move to CPU and convert to float
return self.estimate_depth(
float_image, model_name, input_size, blur_radius, median_size,
apply_auto_contrast, apply_gamma, True, True
)
except Exception as float_fallback_error:
logger.error(f"Float fallback also failed: {str(float_fallback_error)}")
# Check specifically for CUDA out-of-memory errors
elif "CUDA out of memory" in error_msg:
error_msg = (
f"CUDA out of memory while processing depth map. "
f"Try using a smaller model or reducing image size."
)
logger.error(error_msg)
# Try to fall back to CPU if we hit OOM
if not force_cpu:
logger.info("Attempting to fall back to CPU due to CUDA OOM error")
try:
return self.estimate_depth(
image, model_name, input_size, blur_radius, median_size,
apply_auto_contrast, apply_gamma, True, True
)
except Exception as cpu_fallback_error:
logger.error(f"CPU fallback also failed: {str(cpu_fallback_error)}")
self._add_error_text_to_image(error_image, "CUDA Out of Memory. Try a smaller model.")
return (error_image,)
else:
# Other runtime errors
error_msg = f"Runtime error during depth estimation: {str(rt_error)}"
logger.error(error_msg)
logger.debug(traceback.format_exc())
self._add_error_text_to_image(error_image, f"Runtime Error: {str(rt_error)[:100]}...")
return (error_image,)
except Exception as e:
# General exceptions
error_msg = f"Depth estimation failed: {str(e)}"
logger.error(error_msg)
logger.debug(traceback.format_exc())
self._add_error_text_to_image(error_image, f"Error: {str(e)[:100]}...")
return (error_image,)
# Check for NaN values in depth map
if np.isnan(depth_map).any():
logger.warning("Depth map contains NaN values. Replacing with zeros.")
depth_map = np.nan_to_num(depth_map, nan=0.0)
# Continue with the normal depth map processing
try:
# Normalize depth values
depth_min, depth_max = depth_map.min(), depth_map.max()
if depth_max > depth_min:
depth_map = ((depth_map - depth_min) / (depth_max - depth_min) * 255.0)
depth_map = depth_map.astype(np.uint8)
# Create PIL image explicitly with L mode (grayscale)
depth_pil = Image.fromarray(depth_map, mode='L')
# Apply post-processing
if blur_radius > 0:
depth_pil = depth_pil.filter(ImageFilter.GaussianBlur(radius=blur_radius))
if int(median_size) > 0:
depth_pil = depth_pil.filter(ImageFilter.MedianFilter(size=int(median_size)))
if apply_auto_contrast:
depth_pil = ImageOps.autocontrast(depth_pil)
if apply_gamma:
depth_array = np.array(depth_pil).astype(np.float32) / 255.0
mean_luminance = np.mean(depth_array)
if mean_luminance > 0:
gamma = np.log(0.5) / np.log(mean_luminance)
# Use direct numpy operations for gamma correction
corrected = np.power(depth_array, 1.0/gamma) * 255.0
depth_pil = Image.fromarray(corrected.astype(np.uint8), mode='L')
# Fix the tensor conversion:
depth_array = np.array(depth_pil).astype(np.float32) / 255.0
# Check if depth_array has proper dimensions and isn't just a thin line
h, w = depth_array.shape
if h <= 1 or w <= 1:
logger.error(f"Invalid depth map dimensions: {h}x{w}, using error image instead")
if error_image is not None:
self._add_error_text_to_image(error_image, "Invalid depth map dimensions (thin line)")
return (error_image,)
else:
# Create new error image if one doesn't exist
error_image = self._create_basic_error_image()
self._add_error_text_to_image(error_image, "Invalid depth map dimensions (thin line)")
return (error_image,)
# Make sure we preserve proper dimensions - this is the crucial fix
logger.info(f"Depth map dimensions: {h}x{w}")
depth_rgb = np.stack([depth_array] * 3, axis=-1) # Shape becomes (h, w, 3)
# Convert to tensor and add batch dimension, ensuring float32 type
depth_tensor = torch.from_numpy(depth_rgb).unsqueeze(0).float() # Shape becomes (1, h, w, 3)
if self.device is not None and not force_cpu:
depth_tensor = depth_tensor.to(self.device)
# Make sure it's normalized in [0, 1] range
if depth_tensor.max() > 1.0:
depth_tensor = depth_tensor / 255.0
# Debug: log tensor shape and type
logger.info(f"Output depth tensor shape: {depth_tensor.shape}, dtype: {depth_tensor.dtype}")
processing_time = time.time() - start_time
logger.info(f"Depth processing completed in {processing_time:.2f} seconds")
return (depth_tensor,)
except Exception as post_error:
error_msg = f"Error during depth map post-processing: {str(post_error)}"
logger.error(error_msg)
logger.debug(traceback.format_exc())
self._add_error_text_to_image(error_image, f"Post-processing Error: {str(post_error)[:100]}...")
return (error_image,)
except Exception as e:
# Catch-all for any other exceptions
error_msg = f"Depth estimation failed: {str(e)}"
logger.error(error_msg)
logger.debug(traceback.format_exc())
# If error_image hasn't been created yet, create a basic one
if error_image is None:
error_image = self._create_basic_error_image()
self._add_error_text_to_image(error_image, f"Unexpected Error: {str(e)[:100]}...")
return (error_image,)
finally:
# Always clean up resources
torch.cuda.empty_cache()
gc.collect()
def gamma_correction(self, img: Image.Image, gamma: float = 1.0) -> Image.Image:
"""Applies gamma correction to the image."""
# Convert to numpy array
img_array = np.array(img)
# Apply gamma correction directly with numpy
corrected = np.power(img_array.astype(np.float32) / 255.0, 1.0/gamma) * 255.0
# Ensure uint8 type and create image with explicit mode
return Image.fromarray(corrected.astype(np.uint8), mode='L')
# Node registration
NODE_CLASS_MAPPINGS = {
"DepthEstimationNode": DepthEstimationNode
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DepthEstimationNode": "Depth Estimation (V2)"
}
Reference in New Issue View Git Blame Copy Permalink