update

config/dataset_depth/data_diode_all.yaml

@@ -0,0 +1,4 @@
+name: diode_depth
+disp_name: diode_depth_val_all
+dir: diode/diode_val.tar
+filenames: data_split/diode_depth/diode_val_all_filename_list.txt

config/dataset_depth/data_eth3d.yaml

@@ -0,0 +1,4 @@
+name: eth3d_depth
+disp_name: eth3d_depth_full
+dir: eth3d/eth3d.tar
+filenames: data_split/eth3d_depth/eth3d_filename_list.txt

config/dataset_depth/data_hypersim_train.yaml

@@ -0,0 +1,4 @@
+name: hypersim_depth
+disp_name: hypersim_depth_train
+dir: hypersim/hypersim_processed_train.tar
+filenames: data_split/hypersim_depth/filename_list_train_filtered.txt

config/dataset_depth/data_hypersim_val.yaml

@@ -0,0 +1,4 @@
+name: hypersim_depth
+disp_name: hypersim_depth_val
+dir: hypersim/hypersim_processed_val.tar
+filenames: data_split/hypersim_depth/filename_list_val_filtered.txt

config/dataset_depth/data_kitti_eigen_test.yaml

@@ -0,0 +1,6 @@
+name: kitti_depth
+disp_name: kitti_depth_eigen_test_full
+dir: kitti/kitti_eigen_split_test.tar
+filenames: data_split/kitti_depth/eigen_test_files_with_gt.txt
+kitti_bm_crop: true
+valid_mask_crop: eigen

config/dataset_depth/data_kitti_val.yaml

@@ -0,0 +1,6 @@
+name: kitti_depth
+disp_name: kitti_depth_val800_from_eigen_train
+dir: kitti/kitti_sampled_val_800.tar
+filenames: data_split/kitti_depth/eigen_val_from_train_800.txt
+kitti_bm_crop: true
+valid_mask_crop: eigen

config/dataset_depth/data_nyu_test.yaml

@@ -0,0 +1,5 @@
+name: nyu_depth
+disp_name: nyu_depth_test_full
+dir: nyuv2/nyu_labeled_extracted.tar
+filenames: data_split/nyu_depth/labeled/filename_list_test.txt
+eigen_valid_mask: true

config/dataset_depth/data_nyu_train.yaml

@@ -0,0 +1,5 @@
+name: nyu_depth
+disp_name: nyu_depth_train_full
+dir: nyuv2/nyu_labeled_extracted.tar
+filenames: data_split/nyu_depth/labeled/filename_list_train.txt
+eigen_valid_mask: true

config/dataset_depth/data_scannet_val.yaml

@@ -0,0 +1,4 @@
+name: scannet_depth
+disp_name: scannet_depth_val_800_1
+dir: scannet/scannet_val_sampled_800_1.tar
+filenames: data_split/scannet_depth/scannet_val_sampled_list_800_1.txt

config/dataset_depth/data_vkitti_train.yaml

@@ -0,0 +1,6 @@
+name: vkitti_depth
+disp_name: vkitti_depth_train
+dir: vkitti/vkitti.tar
+filenames: data_split/vkitti_depth/vkitti_train.txt
+kitti_bm_crop: true
+valid_mask_crop: null  # no valid_mask_crop for training

config/dataset_depth/data_vkitti_val.yaml

@@ -0,0 +1,6 @@
+name: vkitti_depth
+disp_name: vkitti_depth_val
+dir: vkitti/vkitti.tar
+filenames: data_split/vkitti_depth/vkitti_val.txt
+kitti_bm_crop: true
+valid_mask_crop: eigen

config/dataset_depth/dataset_train.yaml

@@ -0,0 +1,18 @@
+dataset:
+  train:
+    name: mixed
+    prob_ls: [0.9, 0.1]
+    dataset_list:
+    - name: hypersim_depth
+      disp_name: hypersim_depth_train
+      dir: hypersim/hypersim_processed_train.tar
+      filenames: data_split/hypersim_depth/filename_list_train_filtered.txt
+      resize_to_hw:
+      - 480
+      - 640
+    - name: vkitti_depth
+      disp_name: vkitti_depth_train
+      dir: vkitti/vkitti.tar
+      filenames: data_split/vkitti_depth/vkitti_train.txt
+      kitti_bm_crop: true
+      valid_mask_crop: null

config/dataset_depth/dataset_val.yaml

@@ -0,0 +1,45 @@
+dataset:
+  val:
+  # - name: hypersim_depth
+  #   disp_name: hypersim_depth_val
+  #   dir: hypersim/hypersim_processed_val.tar
+  #   filenames: data_split/hypersim_depth/filename_list_val_filtered.txt
+  #   resize_to_hw:
+  #   - 480
+  #   - 640
+
+  # - name: nyu_depth
+  #   disp_name: nyu_depth_train_full
+  #   dir: nyuv2/nyu_labeled_extracted.tar
+  #   filenames: data_split/nyu_depth/labeled/filename_list_train.txt
+  #   eigen_valid_mask: true
+
+  # - name: kitti_depth
+  #   disp_name: kitti_depth_val800_from_eigen_train
+  #   dir: kitti/kitti_depth_sampled_val_800.tar
+  #   filenames: data_split/kitti_depth/eigen_val_from_train_800.txt
+  #   kitti_bm_crop: true
+  #   valid_mask_crop: eigen
+
+  # Smaller subsets for faster validation during training
+  # The first dataset is used to calculate main eval metric.
+  - name: hypersim_depth
+    disp_name: hypersim_depth_val_small_80
+    dir: hypersim/hypersim_processed_val.tar
+    filenames: data_split/hypersim_depth/filename_list_val_filtered_small_80.txt
+    resize_to_hw:
+    - 480
+    - 640
+
+  - name: nyu_depth
+    disp_name: nyu_depth_train_small_100
+    dir: nyuv2/nyu_labeled_extracted.tar
+    filenames: data_split/nyu_depth/labeled/filename_list_train_small_100.txt
+    eigen_valid_mask: true
+
+  - name: kitti_depth
+    disp_name: kitti_depth_val_from_train_sub_100
+    dir: kitti/kitti_sampled_val_800.tar
+    filenames: data_split/kitti_depth/eigen_val_from_train_sub_100.txt
+    kitti_bm_crop: true
+    valid_mask_crop: eigen

config/dataset_depth/dataset_vis.yaml

@@ -0,0 +1,9 @@
+dataset:
+  vis:
+  - name: hypersim_depth
+    disp_name: hypersim_depth_vis
+    dir: hypersim/hypersim_processed_val.tar
+    filenames: data_split/hypersim_depth/selected_vis_sample.txt
+    resize_to_hw:
+    - 480
+    - 640

config/dataset_iid/data_appearance_interiorverse_test.yaml

@@ -0,0 +1,4 @@
+name: interiorverse_iid
+disp_name: interiorverse_iid_appearance_test
+dir: interiorverse/InteriorVerse.tar
+filenames: data_split/interiorverse_iid/interiorverse_test_scenes_85.txt

config/dataset_iid/data_appearance_synthetic_test.yaml

@@ -0,0 +1,4 @@
+name: interiorverse_iid
+disp_name: interiorverse_iid_appearance_test
+dir: synthetic
+filenames: data_split/osu/osu_test_scenes_85.txt

config/dataset_iid/data_art_test.yaml

@@ -0,0 +1,4 @@
+name: interiorverse_iid
+disp_name: interiorverse_iid_appearance_test
+dir: art
+filenames: data_split/osu/art_test_scenes.txt

config/dataset_iid/data_lighting_hypersim_test.yaml

@@ -0,0 +1,4 @@
+name: hypersim_iid
+disp_name: hypersim_iid_lighting_test
+dir: hypersim
+filenames: data_split/hypersim_iid/hypersim_test.txt

config/dataset_iid/dataset_appearance_train.yaml

@@ -0,0 +1,9 @@
+dataset:
+  train:
+    name: mixed
+    prob_ls: [1.0]
+    dataset_list:
+    - name: interiorverse_iid
+      disp_name: interiorverse_iid_appearance_train
+      dir: osu_albedo_new
+      filenames: data_split/osu/osu_train_scenes_85.txt

config/dataset_iid/dataset_appearance_val.yaml

@@ -0,0 +1,6 @@
+dataset:
+  val:
+  - name: interiorverse_iid
+    disp_name: interiorverse_iid_appearance_val
+    dir: synthetic
+    filenames: data_split/MatrixCity/matrixcity_val_scenes_small.txt

config/dataset_iid/dataset_appearance_vis.yaml

@@ -0,0 +1,6 @@
+dataset:
+  vis:
+  - name: interiorverse_iid
+    disp_name: interiorverse_iid_appearance_vis
+    dir: synthetic
+    filenames: data_split/MatrixCity/matrixcity_vis_scenes.txt

config/dataset_iid/dataset_lighting_train.yaml

@@ -0,0 +1,12 @@
+dataset:
+  train:
+    name: mixed
+    prob_ls: [1.0]
+    dataset_list:
+    - name: hypersim_iid
+      disp_name: hypersim_iid_lighting_train
+      dir: hypersim
+      filenames: data_split/hypersim_iid/hypersim_train_filtered.txt
+      resize_to_hw:
+      - 480
+      - 640

config/dataset_iid/dataset_lighting_val.yaml

@@ -0,0 +1,6 @@
+dataset:
+  val:
+  - name: hypersim_iid
+    disp_name: hypersim_iid_lighting_val
+    dir: hypersim
+    filenames: data_split/hypersim_iid/hypersim_val.txt

