guangkaixu commited on
Commit
d2c1576
1 Parent(s): 72f6757
Files changed (1) hide show
  1. genpercept/genpercept_pipeline.py +1 -1
genpercept/genpercept_pipeline.py CHANGED
@@ -315,7 +315,7 @@ class GenPerceptPipeline(DiffusionPipeline):
315
 
316
  # Colorize
317
  if color_map is not None:
318
- assert self.mode == 'depth'
319
  pred_colored = colorize_depth_maps(
320
  pipe_pred, 0, 1, cmap=color_map
321
  ).squeeze() # [3, H, W], value in (0, 1)
 
315
 
316
  # Colorize
317
  if color_map is not None:
318
+ assert self.mode in ['depth', 'disparity']
319
  pred_colored = colorize_depth_maps(
320
  pipe_pred, 0, 1, cmap=color_map
321
  ).squeeze() # [3, H, W], value in (0, 1)