From 199b6a7a6cd6fca86319c02e405fe0d9817b2067 Mon Sep 17 00:00:00 2001 From: Mandlin Sarah Date: Sat, 31 Aug 2024 03:15:35 -0700 Subject: [PATCH] Added type hints and improved error message in append_dims function --- ldm/models/diffusion/sampling_util.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/ldm/models/diffusion/sampling_util.py b/ldm/models/diffusion/sampling_util.py index 7eff02b..f2b5d81 100644 --- a/ldm/models/diffusion/sampling_util.py +++ b/ldm/models/diffusion/sampling_util.py @@ -2,21 +2,22 @@ import torch import numpy as np -def append_dims(x, target_dims): +def append_dims(x: torch.Tensor, target_dims: int) -> torch.Tensor: """Appends dimensions to the end of a tensor until it has target_dims dimensions. From https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/utils.py""" dims_to_append = target_dims - x.ndim if dims_to_append < 0: - raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less') + raise ValueError(f'Input tensor has {x.ndim} dimensions but target_dims is {target_dims}, which is less than the number of dimensions in the input tensor.') return x[(...,) + (None,) * dims_to_append] -def norm_thresholding(x0, value): +def norm_thresholding(x0: torch.Tensor, value: float) -> torch.Tensor: s = append_dims(x0.pow(2).flatten(1).mean(1).sqrt().clamp(min=value), x0.ndim) return x0 * (value / s) -def spatial_norm_thresholding(x0, value): +def spatial_norm_thresholding(x0: torch.Tensor, value: float) -> torch.Tensor: # b c h w s = x0.pow(2).mean(1, keepdim=True).sqrt().clamp(min=value) - return x0 * (value / s) \ No newline at end of file + return x0 * (value / s) +