config/dataset_iid/dataset_lighting_vis.yaml

@@ -0,0 +1,6 @@
+dataset:
+  vis:
+  - name: hypersim_iid
+    disp_name: hypersim_iid_lighting_vis
+    dir: hypersim
+    filenames: data_split/hypersim_iid/hypersim_vis.txt

config/dataset_iid/osu_data_appearance_interiorverse_test.yaml

@@ -0,0 +1,4 @@
+name: interiorverse_iid
+disp_name: interiorverse_iid_appearance_test
+dir: synthetic
+filenames: data_split/osu/osu_test_scenes_85.txt

config/dataset_normals/data_diode_test.yaml

@@ -0,0 +1,4 @@
+name: diode_normals
+disp_name: diode_normals_test
+dir: diode/val
+filenames: data_split/diode_normals/diode_test.txt

config/dataset_normals/data_ibims_test.yaml

@@ -0,0 +1,4 @@
+name: ibims_normals
+disp_name: ibims_normals_test
+dir: ibims/ibims
+filenames: data_split/ibims_normals/ibims_test.txt

config/dataset_normals/data_nyu_test.yaml

@@ -0,0 +1,4 @@
+name: nyu_normals
+disp_name: nyu_normals_test
+dir: nyuv2/test
+filenames: data_split/nyu_normals/nyuv2_test.txt

config/dataset_normals/data_oasis_test.yaml

@@ -0,0 +1,4 @@
+name: oasis_normals
+disp_name: oasis_normals_test
+dir: oasis/val
+filenames: data_split/oasis_normals/oasis_test.txt

config/dataset_normals/data_scannet_test.yaml

@@ -0,0 +1,4 @@
+name: scannet_normals
+disp_name: scannet_normals_test
+dir: scannet
+filenames: data_split/scannet_normals/scannet_test.txt

config/dataset_normals/dataset_train.yaml

@@ -0,0 +1,25 @@
+dataset:
+  train:
+    name: mixed
+    prob_ls: [0.5, 0.49, 0.01]
+    dataset_list:
+    - name: hypersim_normals
+      disp_name: hypersim_normals_train
+      dir: hypersim
+      filenames: data_split/hypersim_normals/hypersim_filtered_all.txt
+      resize_to_hw:
+      - 480
+      - 640
+    - name: interiorverse_normals
+      disp_name: interiorverse_normals_train
+      dir: interiorverse/scenes_85
+      filenames: data_split/interiorverse_normals/interiorverse_filtered_all.txt
+      resize_to_hw: null
+    - name: sintel_normals
+      disp_name: sintel_normals_train
+      dir: sintel
+      filenames: data_split/sintel_normals/sintel_filtered.txt
+      resize_to_hw:
+      - 480
+      - 640
+      center_crop: true

config/dataset_normals/dataset_val.yaml

@@ -0,0 +1,7 @@
+dataset:
+  val:
+  - name: hypersim_normals
+    disp_name: hypersim_normals_val_small_100
+    dir: hypersim
+    filenames: data_split/hypersim_normals/hypersim_filtered_val_100.txt
+    resize_to_hw: null

config/dataset_normals/dataset_vis.yaml

@@ -0,0 +1,7 @@
+dataset:
+  vis:
+  - name: hypersim_normals
+    disp_name: hypersim_normals_vis
+    dir: hypersim
+    filenames: data_split/hypersim_normals/hypersim_filtered_vis_20.txt
+    resize_to_hw: null

config/logging.yaml

@@ -0,0 +1,5 @@
+logging:
+  filename: logging.log
+  format: ' %(asctime)s - %(levelname)s -%(filename)s - %(funcName)s >> %(message)s'
+  console_level: 20
+  file_level: 10

config/model_sdv2.yaml

@@ -0,0 +1,4 @@
+model:
+  name: marigold_pipeline
+  pretrained_path: stable-diffusion-2
+  latent_scale_factor: 0.18215

config/train_debug_depth.yaml

@@ -0,0 +1,10 @@
+base_config:
+  - config/train_marigold_depth.yaml
+
+trainer:
+  save_period: 5
+  backup_period: 10
+  validation_period: 5
+  visualization_period: 5
+
+max_iter: 50

config/train_debug_iid.yaml

@@ -0,0 +1,11 @@
+base_config:
+  # - config/train_marigold_iid_lighting.yaml
+  - config/train_marigold_iid_appearance.yaml
+
+trainer:
+  save_period: 10
+  backup_period: 10
+  validation_period: 5
+  visualization_period: 5
+
+max_iter: 50

config/train_debug_normals.yaml

@@ -0,0 +1,10 @@
+base_config:
+  - config/train_marigold_normals.yaml
+
+trainer:
+  save_period: 5
+  backup_period: 10
+  validation_period: 5
+  visualization_period: 5
+
+max_iter: 50

config/train_marigold_depth.yaml

@@ -0,0 +1,94 @@
+base_config:
+- config/logging.yaml
+- config/wandb.yaml
+- config/dataset_depth/dataset_train.yaml
+- config/dataset_depth/dataset_val.yaml
+- config/dataset_depth/dataset_vis.yaml
+- config/model_sdv2.yaml
+
+pipeline:
+  name: MarigoldDepthPipeline
+  kwargs:
+    scale_invariant: true
+    shift_invariant: true
+    default_denoising_steps: 4
+    default_processing_resolution: 768
+
+depth_normalization:
+  type: scale_shift_depth
+  clip: true
+  norm_min: -1.0
+  norm_max: 1.0
+  min_max_quantile: 0.02
+
+augmentation:
+  lr_flip_p: 0.5
+
+dataloader:
+  num_workers: 2
+  effective_batch_size: 32
+  max_train_batch_size: 2
+  seed: 2024  # to ensure continuity when resuming from checkpoint
+
+trainer:
+  name: MarigoldDepthTrainer
+  training_noise_scheduler:
+    pretrained_path: stable-diffusion-2
+  init_seed: 2024  # use null to train w/o seeding
+  save_period: 50
+  backup_period: 2000
+  validation_period: 500
+  visualization_period: 1000
+
+multi_res_noise:
+  strength: 0.9
+  annealed: true
+  downscale_strategy: original
+
+gt_depth_type: depth_raw_norm
+gt_mask_type: valid_mask_raw
+
+max_epoch: 10000  # a large enough number
+max_iter: 30000  # usually converges at around 20k
+
+optimizer:
+  name: Adam
+
+loss:
+  name: mse_loss
+  kwargs:
+    reduction: mean
+
+lr: 3.0e-05
+lr_scheduler:
+  name: IterExponential
+  kwargs:
+    total_iter: 25000
+    final_ratio: 0.01
+    warmup_steps: 100
+
+# Light setting for the in-training validation and visualization
+validation:
+  denoising_steps: 1
+  ensemble_size: 1
+  processing_res: 0
+  match_input_res: false
+  resample_method: bilinear
+  main_val_metric: abs_relative_difference
+  main_val_metric_goal: minimize
+  init_seed: 2024
+
+eval:
+  alignment: least_square
+  align_max_res: null
+  eval_metrics:
+  - abs_relative_difference
+  - squared_relative_difference
+  - rmse_linear
+  - rmse_log
+  - log10
+  - delta1_acc
+  - delta2_acc
+  - delta3_acc
+  - i_rmse
+  - silog_rmse

config/train_marigold_iid_appearance.yaml

@@ -0,0 +1,81 @@
+base_config:
+- config/logging.yaml
+- config/wandb.yaml
+- config/dataset_iid/dataset_appearance_train.yaml
+- config/dataset_iid/dataset_appearance_val.yaml
+- config/dataset_iid/dataset_appearance_vis.yaml
+- config/model_sdv2.yaml
+
+pipeline:
+  name: MarigoldIIDPipeline
+  kwargs:
+    default_denoising_steps: 4
+    default_processing_resolution: 768
+    target_properties:
+      target_names:
+          - albedo
+      albedo: 
+        prediction_space: srgb
+
+augmentation:
+  lr_flip_p: 0.5
+
+dataloader:
+  num_workers: 2
+  effective_batch_size: 32
+  max_train_batch_size: 8
+  seed: 2024  # to ensure continuity when resuming from checkpoint
+
+trainer:
+  name: MarigoldIIDTrainer
+  training_noise_scheduler:
+    pretrained_path: stable-diffusion-2
+  init_seed: 2024  # use null to train w/o seeding
+  save_period: 50
+  backup_period: 2000
+  validation_period: 500
+  visualization_period: 1000
+
+multi_res_noise:
+  strength: 0.9
+  annealed: true
+  downscale_strategy: original
+
+gt_mask_type: mask
+
+max_epoch: 10000  # a large enough number
+max_iter: 10000  # usually converges at around 40k
+
+optimizer:
+  name: Adam
+
+loss:
+  name: mse_loss
+  kwargs:
+    reduction: mean
+
+lr: 2.0e-05
+lr_scheduler:
+  name: IterExponential
+  kwargs:
+    total_iter: 5000
+    final_ratio: 0.01
+    warmup_steps: 100
+
+# Light setting for the in-training validation and visualization
+validation:
+  denoising_steps: 4
+  ensemble_size: 1
+  processing_res: 0
+  match_input_res: true
+  resample_method: bilinear
+  main_val_metric: psnr
+  main_val_metric_goal: maximize
+  init_seed: 2024
+  use_mask: false
+
+eval:
+  eval_metrics:
+  - psnr
+  targets_to_eval_in_linear_space:
+  - material

