MarigoldDepthEstimation_v2¶
Documentation¶
- Class name:
MarigoldDepthEstimation_v2
- Category:
Marigold
- Output node:
False
The MarigoldDepthEstimation_v2 node is designed for diffusion-based monocular depth estimation, leveraging the Marigold model to generate depth maps from single images. This node focuses on enhancing the accuracy of depth predictions through iterative processing and ensemble techniques, offering customizable parameters for fine-tuning the depth estimation process.
Input types¶
Required¶
marigold_model
- Selects the specific Marigold model to be used for depth estimation, affecting the depth map's accuracy and appearance.
- Comfy dtype:
MARIGOLDMODEL
- Python dtype:
str
image
- The input image for which the depth map is to be estimated, serving as the basis for the depth estimation process.
- Comfy dtype:
IMAGE
- Python dtype:
torch.Tensor
seed
- Sets the seed for random number generation, ensuring reproducibility of the depth estimation.
- Comfy dtype:
INT
- Python dtype:
int
denoise_steps
- Specifies the number of steps per depth map, allowing users to balance between accuracy and processing time.
- Comfy dtype:
INT
- Python dtype:
int
ensemble_size
- Determines the number of depth maps to be ensembled into a single output, enhancing the depth estimation's accuracy.
- Comfy dtype:
INT
- Python dtype:
int
processing_resolution
- Defines the resolution at which the depth estimation process is performed, impacting the detail and quality of the depth map.
- Comfy dtype:
INT
- Python dtype:
int
scheduler
- Chooses the scheduler for the depth estimation process, influencing the final depth map's characteristics.
- Comfy dtype:
COMBO[STRING]
- Python dtype:
str
use_taesd_vae
- Determines whether to use the TAESD VAE model for depth estimation, potentially improving the quality of the depth map.
- Comfy dtype:
BOOLEAN
- Python dtype:
bool
Output types¶
image
- Comfy dtype:
IMAGE
- Produces a depth map image, representing the estimated depth of the scene from a single input image.
- Python dtype:
torch.Tensor
- Comfy dtype:
Usage tips¶
- Infra type:
GPU
- Common nodes: unknown
Source code¶
class MarigoldDepthEstimation_v2:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"marigold_model": ("MARIGOLDMODEL",),
"image": ("IMAGE", ),
"seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}),
"denoise_steps": ("INT", {"default": 4, "min": 1, "max": 4096, "step": 1}),
"ensemble_size": ("INT", {"default": 3, "min": 1, "max": 4096, "step": 1}),
"processing_resolution": ("INT", {"default": 768, "min": 64, "max": 4096, "step": 8}),
"scheduler": (
["DDIMScheduler", "LCMScheduler",],
{
"default": 'LCMScheduler'
}),
"use_taesd_vae": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES =("image",)
FUNCTION = "process"
CATEGORY = "Marigold"
DESCRIPTION = """
Diffusion-based monocular depth estimation:
https://github.com/prs-eth/Marigold
Uses Diffusers 0.28.0 Marigold pipelines.
"""
def process(self, marigold_model, image, seed, denoise_steps, processing_resolution, ensemble_size, scheduler, use_taesd_vae):
try:
from diffusers import AutoencoderTiny
except:
raise Exception("diffusers==0.28 is required for v2 nodes")
batch_size = image.shape[0]
device = model_management.get_torch_device()
torch.manual_seed(seed)
image = image.permute(0, 3, 1, 2).to(device)
pipeline = marigold_model['pipeline']
pred_type = marigold_model['modeltype']
if use_taesd_vae:
pipeline.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd", torch_dtype=torch.float16).to(device)
pbar = comfy.utils.ProgressBar(batch_size)
scheduler_kwargs = {
DDIMScheduler: {
"num_inference_steps": denoise_steps,
"ensemble_size": ensemble_size,
},
LCMScheduler: {
"num_inference_steps": denoise_steps,
"ensemble_size": ensemble_size,
},
}
if scheduler == 'DDIMScheduler':
pipe_kwargs = scheduler_kwargs[DDIMScheduler]
elif scheduler == 'LCMScheduler':
pipe_kwargs = scheduler_kwargs[LCMScheduler]
generator = torch.Generator(device).manual_seed(seed)
processed_out = []
for i in range(batch_size):
processed = pipeline(
image[i],
output_type = "pt",
generator = generator,
processing_resolution = processing_resolution,
**pipe_kwargs
)
pbar.update(1)
if pred_type == "normals":
normals = pipeline.image_processor.visualize_normals(processed.prediction)
normals_tensor = transforms.ToTensor()(normals[0])
processed_out.append(normals_tensor)
else:
processed_out.append(processed[0])
if pred_type == "normals":
processed_out = torch.stack(processed_out, dim=0)
processed_out = processed_out.permute(0, 2, 3, 1).cpu().float()
else:
processed_out = torch.cat(processed_out, dim=0)
processed_out = processed_out.permute(0, 2, 3, 1).repeat(1, 1, 1, 3).cpu().float()
processed_out = 1.0 - processed_out
return (processed_out,)