JaehaerysNL/StableDiffusion
1
0
Fork 0
mirror of https://github.com/Stability-AI/stablediffusion.git synced 2024-12-22 07:34:58 +00:00
Code Issues Projects Releases Packages Wiki Activity

Added type hints and improved error message in append_dims function

Browse source
This commit is contained in:
Mandlin Sarah 2024-08-31 03:15:35 -07:00
parent cf1d67a6fd
commit 199b6a7a6c
1 changed files with 6 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
ldm/models/diffusion/sampling_util.py
View file

@ -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)
return x0 * (value / s)