config/train_marigold_iid_appearance_finetuned.yaml

@@ -0,0 +1,81 @@
+base_config:
+- config/logging.yaml
+- config/wandb.yaml
+- config/dataset_iid/dataset_appearance_train.yaml
+- config/dataset_iid/dataset_appearance_val.yaml
+- config/dataset_iid/dataset_appearance_vis.yaml
+- config/model_sdv2.yaml
+
+pipeline:
+  name: MarigoldIIDPipeline
+  kwargs:
+    default_denoising_steps: 4
+    default_processing_resolution: 768
+    target_properties:
+      target_names:
+          - albedo
+      albedo: 
+        prediction_space: srgb
+
+augmentation:
+  lr_flip_p: 0.5
+
+dataloader:
+  num_workers: 2
+  effective_batch_size: 32
+  max_train_batch_size: 8
+  seed: 2024  # to ensure continuity when resuming from checkpoint
+
+trainer:
+  name: MarigoldIIDTrainer
+  training_noise_scheduler:
+    pretrained_path: stable-diffusion-2
+  init_seed: 2024  # use null to train w/o seeding
+  save_period: 50
+  backup_period: 2000
+  validation_period: 177
+  visualization_period: 177
+
+multi_res_noise:
+  strength: 0.9
+  annealed: true
+  downscale_strategy: original
+
+gt_mask_type: null
+
+max_epoch: 10000  # a large enough number
+max_iter: 5000  # usually converges at around 40k
+
+optimizer:
+  name: Adam
+
+loss:
+  name: mse_loss
+  kwargs:
+    reduction: mean
+
+lr: 5.0e-07
+lr_scheduler:
+  name: IterExponential
+  kwargs:
+    total_iter: 2500
+    final_ratio: 0.01
+    warmup_steps: 100
+
+# Light setting for the in-training validation and visualization
+validation:
+  denoising_steps: 4
+  ensemble_size: 1
+  processing_res: 1000
+  match_input_res: true
+  resample_method: bilinear
+  main_val_metric: psnr
+  main_val_metric_goal: maximize
+  init_seed: 2024
+  use_mask: false
+
+eval:
+  eval_metrics:
+  - psnr
+  targets_to_eval_in_linear_space:
+  - material

config/train_marigold_iid_lighting.yaml

@@ -0,0 +1,82 @@
+base_config:
+- config/logging.yaml
+- config/wandb.yaml
+- config/dataset_iid/dataset_lighting_train.yaml
+- config/dataset_iid/dataset_lighting_val.yaml
+- config/dataset_iid/dataset_lighting_vis.yaml
+- config/model_sdv2.yaml
+
+pipeline:
+  name: MarigoldIIDPipeline
+  kwargs:
+    default_denoising_steps: 4
+    default_processing_resolution: 768
+    target_properties:
+      target_names:
+      - albedo
+      albedo:
+        prediction_space: linear
+        up_to_scale: false
+
+augmentation:
+  lr_flip_p: 0.5
+
+dataloader:
+  num_workers: 2
+  effective_batch_size: 32
+  max_train_batch_size: 8
+  seed: 2024  # to ensure continuity when resuming from checkpoint
+
+trainer:
+  name: MarigoldIIDTrainer
+  training_noise_scheduler:
+    pretrained_path: stable-diffusion-2
+  init_seed: 2024  # use null to train w/o seeding
+  save_period: 50
+  backup_period: 2000
+  validation_period: 500
+  visualization_period: 1000
+
+multi_res_noise:
+  strength: 0.9
+  annealed: true
+  downscale_strategy: original
+
+gt_mask_type: mask
+
+max_epoch: 10000  # a large enough number
+max_iter: 50000  # usually converges at around 34k
+
+optimizer:
+  name: Adam
+
+loss:
+  name: mse_loss
+  kwargs:
+    reduction: mean
+
+lr: 8e-05
+lr_scheduler:
+  name: IterExponential
+  kwargs:
+    total_iter: 45000
+    final_ratio: 0.01
+    warmup_steps: 100
+
+# Light setting for the in-training validation and visualization
+validation:
+  denoising_steps: 4
+  ensemble_size: 1
+  processing_res: 0
+  match_input_res: true
+  resample_method: bilinear
+  main_val_metric: psnr
+  main_val_metric_goal: maximize
+  init_seed: 2024
+  use_mask: false
+
+eval:
+  eval_metrics:
+  - psnr
+  targets_to_eval_in_linear_space:
+  - None

config/train_marigold_normals.yaml

@@ -0,0 +1,86 @@
+base_config:
+- config/logging.yaml
+- config/wandb.yaml
+- config/dataset_normals/dataset_train.yaml
+- config/dataset_normals/dataset_val.yaml
+- config/dataset_normals/dataset_vis.yaml
+- config/model_sdv2.yaml
+
+pipeline:
+  name: MarigoldNormalsPipeline
+  kwargs:
+    default_denoising_steps: 4
+    default_processing_resolution: 768
+
+augmentation:
+  lr_flip_p: 0.5
+  color_jitter_p: 0.3
+  gaussian_blur_p: 0.3
+  motion_blur_p: 0.3
+  gaussian_blur_sigma: 4
+  motion_blur_kernel_size: 11
+  motion_blur_angle_range: 360
+  jitter_brightness_factor: 0.5
+  jitter_contrast_factor: 0.5
+  jitter_saturation_factor: 0.5
+  jitter_hue_factor: 0.2
+
+dataloader:
+  num_workers: 2
+  effective_batch_size: 32
+  max_train_batch_size: 2
+  seed: 2024  # to ensure continuity when resuming from checkpoint
+
+trainer:
+  name: MarigoldNormalsTrainer
+  training_noise_scheduler:
+    pretrained_path: stable-diffusion-2
+  init_seed: 2024  # use null to train w/o seeding
+  save_period: 50
+  backup_period: 2000
+  validation_period: 500
+  visualization_period: 1000
+
+multi_res_noise:
+  strength: 0.9
+  annealed: true
+  downscale_strategy: original
+
+gt_normals_type: normals
+gt_mask_type: null
+
+max_epoch: 10000  # a large enough number
+max_iter: 30000  # usually converges at around 26k
+
+optimizer:
+  name: Adam
+
+loss:
+  name: mse_loss
+  kwargs:
+    reduction: mean
+
+lr: 6.0e-05
+lr_scheduler:
+  name: IterExponential
+  kwargs:
+    total_iter: 25000
+    final_ratio: 0.01
+    warmup_steps: 100
+
+# Light setting for the in-training validation and visualization
+validation:
+  denoising_steps: 4
+  ensemble_size: 1
+  processing_res: 768
+  match_input_res: true
+  resample_method: bilinear
+  main_val_metric: mean_angular_error
+  main_val_metric_goal: minimize
+  init_seed: 0
+
+eval:
+  align_max_res: null
+  eval_metrics:
+  - mean_angular_error
+  - sub11_25_error

config/wandb.yaml

@@ -0,0 +1,3 @@
+wandb:
+  # entity: your_entity
+  project: marigold

marigold/marigold_depth_pipeline.py

