Update models/depth_normal_pipeline_clip_cfg_1.py
Browse files
models/depth_normal_pipeline_clip_cfg_1.py
CHANGED
|
@@ -73,12 +73,13 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
|
|
| 73 |
@torch.no_grad()
|
| 74 |
def __call__(self,
|
| 75 |
input_image:Image,
|
| 76 |
-
|
| 77 |
ensemble_size: int = 10,
|
| 78 |
processing_res: int = 768,
|
| 79 |
match_input_res:bool =True,
|
| 80 |
batch_size:int = 0,
|
| 81 |
domain: str = "indoor",
|
|
|
|
| 82 |
color_map: str="Spectral",
|
| 83 |
show_progress_bar:bool = True,
|
| 84 |
ensemble_kwargs: Dict = None,
|
|
@@ -95,7 +96,7 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
|
|
| 95 |
)," Value Error: `resize_output_back` is only valid with "
|
| 96 |
|
| 97 |
assert processing_res >=0
|
| 98 |
-
assert
|
| 99 |
assert ensemble_size >=1
|
| 100 |
|
| 101 |
# --------------- Image Processing ------------------------
|
|
@@ -145,8 +146,9 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
|
|
| 145 |
|
| 146 |
depth_pred_raw, normal_pred_raw = self.single_infer(
|
| 147 |
input_rgb=batched_image,
|
| 148 |
-
num_inference_steps=
|
| 149 |
domain=domain,
|
|
|
|
| 150 |
show_pbar=show_progress_bar,
|
| 151 |
)
|
| 152 |
depth_pred_ls.append(depth_pred_raw.detach().clone())
|
|
@@ -230,6 +232,7 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
|
|
| 230 |
def single_infer(self,input_rgb:torch.Tensor,
|
| 231 |
num_inference_steps:int,
|
| 232 |
domain:str,
|
|
|
|
| 233 |
show_pbar:bool,):
|
| 234 |
|
| 235 |
device = input_rgb.device
|
|
@@ -286,7 +289,7 @@ class DepthNormalEstimationPipeline(DiffusionPipeline):
|
|
| 286 |
# predict the noise residual
|
| 287 |
noise_pred = self.unet(unet_input, t.repeat(4), encoder_hidden_states=batch_img_embed, class_labels=class_embedding.repeat(2,1)).sample
|
| 288 |
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
|
| 289 |
-
guidance_scale =
|
| 290 |
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
|
| 291 |
|
| 292 |
# compute the previous noisy sample x_t -> x_t-1
|
|
|
|
| 73 |
@torch.no_grad()
|
| 74 |
def __call__(self,
|
| 75 |
input_image:Image,
|
| 76 |
+
denoising_steps: int = 10,
|
| 77 |
ensemble_size: int = 10,
|
| 78 |
processing_res: int = 768,
|
| 79 |
match_input_res:bool =True,
|
| 80 |
batch_size:int = 0,
|
| 81 |
domain: str = "indoor",
|
| 82 |
+
guidance_scale: int = 3,
|
| 83 |
color_map: str="Spectral",
|
| 84 |
show_progress_bar:bool = True,
|
| 85 |
ensemble_kwargs: Dict = None,
|
|
|
|
| 96 |
)," Value Error: `resize_output_back` is only valid with "
|
| 97 |
|
| 98 |
assert processing_res >=0
|
| 99 |
+
assert denoising_steps >=1
|
| 100 |
assert ensemble_size >=1
|
| 101 |
|
| 102 |
# --------------- Image Processing ------------------------
|
|
|
|
| 146 |
|
| 147 |
depth_pred_raw, normal_pred_raw = self.single_infer(
|
| 148 |
input_rgb=batched_image,
|
| 149 |
+
num_inference_steps=denoising_steps,
|
| 150 |
domain=domain,
|
| 151 |
+
guidance_scale=guidance_scale,
|
| 152 |
show_pbar=show_progress_bar,
|
| 153 |
)
|
| 154 |
depth_pred_ls.append(depth_pred_raw.detach().clone())
|
|
|
|
| 232 |
def single_infer(self,input_rgb:torch.Tensor,
|
| 233 |
num_inference_steps:int,
|
| 234 |
domain:str,
|
| 235 |
+
guidance_scale:int,
|
| 236 |
show_pbar:bool,):
|
| 237 |
|
| 238 |
device = input_rgb.device
|
|
|
|
| 289 |
# predict the noise residual
|
| 290 |
noise_pred = self.unet(unet_input, t.repeat(4), encoder_hidden_states=batch_img_embed, class_labels=class_embedding.repeat(2,1)).sample
|
| 291 |
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
|
| 292 |
+
guidance_scale = guidance_scale
|
| 293 |
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
|
| 294 |
|
| 295 |
# compute the previous noisy sample x_t -> x_t-1
|