@@ -0,0 +1,516 @@
+# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# --------------------------------------------------------------------------
+# More information about Marigold:
+#   https://marigoldmonodepth.github.io
+#   https://marigoldcomputervision.github.io
+# Efficient inference pipelines are now part of diffusers:
+#   https://huggingface.co/docs/diffusers/using-diffusers/marigold_usage
+#   https://huggingface.co/docs/diffusers/api/pipelines/marigold
+# Examples of trained models and live demos:
+#   https://huggingface.co/prs-eth
+# Related projects:
+#   https://rollingdepth.github.io/
+#   https://marigolddepthcompletion.github.io/
+# Citation (BibTeX):
+#   https://github.com/prs-eth/Marigold#-citation
+# If you find Marigold useful, we kindly ask you to cite our papers.
+# --------------------------------------------------------------------------
+
+import logging
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import (
+    AutoencoderKL,
+    DDIMScheduler,
+    DiffusionPipeline,
+    LCMScheduler,
+    UNet2DConditionModel,
+)
+from diffusers.utils import BaseOutput
+from torch.utils.data import DataLoader, TensorDataset
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms.functional import pil_to_tensor, resize
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+from typing import Dict, Optional, Union
+
+from .util.batchsize import find_batch_size
+from .util.ensemble import ensemble_depth
+from .util.image_util import (
+    chw2hwc,
+    colorize_depth_maps,
+    get_tv_resample_method,
+    resize_max_res,
+)
+
+
+class MarigoldDepthOutput(BaseOutput):
+    """
+    Output class for Marigold Monocular Depth Estimation pipeline.
+
+    Args:
+        depth_np (`np.ndarray`):
+            Predicted depth map, with depth values in the range of [0, 1].
+        depth_colored (`PIL.Image.Image`):
+            Colorized depth map, with the shape of [H, W, 3] and values in [0, 255].
+        uncertainty (`None` or `np.ndarray`):
+            Uncalibrated uncertainty(MAD, median absolute deviation) coming from ensembling.
+    """
+
+    depth_np: np.ndarray
+    depth_colored: Union[None, Image.Image]
+    uncertainty: Union[None, np.ndarray]
+
+
+class MarigoldDepthPipeline(DiffusionPipeline):
+    """
+    Pipeline for Marigold Monocular Depth Estimation: https://marigoldcomputervision.github.io.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Args:
+        unet (`UNet2DConditionModel`):
+            Conditional U-Net to denoise the prediction latent, conditioned on image latent.
+        vae (`AutoencoderKL`):
+            Variational Auto-Encoder (VAE) Model to encode and decode images and predictions
+            to and from latent representations.
+        scheduler (`DDIMScheduler`):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+        text_encoder (`CLIPTextModel`):
+            Text-encoder, for empty text embedding.
+        tokenizer (`CLIPTokenizer`):
+            CLIP tokenizer.
+        scale_invariant (`bool`, *optional*):
+            A model property specifying whether the predicted depth maps are scale-invariant. This value must be set in
+            the model config. When used together with the `shift_invariant=True` flag, the model is also called
+            "affine-invariant". NB: overriding this value is not supported.
+        shift_invariant (`bool`, *optional*):
+            A model property specifying whether the predicted depth maps are shift-invariant. This value must be set in
+            the model config. When used together with the `scale_invariant=True` flag, the model is also called
+            "affine-invariant". NB: overriding this value is not supported.
+        default_denoising_steps (`int`, *optional*):
+            The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable
+            quality with the given model. This value must be set in the model config. When the pipeline is called
+            without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure
+            reasonable results with various model flavors compatible with the pipeline, such as those relying on very
+            short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`).
+        default_processing_resolution (`int`, *optional*):
+            The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in
+            the model config. When the pipeline is called without explicitly setting `processing_resolution`, the
+            default value is used. This is required to ensure reasonable results with various model flavors trained
+            with varying optimal processing resolution values.
+    """
+
+    latent_scale_factor = 0.18215
+
+    def __init__(
+        self,
+        unet: UNet2DConditionModel,
+        vae: AutoencoderKL,
+        scheduler: Union[DDIMScheduler, LCMScheduler],
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        scale_invariant: Optional[bool] = True,
+        shift_invariant: Optional[bool] = True,
+        default_denoising_steps: Optional[int] = None,
+        default_processing_resolution: Optional[int] = None,
+    ):
+        super().__init__()
+        self.register_modules(
+            unet=unet,
+            vae=vae,
+            scheduler=scheduler,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+        )
+        self.register_to_config(
+            scale_invariant=scale_invariant,
+            shift_invariant=shift_invariant,
+            default_denoising_steps=default_denoising_steps,
+            default_processing_resolution=default_processing_resolution,
+        )
+
+        self.scale_invariant = scale_invariant
+        self.shift_invariant = shift_invariant
+        self.default_denoising_steps = default_denoising_steps
+        self.default_processing_resolution = default_processing_resolution
+
+        self.empty_text_embed = None
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        input_image: Union[Image.Image, torch.Tensor],
+        denoising_steps: Optional[int] = None,
+        ensemble_size: int = 1,
+        processing_res: Optional[int] = None,
+        match_input_res: bool = True,
+        resample_method: str = "bilinear",
+        batch_size: int = 0,
+        generator: Union[torch.Generator, None] = None,
+        color_map: str = "Spectral",
+        show_progress_bar: bool = True,
+        ensemble_kwargs: Dict = None,
+    ) -> MarigoldDepthOutput:
+        """
+        Function invoked when calling the pipeline.
+
+        Args:
+            input_image (`Image`):
+                Input RGB (or gray-scale) image.
+            denoising_steps (`int`, *optional*, defaults to `None`):
+                Number of denoising diffusion steps during inference. The default value `None` results in automatic
+                selection.
+            ensemble_size (`int`, *optional*, defaults to `1`):
+                Number of predictions to be ensembled.
+            processing_res (`int`, *optional*, defaults to `None`):
+                Effective processing resolution. When set to `0`, processes at the original image resolution. This
+                produces crisper predictions, but may also lead to the overall loss of global context. The default
+                value `None` resolves to the optimal value from the model config.
+            match_input_res (`bool`, *optional*, defaults to `True`):
+                Resize the prediction to match the input resolution.
+                Only valid if `processing_res` > 0.
+            resample_method: (`str`, *optional*, defaults to `bilinear`):
+                Resampling method used to resize images and predictions. This can be one of `bilinear`, `bicubic` or
+                `nearest`, defaults to: `bilinear`.
+            batch_size (`int`, *optional*, defaults to `0`):
+                Inference batch size, no bigger than `num_ensemble`.
+                If set to 0, the script will automatically decide the proper batch size.
+            generator (`torch.Generator`, *optional*, defaults to `None`)
+                Random generator for initial noise generation.
+            show_progress_bar (`bool`, *optional*, defaults to `True`):
+                Display a progress bar of diffusion denoising.
+            color_map (`str`, *optional*, defaults to `"Spectral"`, pass `None` to skip colorized depth map generation):
+                Colormap used to colorize the depth map.
+            scale_invariant (`str`, *optional*, defaults to `True`):
+                Flag of scale-invariant prediction, if True, scale will be adjusted from the raw prediction.
+            shift_invariant (`str`, *optional*, defaults to `True`):
+                Flag of shift-invariant prediction, if True, shift will be adjusted from the raw prediction, if False,
+                near plane will be fixed at 0m.
+            ensemble_kwargs (`dict`, *optional*, defaults to `None`):
+                Arguments for detailed ensembling settings.
+        Returns:
+            `MarigoldDepthOutput`: Output class for Marigold monocular depth prediction pipeline, including:
+            - **depth_np** (`np.ndarray`) Predicted depth map with depth values in the range of [0, 1]
+            - **depth_colored** (`PIL.Image.Image`) Colorized depth map, with the shape of [H, W, 3] and values in [0, 255], None if `color_map` is `None`
+            - **uncertainty** (`None` or `np.ndarray`) Uncalibrated uncertainty(MAD, median absolute deviation)
+                    coming from ensembling. None if `ensemble_size = 1`
+        """
+        # Model-specific optimal default values leading to fast and reasonable results.
+        if denoising_steps is None:
+            denoising_steps = self.default_denoising_steps
+        if processing_res is None:
+            processing_res = self.default_processing_resolution
+
+        assert processing_res >= 0
+        assert ensemble_size >= 1
+
+        # Check if denoising step is reasonable
+        self._check_inference_step(denoising_steps)
+
+        resample_method: InterpolationMode = get_tv_resample_method(resample_method)
+
+        # ----------------- Image Preprocess -----------------
+        # Convert to torch tensor
+        if isinstance(input_image, Image.Image):
+            input_image = input_image.convert("RGB")
+            # convert to torch tensor [H, W, rgb] -> [rgb, H, W]
+            rgb = pil_to_tensor(input_image)
+            rgb = rgb.unsqueeze(0)  # [1, rgb, H, W]
+        elif isinstance(input_image, torch.Tensor):
+            rgb = input_image
+        else:
+            raise TypeError(f"Unknown input type: {type(input_image) = }")
+        input_size = rgb.shape
+        assert (
+            4 == rgb.dim() and 3 == input_size[-3]
+        ), f"Wrong input shape {input_size}, expected [1, rgb, H, W]"
+
+        # Resize image
+        if processing_res > 0:
+            rgb = resize_max_res(
+                rgb,
+                max_edge_resolution=processing_res,
+                resample_method=resample_method,
+            )
+
+        # Normalize rgb values
+        rgb_norm: torch.Tensor = rgb / 255.0 * 2.0 - 1.0  #  [0, 255] -> [-1, 1]
+        rgb_norm = rgb_norm.to(self.dtype)
+        assert rgb_norm.min() >= -1.0 and rgb_norm.max() <= 1.0
+
+        # ----------------- Predicting depth -----------------
+        # Batch repeated input image
+        duplicated_rgb = rgb_norm.expand(ensemble_size, -1, -1, -1)
+        single_rgb_dataset = TensorDataset(duplicated_rgb)
+        if batch_size > 0:
+            _bs = batch_size
+        else:
+            _bs = find_batch_size(
+                ensemble_size=ensemble_size,
+                input_res=max(rgb_norm.shape[1:]),
+                dtype=self.dtype,
+            )
+
+        single_rgb_loader = DataLoader(
+            single_rgb_dataset, batch_size=_bs, shuffle=False
+        )
+
+        # Predict depth maps (batched)
+        target_pred_ls = []
+        if show_progress_bar:
+            iterable = tqdm(
+                single_rgb_loader, desc=" " * 2 + "Inference batches", leave=False
+            )
+        else:
+            iterable = single_rgb_loader
+        for batch in iterable:
+            (batched_img,) = batch
+            target_pred_raw = self.single_infer(
+                rgb_in=batched_img,
+                num_inference_steps=denoising_steps,
+                show_pbar=show_progress_bar,
+                generator=generator,
+            )
+            target_pred_ls.append(target_pred_raw.detach())
+        target_preds = torch.concat(target_pred_ls, dim=0)
+        torch.cuda.empty_cache()  # clear vram cache for ensembling
+
+        # ----------------- Test-time ensembling -----------------
+        if ensemble_size > 1:
+            final_pred, pred_uncert = ensemble_depth(
+                target_preds,
+                scale_invariant=self.scale_invariant,
+                shift_invariant=self.shift_invariant,
+                **(ensemble_kwargs or {}),
+            )
+        else:
+            final_pred = target_preds
+            pred_uncert = None
+
+        # Resize back to original resolution
+        if match_input_res:
+            final_pred = resize(
+                final_pred,
+                input_size[-2:],
+                interpolation=resample_method,
+                antialias=True,
+            )
+
+        # Convert to numpy
+        final_pred = final_pred.squeeze()
+        final_pred = final_pred.cpu().numpy()
+        if pred_uncert is not None:
+            pred_uncert = pred_uncert.squeeze().cpu().numpy()
+
+        # Clip output range
+        final_pred = final_pred.clip(0, 1)
+
+        # Colorize
+        if color_map is not None:
+            depth_colored = colorize_depth_maps(
+                final_pred, 0, 1, cmap=color_map
+            ).squeeze()  # [3, H, W], value in (0, 1)
+            depth_colored = (depth_colored * 255).astype(np.uint8)
+            depth_colored_hwc = chw2hwc(depth_colored)
+            depth_colored_img = Image.fromarray(depth_colored_hwc)
+        else:
+            depth_colored_img = None
+
+        return MarigoldDepthOutput(
+            depth_np=final_pred,
+            depth_colored=depth_colored_img,
+            uncertainty=pred_uncert,
+        )
+
+    def _check_inference_step(self, n_step: int) -> None:
+        """
+        Check if denoising step is reasonable
+        Args:
+            n_step (`int`): denoising steps
+        """
+        assert n_step >= 1
+
+        if isinstance(self.scheduler, DDIMScheduler):
+            if "trailing" != self.scheduler.config.timestep_spacing:
+                logging.warning(
+                    f"The loaded `DDIMScheduler` is configured with `timestep_spacing="
+                    f'"{self.scheduler.config.timestep_spacing}"`; the recommended setting is `"trailing"`. '
+                    f"This change is backward-compatible and yields better results. "
+                    f"Consider using `prs-eth/marigold-depth-v1-1` for the best experience."
+                )
+            else:
+                if n_step > 10:
+                    logging.warning(
+                        f"Setting too many denoising steps ({n_step}) may degrade the prediction; consider relying on "
+                        f"the default values."
+                    )
+            if not self.scheduler.config.rescale_betas_zero_snr:
+                logging.warning(
+                    f"The loaded `DDIMScheduler` is configured with `rescale_betas_zero_snr="
+                    f"{self.scheduler.config.rescale_betas_zero_snr}`; the recommended setting is True. "
+                    f"Consider using `prs-eth/marigold-depth-v1-1` for the best experience."
+                )
+        elif isinstance(self.scheduler, LCMScheduler):
+            logging.warning(
+                "DeprecationWarning: LCMScheduler will not be supported in the future. "
+                "Consider using `prs-eth/marigold-depth-v1-1` for the best experience."
+            )
+            if n_step > 10:
+                logging.warning(
+                    f"Setting too many denoising steps ({n_step}) may degrade the prediction; consider relying on "
+                    f"the default values."
+                )
+        else:
+            raise RuntimeError(f"Unsupported scheduler type: {type(self.scheduler)}")
+
+    def encode_empty_text(self):
+        """
+        Encode text embedding for empty prompt
+        """
+        prompt = ""
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="do_not_pad",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids.to(self.text_encoder.device)
+        self.empty_text_embed = self.text_encoder(text_input_ids)[0].to(self.dtype)
+
+    @torch.no_grad()
+    def single_infer(
+        self,
+        rgb_in: torch.Tensor,
+        num_inference_steps: int,
+        generator: Union[torch.Generator, None],
+        show_pbar: bool,
+    ) -> torch.Tensor:
+        """
+        Perform a single prediction without ensembling.
+
+        Args:
+            rgb_in (`torch.Tensor`):
+                Input RGB image.
+            num_inference_steps (`int`):
+                Number of diffusion denoisign steps (DDIM) during inference.
+            show_pbar (`bool`):
+                Display a progress bar of diffusion denoising.
+            generator (`torch.Generator`)
+                Random generator for initial noise generation.
+        Returns:
+            `torch.Tensor`: Predicted targets.
+        """
+        device = self.device
+        rgb_in = rgb_in.to(device)
+
+        # Set timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps  # [T]
+
+        # Encode image
+        rgb_latent = self.encode_rgb(rgb_in)  # [B, 4, h, w]
+
+        # Noisy latent for outputs
+        target_latent = torch.randn(
+            rgb_latent.shape,
+            device=device,
+            dtype=self.dtype,
+            generator=generator,
+        )  # [B, 4, h, w]
+
+        # Batched empty text embedding
+        if self.empty_text_embed is None:
+            self.encode_empty_text()
+        batch_empty_text_embed = self.empty_text_embed.repeat(
+            (rgb_latent.shape[0], 1, 1)
+        ).to(device)  # [B, 2, 1024]
+
+        # Denoising loop
+        if show_pbar:
+            iterable = tqdm(
+                enumerate(timesteps),
+                total=len(timesteps),
+                leave=False,
+                desc=" " * 4 + "Diffusion denoising",
+            )
+        else:
+            iterable = enumerate(timesteps)
+
+        for i, t in iterable:
+            unet_input = torch.cat(
+                [rgb_latent, target_latent], dim=1
+            )  # this order is important
+
+            # predict the noise residual
+            noise_pred = self.unet(
+                unet_input, t, encoder_hidden_states=batch_empty_text_embed
+            ).sample  # [B, 4, h, w]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            target_latent = self.scheduler.step(
+                noise_pred, t, target_latent, generator=generator
+            ).prev_sample
+
+        depth = self.decode_depth(target_latent)  # [B,3,H,W]
+
+        # clip prediction
+        depth = torch.clip(depth, -1.0, 1.0)
+        # shift to [0, 1]
+        depth = (depth + 1.0) / 2.0
+
+        return depth
+
+    def encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor:
+        """
+        Encode RGB image into latent.
+
+        Args:
+            rgb_in (`torch.Tensor`):
+                Input RGB image to be encoded.
+
+        Returns:
+            `torch.Tensor`: Image latent.
+        """
+        # encode
+        h = self.vae.encoder(rgb_in)
+        moments = self.vae.quant_conv(h)
+        mean, logvar = torch.chunk(moments, 2, dim=1)
+        # scale latent
+        rgb_latent = mean * self.latent_scale_factor
+        return rgb_latent
+
+    def decode_depth(self, depth_latent: torch.Tensor) -> torch.Tensor:
+        """
+        Decode depth latent into depth map.
+
+        Args:
+            depth_latent (`torch.Tensor`):
+                Depth latent to be decoded.
+
+        Returns:
+            `torch.Tensor`: Decoded depth map.
+        """
+        # scale latent
+        depth_latent = depth_latent / self.latent_scale_factor
+        # decode
+        z = self.vae.post_quant_conv(depth_latent)
+        stacked = self.vae.decoder(z)
+        # mean of output channels
+        depth_mean = stacked.mean(dim=1, keepdim=True)
+        return depth_mean

marigold/marigold_normals_pipeline.py

@@ -0,0 +1,479 @@
+# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# --------------------------------------------------------------------------
+# More information about Marigold:
+#   https://marigoldmonodepth.github.io
+#   https://marigoldcomputervision.github.io
+# Efficient inference pipelines are now part of diffusers:
+#   https://huggingface.co/docs/diffusers/using-diffusers/marigold_usage
+#   https://huggingface.co/docs/diffusers/api/pipelines/marigold
+# Examples of trained models and live demos:
+#   https://huggingface.co/prs-eth
+# Related projects:
+#   https://rollingdepth.github.io/
+#   https://marigolddepthcompletion.github.io/
+# Citation (BibTeX):
+#   https://github.com/prs-eth/Marigold#-citation
+# If you find Marigold useful, we kindly ask you to cite our papers.
+# --------------------------------------------------------------------------
+
+import logging
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import (
+    AutoencoderKL,
+    DDIMScheduler,
+    DiffusionPipeline,
+    LCMScheduler,
+    UNet2DConditionModel,
+)
+from diffusers.utils import BaseOutput
+from torch.utils.data import DataLoader, TensorDataset
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms.functional import pil_to_tensor, resize
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+from typing import Dict, Optional, Union
+
+from .util.batchsize import find_batch_size
+from .util.ensemble import ensemble_normals
+from .util.image_util import (
+    chw2hwc,
+    get_tv_resample_method,
+    resize_max_res,
+)
+
+
+class MarigoldNormalsOutput(BaseOutput):
+    """
+    Output class for Marigold Surface Normals Estimation pipeline.
+
+    Args:
+        normals_np (`np.ndarray`):
+            Predicted normals map of shape [3, H, W] with values in the range of [-1, 1] (unit length vectors).
+        normals_img (`PIL.Image.Image`):
+            Normals image, with the shape of [H, W, 3] and values in [0, 255].
+        uncertainty (`None` or `np.ndarray`):
+            Uncalibrated uncertainty(MAD, median absolute deviation) coming from ensembling.
+    """
+
+    normals_np: np.ndarray
+    normals_img: Image.Image
+    uncertainty: Union[None, np.ndarray]
+
+
+class MarigoldNormalsPipeline(DiffusionPipeline):
+    """
+    Pipeline for Marigold Surface Normals Estimation: https://marigoldcomputervision.github.io.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Args:
+        unet (`UNet2DConditionModel`):
+            Conditional U-Net to denoise the prediction latent, conditioned on image latent.
+        vae (`AutoencoderKL`):
+            Variational Auto-Encoder (VAE) Model to encode and decode images and predictions
+            to and from latent representations.
+        scheduler (`DDIMScheduler`):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+        text_encoder (`CLIPTextModel`):
+            Text-encoder, for empty text embedding.
+        tokenizer (`CLIPTokenizer`):
+            CLIP tokenizer.
+        default_denoising_steps (`int`, *optional*):
+            The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable
+            quality with the given model. This value must be set in the model config. When the pipeline is called
+            without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure
+            reasonable results with various model flavors compatible with the pipeline, such as those relying on very
+            short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`).
+        default_processing_resolution (`int`, *optional*):
+            The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in
+            the model config. When the pipeline is called without explicitly setting `processing_resolution`, the
+            default value is used. This is required to ensure reasonable results with various model flavors trained
+            with varying optimal processing resolution values.
+    """
+
+    latent_scale_factor = 0.18215
+
+    def __init__(
+        self,
+        unet: UNet2DConditionModel,
+        vae: AutoencoderKL,
+        scheduler: Union[DDIMScheduler, LCMScheduler],
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        default_denoising_steps: Optional[int] = None,
+        default_processing_resolution: Optional[int] = None,
+    ):
+        super().__init__()
+        self.register_modules(
+            unet=unet,
+            vae=vae,
+            scheduler=scheduler,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+        )
+        self.register_to_config(
+            default_denoising_steps=default_denoising_steps,
+            default_processing_resolution=default_processing_resolution,
+        )
+
+        self.default_denoising_steps = default_denoising_steps
+        self.default_processing_resolution = default_processing_resolution
+
+        self.empty_text_embed = None
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        input_image: Union[Image.Image, torch.Tensor],
+        denoising_steps: Optional[int] = None,
+        ensemble_size: int = 1,
+        processing_res: Optional[int] = None,
+        match_input_res: bool = True,
+        resample_method: str = "bilinear",
+        batch_size: int = 0,
+        generator: Union[torch.Generator, None] = None,
+        show_progress_bar: bool = True,
+        ensemble_kwargs: Dict = None,
+    ) -> MarigoldNormalsOutput:
+        """
+        Function invoked when calling the pipeline.
+
+        Args:
+            input_image (`Image`):
+                Input RGB (or gray-scale) image.
+            denoising_steps (`int`, *optional*, defaults to `None`):
+                Number of denoising diffusion steps during inference. The default value `None` results in automatic
+                selection.
+            ensemble_size (`int`, *optional*, defaults to `1`):
+                Number of predictions to be ensembled.
+            processing_res (`int`, *optional*, defaults to `None`):
+                Effective processing resolution. When set to `0`, processes at the original image resolution. This
+                produces crisper predictions, but may also lead to the overall loss of global context. The default
+                value `None` resolves to the optimal value from the model config.
+            match_input_res (`bool`, *optional*, defaults to `True`):
+                Resize the prediction to match the input resolution.
+                Only valid if `processing_res` > 0.
+            resample_method: (`str`, *optional*, defaults to `bilinear`):
+                Resampling method used to resize images and predictions. This can be one of `bilinear`, `bicubic` or
+                `nearest`, defaults to: `bilinear`.
+            batch_size (`int`, *optional*, defaults to `0`):
+                Inference batch size, no bigger than `num_ensemble`.
+                If set to 0, the script will automatically decide the proper batch size.
+            generator (`torch.Generator`, *optional*, defaults to `None`)
+                Random generator for initial noise generation.
+            show_progress_bar (`bool`, *optional*, defaults to `True`):
+                Display a progress bar of diffusion denoising.
+            ensemble_kwargs (`dict`, *optional*, defaults to `None`):
+                Arguments for detailed ensembling settings.
+        Returns:
+            `MarigoldNormalsOutput`: Output class for Marigold monocular surface normals estimation pipeline, including:
+            - **normals_np** (`np.ndarray`) Predicted normals map of shape [3, H, W] with values in the range of [-1, 1]
+                    (unit length vectors)
+            - **normals_img** (`PIL.Image.Image`) Normals image, with the shape of [H, W, 3] and values in [0, 255]
+            - **uncertainty** (`None` or `np.ndarray`) Uncalibrated uncertainty(MAD, median absolute deviation)
+                    coming from ensembling. None if `ensemble_size = 1`
+        """
+        # Model-specific optimal default values leading to fast and reasonable results.
+        if denoising_steps is None:
+            denoising_steps = self.default_denoising_steps
+        if processing_res is None:
+            processing_res = self.default_processing_resolution
+
+        assert processing_res >= 0
+        assert ensemble_size >= 1
+
+        # Check if denoising step is reasonable
+        self._check_inference_step(denoising_steps)
+
+        resample_method: InterpolationMode = get_tv_resample_method(resample_method)
+
+        # ----------------- Image Preprocess -----------------
+        # Convert to torch tensor
+        if isinstance(input_image, Image.Image):
+            input_image = input_image.convert("RGB")
+            # convert to torch tensor [H, W, rgb] -> [rgb, H, W]
+            rgb = pil_to_tensor(input_image)
+            rgb = rgb.unsqueeze(0)  # [1, rgb, H, W]
+        elif isinstance(input_image, torch.Tensor):
+            rgb = input_image
+        else:
+            raise TypeError(f"Unknown input type: {type(input_image) = }")
+        input_size = rgb.shape
+        assert (
+            4 == rgb.dim() and 3 == input_size[-3]
+        ), f"Wrong input shape {input_size}, expected [1, rgb, H, W]"
+
+        # Resize image
+        if processing_res > 0:
+            rgb = resize_max_res(
+                rgb,
+                max_edge_resolution=processing_res,
+                resample_method=resample_method,
+            )
+
+        # Normalize rgb values
+        rgb_norm: torch.Tensor = rgb / 255.0 * 2.0 - 1.0  #  [0, 255] -> [-1, 1]
+        rgb_norm = rgb_norm.to(self.dtype)
+        assert rgb_norm.min() >= -1.0 and rgb_norm.max() <= 1.0
+
+        # ----------------- Predicting normals -----------------
+        # Batch repeated input image
+        duplicated_rgb = rgb_norm.expand(ensemble_size, -1, -1, -1)
+        single_rgb_dataset = TensorDataset(duplicated_rgb)
+        if batch_size > 0:
+            _bs = batch_size
+        else:
+            _bs = find_batch_size(
+                ensemble_size=ensemble_size,
+                input_res=max(rgb_norm.shape[1:]),
+                dtype=self.dtype,
+            )
+
+        single_rgb_loader = DataLoader(
+            single_rgb_dataset, batch_size=_bs, shuffle=False
+        )
+
+        # Predict normals maps (batched)
+        target_pred_ls = []
+        if show_progress_bar:
+            iterable = tqdm(
+                single_rgb_loader, desc=" " * 2 + "Inference batches", leave=False
+            )
+        else:
+            iterable = single_rgb_loader
+        for batch in iterable:
+            (batched_img,) = batch
+            target_pred_raw = self.single_infer(
+                rgb_in=batched_img,
+                num_inference_steps=denoising_steps,
+                show_pbar=show_progress_bar,
+                generator=generator,
+            )
+            target_pred_ls.append(target_pred_raw.detach())
+        target_preds = torch.concat(target_pred_ls, dim=0)
+        torch.cuda.empty_cache()  # clear vram cache for ensembling
+
+        # ----------------- Test-time ensembling -----------------
+        if ensemble_size > 1:
+            final_pred, pred_uncert = ensemble_normals(
+                target_preds,
+                **(ensemble_kwargs or {}),
+            )
+        else:
+            final_pred = target_preds
+            pred_uncert = None
+
+        # Resize back to original resolution
+        if match_input_res:
+            final_pred = resize(
+                final_pred,
+                input_size[-2:],
+                interpolation=resample_method,
+                antialias=True,
+            )
+
+        # Convert to numpy
+        final_pred = final_pred.squeeze()
+        final_pred = final_pred.cpu().numpy()
+        if pred_uncert is not None:
+            pred_uncert = pred_uncert.squeeze().cpu().numpy()
+
+        # Clip output range
+        final_pred = final_pred.clip(-1, 1)
+
+        # Colorize
+        normals_img = ((final_pred + 1) * 127.5).astype(np.uint8)
+        normals_img = chw2hwc(normals_img)
+        normals_img = Image.fromarray(normals_img)
+
+        return MarigoldNormalsOutput(
+            normals_np=final_pred,
+            normals_img=normals_img,
+            uncertainty=pred_uncert,
+        )
+
+    def _check_inference_step(self, n_step: int) -> None:
+        """
+        Check if denoising step is reasonable
+        Args:
+            n_step (`int`): denoising steps
+        """
+        assert n_step >= 1
+
+        if isinstance(self.scheduler, DDIMScheduler):
+            if "trailing" != self.scheduler.config.timestep_spacing:
+                logging.warning(
+                    f"The loaded `DDIMScheduler` is configured with `timestep_spacing="
+                    f'"{self.scheduler.config.timestep_spacing}"`; the recommended setting is `"trailing"`. '
+                    f"This change is backward-compatible and yields better results. "
+                    f"Consider using `prs-eth/marigold-normals-v1-1` for the best experience."
+                )
+            else:
+                if n_step > 10:
+                    logging.warning(
+                        f"Setting too many denoising steps ({n_step}) may degrade the prediction; consider relying on "
+                        f"the default values."
+                    )
+            if not self.scheduler.config.rescale_betas_zero_snr:
+                logging.warning(
+                    f"The loaded `DDIMScheduler` is configured with `rescale_betas_zero_snr="
+                    f"{self.scheduler.config.rescale_betas_zero_snr}`; the recommended setting is True. "
+                    f"Consider using `prs-eth/marigold-normals-v1-1` for the best experience."
+                )
+        elif isinstance(self.scheduler, LCMScheduler):
+            raise RuntimeError(
+                "This pipeline implementation does not support the LCMScheduler. Please refer to the project "
+                "README.md for instructions about using LCM."
+            )
+        else:
+            raise RuntimeError(f"Unsupported scheduler type: {type(self.scheduler)}")
+
+    def encode_empty_text(self):
+        """
+        Encode text embedding for empty prompt
+        """
+        prompt = ""
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="do_not_pad",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids.to(self.text_encoder.device)
+        self.empty_text_embed = self.text_encoder(text_input_ids)[0].to(self.dtype)
+
+    @torch.no_grad()
+    def single_infer(
+        self,
+        rgb_in: torch.Tensor,
+        num_inference_steps: int,
+        generator: Union[torch.Generator, None],
+        show_pbar: bool,
+    ) -> torch.Tensor:
+        """
+        Perform a single prediction without ensembling.
+
+        Args:
+            rgb_in (`torch.Tensor`):
+                Input RGB image.
+            num_inference_steps (`int`):
+                Number of diffusion denoisign steps (DDIM) during inference.
+            show_pbar (`bool`):
+                Display a progress bar of diffusion denoising.
+            generator (`torch.Generator`)
+                Random generator for initial noise generation.
+        Returns:
+            `torch.Tensor`: Predicted targets.
+        """
+        device = self.device
+        rgb_in = rgb_in.to(device)
+
+        # Set timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps  # [T]
+
+        # Encode image
+        rgb_latent = self.encode_rgb(rgb_in)  # [B, 4, h, w]
+
+        # Noisy latent for outputs
+        target_latent = torch.randn(
+            rgb_latent.shape,
+            device=device,
+            dtype=self.dtype,
+            generator=generator,
+        )  # [B, 4, h, w]
+
+        # Batched empty text embedding
+        if self.empty_text_embed is None:
+            self.encode_empty_text()
+        batch_empty_text_embed = self.empty_text_embed.repeat(
+            (rgb_latent.shape[0], 1, 1)
+        ).to(device)  # [B, 2, 1024]
+
+        # Denoising loop
+        if show_pbar:
+            iterable = tqdm(
+                enumerate(timesteps),
+                total=len(timesteps),
+                leave=False,
+                desc=" " * 4 + "Diffusion denoising",
+            )
+        else:
+            iterable = enumerate(timesteps)
+
+        for i, t in iterable:
+            unet_input = torch.cat(
+                [rgb_latent, target_latent], dim=1
+            )  # this order is important
+
+            # predict the noise residual
+            noise_pred = self.unet(
+                unet_input, t, encoder_hidden_states=batch_empty_text_embed
+            ).sample  # [B, 4, h, w]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            target_latent = self.scheduler.step(
+                noise_pred, t, target_latent, generator=generator
+            ).prev_sample
+
+        normals = self.decode_normals(target_latent)  # [B,3,H,W]
+
+        # clip prediction
+        normals = torch.clip(normals, -1.0, 1.0)
+        norm = torch.norm(normals, dim=1, keepdim=True)
+        normals /= norm.clamp(min=1e-6)  # [B,3,H,W]
+
+        return normals
+
+    def encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor:
+        """
+        Encode RGB image into latent.
+
+        Args:
+            rgb_in (`torch.Tensor`):
+                Input RGB image to be encoded.
+
+        Returns:
+            `torch.Tensor`: Image latent.
+        """
+        # encode
+        h = self.vae.encoder(rgb_in)
+        moments = self.vae.quant_conv(h)
+        mean, logvar = torch.chunk(moments, 2, dim=1)
+        # scale latent
+        rgb_latent = mean * self.latent_scale_factor
+        return rgb_latent
+
+    def decode_normals(self, normals_latent: torch.Tensor) -> torch.Tensor:
+        """
+        Decode normals latent into normals map.
+
+        Args:
+            normals_latent (`torch.Tensor`):
+                Normals latent to be decoded.
+
+        Returns:
+            `torch.Tensor`: Decoded normals map.
+        """
+        # scale latent
+        normals_latent = normals_latent / self.latent_scale_factor
+        # decode
+        z = self.vae.post_quant_conv(normals_latent)
+        stacked = self.vae.decoder(z)
+        return stacked

marigold/util/batchsize.py

@@ -0,0 +1,90 @@
+# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# --------------------------------------------------------------------------
+# More information about Marigold:
+#   https://marigoldmonodepth.github.io
+#   https://marigoldcomputervision.github.io
+# Efficient inference pipelines are now part of diffusers:
+#   https://huggingface.co/docs/diffusers/using-diffusers/marigold_usage
+#   https://huggingface.co/docs/diffusers/api/pipelines/marigold
+# Examples of trained models and live demos:
+#   https://huggingface.co/prs-eth
+# Related projects:
+#   https://rollingdepth.github.io/
+#   https://marigolddepthcompletion.github.io/
+# Citation (BibTeX):
+#   https://github.com/prs-eth/Marigold#-citation
+# If you find Marigold useful, we kindly ask you to cite our papers.
+# --------------------------------------------------------------------------
+
+import math
+import torch
+
+# Search table for suggested max. inference batch size
+bs_search_table = [
+    # tested on A100-PCIE-80GB
+    {"res": 768, "total_vram": 79, "bs": 35, "dtype": torch.float32},
+    {"res": 1024, "total_vram": 79, "bs": 20, "dtype": torch.float32},
+    # tested on A100-PCIE-40GB
+    {"res": 768, "total_vram": 39, "bs": 15, "dtype": torch.float32},
+    {"res": 1024, "total_vram": 39, "bs": 8, "dtype": torch.float32},
+    {"res": 768, "total_vram": 39, "bs": 30, "dtype": torch.float16},
+    {"res": 1024, "total_vram": 39, "bs": 15, "dtype": torch.float16},
+    # tested on RTX3090, RTX4090
+    {"res": 512, "total_vram": 23, "bs": 20, "dtype": torch.float32},
+    {"res": 768, "total_vram": 23, "bs": 7, "dtype": torch.float32},
+    {"res": 1024, "total_vram": 23, "bs": 3, "dtype": torch.float32},
+    {"res": 512, "total_vram": 23, "bs": 40, "dtype": torch.float16},
+    {"res": 768, "total_vram": 23, "bs": 18, "dtype": torch.float16},
+    {"res": 1024, "total_vram": 23, "bs": 10, "dtype": torch.float16},
+    # tested on GTX1080Ti
+    {"res": 512, "total_vram": 10, "bs": 5, "dtype": torch.float32},
+    {"res": 768, "total_vram": 10, "bs": 2, "dtype": torch.float32},
+    {"res": 512, "total_vram": 10, "bs": 10, "dtype": torch.float16},
+    {"res": 768, "total_vram": 10, "bs": 5, "dtype": torch.float16},
+    {"res": 1024, "total_vram": 10, "bs": 3, "dtype": torch.float16},
+]
+
+
+def find_batch_size(ensemble_size: int, input_res: int, dtype: torch.dtype) -> int:
+    """
+    Automatically search for suitable operating batch size.
+
+    Args:
+        ensemble_size (`int`):
+            Number of predictions to be ensembled.
+        input_res (`int`):
+            Operating resolution of the input image.
+
+    Returns:
+        `int`: Operating batch size.
+    """
+    if not torch.cuda.is_available():
+        return 1
+
+    total_vram = torch.cuda.mem_get_info()[1] / 1024.0**3
+    filtered_bs_search_table = [s for s in bs_search_table if s["dtype"] == dtype]
+    for settings in sorted(
+        filtered_bs_search_table,
+        key=lambda k: (k["res"], -k["total_vram"]),
+    ):
+        if input_res <= settings["res"] and total_vram >= settings["total_vram"]:
+            bs = settings["bs"]
+            if bs > ensemble_size:
+                bs = ensemble_size
+            elif bs > math.ceil(ensemble_size / 2) and bs < ensemble_size:
+                bs = math.ceil(ensemble_size / 2)
+            return bs
+
+    return 1

marigold/util/ensemble.py

@@ -0,0 +1,270 @@
+# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# --------------------------------------------------------------------------
+# More information about Marigold:
+#   https://marigoldmonodepth.github.io
+#   https://marigoldcomputervision.github.io
+# Efficient inference pipelines are now part of diffusers:
+#   https://huggingface.co/docs/diffusers/using-diffusers/marigold_usage
+#   https://huggingface.co/docs/diffusers/api/pipelines/marigold
+# Examples of trained models and live demos:
+#   https://huggingface.co/prs-eth
+# Related projects:
+#   https://rollingdepth.github.io/
+#   https://marigolddepthcompletion.github.io/
+# Citation (BibTeX):
+#   https://github.com/prs-eth/Marigold#-citation
+# If you find Marigold useful, we kindly ask you to cite our papers.
+# --------------------------------------------------------------------------
+
+import numpy as np
+import torch
+from functools import partial
+from typing import Optional, Tuple
+
+from .image_util import get_tv_resample_method, resize_max_res
+
+
+def ensemble_depth(
+    depth: torch.Tensor,
+    scale_invariant: bool = True,
+    shift_invariant: bool = True,
+    output_uncertainty: bool = False,
+    reduction: str = "median",
+    regularizer_strength: float = 0.02,
+    max_iter: int = 50,
+    tol: float = 1e-6,
+    max_res: int = 1024,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Ensembles depth maps represented by the `depth` tensor with expected shape `(B, 1, H, W)`, where B is the
+    number of ensemble members for a given prediction of size `(H x W)`. Even though the function is designed for
+    depth maps, it can also be used with disparity maps as long as the input tensor values are non-negative. The
+    alignment happens when the predictions have one or more degrees of freedom, that is when they are either
+    affine-invariant (`scale_invariant=True` and `shift_invariant=True`), or just scale-invariant (only
+    `scale_invariant=True`). For absolute predictions (`scale_invariant=False` and `shift_invariant=False`)
+    alignment is skipped and only ensembling is performed.
+
+    Args:
+        depth (`torch.Tensor`):
+            Input ensemble depth maps.
+        scale_invariant (`bool`, *optional*, defaults to `True`):
+            Whether to treat predictions as scale-invariant.
+        shift_invariant (`bool`, *optional*, defaults to `True`):
+            Whether to treat predictions as shift-invariant.
+        output_uncertainty (`bool`, *optional*, defaults to `False`):
+            Whether to output uncertainty map.
+        reduction (`str`, *optional*, defaults to `"median"`):
+            Reduction method used to ensemble aligned predictions. The accepted values are: `"mean"` and
+            `"median"`.
+        regularizer_strength (`float`, *optional*, defaults to `0.02`):
+            Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1.
+        max_iter (`int`, *optional*, defaults to `2`):
+            Maximum number of the alignment solver steps. Refer to `scipy.optimize.minimize` function, `options`
+            argument.
+        tol (`float`, *optional*, defaults to `1e-3`):
+            Alignment solver tolerance. The solver stops when the tolerance is reached.
+        max_res (`int`, *optional*, defaults to `1024`):
+            Resolution at which the alignment is performed; `None` matches the `processing_resolution`.
+    Returns:
+        A tensor of aligned and ensembled depth maps and optionally a tensor of uncertainties of the same shape:
+        `(1, 1, H, W)`.
+    """
+    if depth.dim() != 4 or depth.shape[1] != 1:
+        raise ValueError(f"Expecting 4D tensor of shape [B,1,H,W]; got {depth.shape}.")
+    if reduction not in ("mean", "median"):
+        raise ValueError(f"Unrecognized reduction method: {reduction}.")
+    if not scale_invariant and shift_invariant:
+        raise ValueError("Pure shift-invariant ensembling is not supported.")
+
+    def init_param(depth: torch.Tensor):
+        init_min = depth.reshape(ensemble_size, -1).min(dim=1).values
+        init_max = depth.reshape(ensemble_size, -1).max(dim=1).values
+
+        if scale_invariant and shift_invariant:
+            init_s = 1.0 / (init_max - init_min).clamp(min=1e-6)
+            init_t = -init_s * init_min
+            param = torch.cat((init_s, init_t)).cpu().numpy()
+        elif scale_invariant:
+            init_s = 1.0 / init_max.clamp(min=1e-6)
+            param = init_s.cpu().numpy()
+        else:
+            raise ValueError("Unrecognized alignment.")
+
+        return param.astype(np.float64)
+
+    def align(depth: torch.Tensor, param: np.ndarray) -> torch.Tensor:
+        if scale_invariant and shift_invariant:
+            s, t = np.split(param, 2)
+            s = torch.from_numpy(s).to(depth).view(ensemble_size, 1, 1, 1)
+            t = torch.from_numpy(t).to(depth).view(ensemble_size, 1, 1, 1)
+            out = depth * s + t
+        elif scale_invariant:
+            s = torch.from_numpy(param).to(depth).view(ensemble_size, 1, 1, 1)
+            out = depth * s
+        else:
+            raise ValueError("Unrecognized alignment.")
+        return out
+
+    def ensemble(
+        depth_aligned: torch.Tensor, return_uncertainty: bool = False
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        uncertainty = None
+        if reduction == "mean":
+            prediction = torch.mean(depth_aligned, dim=0, keepdim=True)
+            if return_uncertainty:
+                uncertainty = torch.std(depth_aligned, dim=0, keepdim=True)
+        elif reduction == "median":
+            prediction = torch.median(depth_aligned, dim=0, keepdim=True).values
+            if return_uncertainty:
+                uncertainty = torch.median(
+                    torch.abs(depth_aligned - prediction), dim=0, keepdim=True
+                ).values
+        else:
+            raise ValueError(f"Unrecognized reduction method: {reduction}.")
+        return prediction, uncertainty
+
+    def cost_fn(param: np.ndarray, depth: torch.Tensor) -> float:
+        cost = 0.0
+        depth_aligned = align(depth, param)
+
+        for i, j in torch.combinations(torch.arange(ensemble_size)):
+            diff = depth_aligned[i] - depth_aligned[j]
+            cost += (diff**2).mean().sqrt().item()
+
+        if regularizer_strength > 0:
+            prediction, _ = ensemble(depth_aligned, return_uncertainty=False)
+            err_near = (0.0 - prediction.min()).abs().item()
+            err_far = (1.0 - prediction.max()).abs().item()
+            cost += (err_near + err_far) * regularizer_strength
+
+        return cost
+
+    def compute_param(depth: torch.Tensor):
+        import scipy
+
+        depth_to_align = depth.to(torch.float32)
+        if max_res is not None and max(depth_to_align.shape[2:]) > max_res:
+            depth_to_align = resize_max_res(
+                depth_to_align, max_res, get_tv_resample_method("nearest-exact")
+            )
+
+        param = init_param(depth_to_align)
+
+        res = scipy.optimize.minimize(
+            partial(cost_fn, depth=depth_to_align),
+            param,
+            method="BFGS",
+            tol=tol,
+            options={"maxiter": max_iter, "disp": False},
+        )
+
+        return res.x
+
+    requires_aligning = scale_invariant or shift_invariant
+    ensemble_size = depth.shape[0]
+
+    if requires_aligning:
+        param = compute_param(depth)
+        depth = align(depth, param)
+
+    depth, uncertainty = ensemble(depth, return_uncertainty=output_uncertainty)
+
+    depth_max = depth.max()
+    if scale_invariant and shift_invariant:
+        depth_min = depth.min()
+    elif scale_invariant:
+        depth_min = 0
+    else:
+        raise ValueError("Unrecognized alignment.")
+    depth_range = (depth_max - depth_min).clamp(min=1e-6)
+    depth = (depth - depth_min) / depth_range
+    if output_uncertainty:
+        uncertainty /= depth_range
+
+    return depth, uncertainty  # [1,1,H,W], [1,1,H,W]
+
+
+def ensemble_normals(
+    normals: torch.Tensor,
+    output_uncertainty: bool = False,
+    reduction: str = "closest",
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
+    the number of ensemble members for a given prediction of size `(H x W)`.
+
+    Args:
+        normals (`torch.Tensor`):
+            Input ensemble normals maps.
+        output_uncertainty (`bool`, *optional*, defaults to `False`):
+            Whether to output uncertainty map.
+        reduction (`str`, *optional*, defaults to `"closest"`):
+            Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and
+            `"mean"`.
+
+    Returns:
+        A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of
+        uncertainties of shape `(1, 1, H, W)`.
+    """
+    if normals.dim() != 4 or normals.shape[1] != 3:
+        raise ValueError(
+            f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}."
+        )
+    if reduction not in ("closest", "mean"):
+        raise ValueError(f"Unrecognized reduction method: {reduction}.")
+
+    mean_normals = normals.mean(dim=0, keepdim=True)  # [1,3,H,W]
+    norm = torch.norm(mean_normals, dim=1, keepdim=True)
+    mean_normals /= norm.clamp(min=1e-6)  # [1,3,H,W]
+
+    sim_cos = None
+    if output_uncertainty or (reduction != "mean"):
+        sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True)  # [E,1,H,W]
+        sim_cos = sim_cos.clamp(-1, 1)  # required to avoid NaN in uncertainty with fp16
+
+    uncertainty = None
+    if output_uncertainty:
+        uncertainty = sim_cos.arccos()  # [E,1,H,W]
+        uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi  # [1,1,H,W]
+
+    if reduction == "mean":
+        return mean_normals, uncertainty  # [1,3,H,W], [1,1,H,W]
+
+    closest_indices = sim_cos.argmax(dim=0, keepdim=True)  # [1,1,H,W]
+    closest_indices = closest_indices.repeat(1, 3, 1, 1)  # [1,3,H,W]
+    closest_normals = torch.gather(normals, 0, closest_indices)  # [1,3,H,W]
+
+    return closest_normals, uncertainty  # [1,3,H,W], [1,1,H,W]
+
+
+def ensemble_iid(
+    targets: torch.Tensor,
+    output_uncertainty: bool = False,
+    reduction: str = "median",
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    uncertainty = None
+    if reduction == "mean":
+        prediction = torch.mean(targets, dim=0, keepdim=True)
+        if output_uncertainty:
+            uncertainty = torch.std(targets, dim=0, keepdim=True)
+    elif reduction == "median":
+        prediction = torch.median(targets, dim=0, keepdim=True).values
+        if output_uncertainty:
+            uncertainty = torch.median(
+                torch.abs(targets - prediction), dim=0, keepdim=True
+            ).values
+    else:
+        raise ValueError(f"Unrecognized reduction method: {reduction}.")
+    return prediction, uncertainty

requirements++.txt

@@ -0,0 +1,7 @@
+h5py
+opencv-python
+tensorboard
+wandb
+scikit-learn
+xformers==0.0.28
+

requirements+.txt

@@ -0,0 +1,4 @@
+omegaconf
+pandas
+tabulate
+torchmetrics