import json import argparse import random import numpy as np import torch.nn.functional as F import torch from torch import nn import cv2 from scipy.spatial import KDTree from functools import partial def save_args(args, to_path): with open(to_path, "w") as f: json.dump(args.__dict__, f, indent=2) def load_args(from_path, is_test=True): parser = argparse.ArgumentParser() args = parser.parse_args() with open(from_path, "r") as f: args.__dict__ = json.load(f) args.is_test = is_test if "E_name" not in args.__dict__.keys(): args.E_name = "basic" return args def tensor2img(x, round=False): ''' x : [BS x c x H x W] or [c x H x W] ''' if x.ndim == 3: x = x.unsqueeze(0) BS, C, H, W = x.shape x = x.permute(0,2,3,1).reshape(-1, W, C).detach().cpu().numpy() # x = (x+1)/2 # x = np.clip(x, 0, 1) x = np.clip(x, -1, 1) x = (x+1)/2 adjust = 0.5 if round else 0 x = np.uint8(x*255.0 + adjust) if x.shape[-1] == 1: # gray sclae x = np.concatenate([x,x,x], axis=-1) return x def resize_mask(m, shape): m = F.interpolate(m, shape) m[m > 0.5] = 1 m[m < 0.5] = 0 return m def remove_overlap(seg_out, warped_cm, inference=False): assert len(warped_cm.shape) == 4 overlapped_region = (torch.cat([seg_out[:, 1:3, :, :], seg_out[:, 5:, :, :]], dim=1)).sum(dim=1, keepdim=True) if inference: overlapped_region = (overlapped_region > 0.5).float() warped_cm = warped_cm - overlapped_region * warped_cm return warped_cm def bilateral_filter(image, kernel_size, sigma_d, sigma_r): """Bilateral filter implementation. Args: image: input float tensor with shape [bsz, ch, height, width] kernel_size: int. we assume it is odd. sigma_d: float. standard deviation for distance. sigma_r: float or tensor. standard deviation for range. """ padding = (kernel_size - 1) // 2 # distance bsz, ch, height, width = image.shape if isinstance(sigma_r, float): sigma_r = torch.tensor([sigma_r]).expand(bsz) sigma_r = sigma_r.to(image.device) height_pad = height + 2 * padding width_pad = width + 2 * padding # gaussian on spacial distance grid_x, grid_y = torch.meshgrid(torch.arange(width_pad), torch.arange(height_pad), indexing='xy') grid_x = grid_x.float().to(image.device) grid_y = grid_y.float().to(image.device) unfold_grid = nn.Unfold(kernel_size=kernel_size) grid_x_unfolded = unfold_grid(grid_x[None, None]) grid_y_unfolded = unfold_grid(grid_y[None, None]) grid_x_unfolded = grid_x_unfolded.transpose(1, 2).reshape(height * width, 1, kernel_size ** 2) grid_y_unfolded = grid_y_unfolded.transpose(1, 2).reshape(height * width, 1, kernel_size ** 2) center_index = kernel_size ** 2 // 2 diff_x_unfolded = grid_x_unfolded - grid_x_unfolded[:, :, center_index][:, :, None] diff_y_unfolded = grid_y_unfolded - grid_y_unfolded[:, :, center_index][:, :, None] dist_unfolded = diff_x_unfolded ** 2 + diff_y_unfolded ** 2 gaussian_dist = torch.exp(-dist_unfolded / (2 * sigma_d ** 2)) # gaussian on range unfold = nn.Unfold(kernel_size=kernel_size, padding=padding) image_unfolded = unfold(image) image_unfolded = image_unfolded.transpose(1, 2).reshape(bsz, height * width, ch, kernel_size ** 2) center_value = image_unfolded[:, :, :, center_index] diff_value = image_unfolded - center_value[:, :, :, None] dist_value = diff_value ** 2 gaussian_value = torch.exp(-dist_value / (2 * sigma_r[:, None, None, None] ** 2)) # bilateral filter bilateral_weight = gaussian_dist[None] * gaussian_value result_unfolded = torch.sum(bilateral_weight * image_unfolded, dim=-1) z_constant = bilateral_weight.sum(dim=-1) result_unfolded = result_unfolded / z_constant result = result_unfolded.transpose(1, 2).reshape(bsz, ch, height, width) return result def ensure_tensor(x): if isinstance(x, torch.Tensor): return x # numpy if isinstance(x, np.ndarray): return torch.from_numpy(x) return torch.tensor(x) def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) # SIFT-VTON SIFT matching and filtering ANGLE_THRES = 45 def filter_angle(keypoints1, keypoints2, match): idx1 = match.queryIdx idx2 = match.trainIdx angle1 = keypoints1[idx1].angle angle2 = keypoints2[idx2].angle diff = abs(angle1 - angle2) if diff > 180: diff = 360 - diff return diff < ANGLE_THRES def round_tuple(t): return tuple(map(lambda x: round(x), t)) def transpose_pt(pt): return (pt[1], pt[0]) SCALE_RATIO_THRES = 2.25 def filter_scale(keypoints1, keypoints2, match): idx1 = match.queryIdx idx2 = match.trainIdx size1 = keypoints1[idx1].size size2 = keypoints2[idx2].size scale_ratio = size1 / size2 if 1/SCALE_RATIO_THRES <= scale_ratio <= SCALE_RATIO_THRES: return True return False # unused PIXEL_DIST_THRES = 200 def filter_value(keypoints1, keypoints2, img1, img2, match): idx1 = match.queryIdx idx2 = match.trainIdx loc1 = transpose_pt(round_tuple(keypoints1[idx1].pt)) loc2 = transpose_pt(round_tuple(keypoints2[idx2].pt)) value1 = img1[loc1].astype(np.int16) value2 = img2[loc2].astype(np.int16) dist = np.linalg.norm(value1 - value2, ord=1) if dist < PIXEL_DIST_THRES: return True return False PIXEL_HUE_DIFF_THRES = 15 PIXEL_SAT_EFFECTIVE_RATIO = 0.1 PIXEL_VAL_DIFF_RATIO = 0.4 PIXEL_VAL_EFFECTIVE_RATIO = 0.2 def filter_hsv(keypoints1, keypoints2, img1, img2, match): idx1 = match.queryIdx idx2 = match.trainIdx loc1 = transpose_pt(round_tuple(keypoints1[idx1].pt)) loc2 = transpose_pt(round_tuple(keypoints2[idx2].pt)) value1 = img1[loc1] value2 = img2[loc2] h1, s1, v1 = cv2.cvtColor(value1[None, None], cv2.COLOR_RGB2HSV)[0, 0] h2, s2, v2 = cv2.cvtColor(value2[None, None], cv2.COLOR_RGB2HSV)[0, 0] v1 = v1.astype(np.int16) v2 = v2.astype(np.int16) diff = abs(v1 - v2) if diff > PIXEL_VAL_DIFF_RATIO * 255: return False if s1 < 255 * PIXEL_SAT_EFFECTIVE_RATIO and s2 < 255 * PIXEL_SAT_EFFECTIVE_RATIO: return True if v1 < 255 * PIXEL_VAL_EFFECTIVE_RATIO and v2 < 255 * PIXEL_VAL_EFFECTIVE_RATIO: return True h1 = h1.astype(np.int16) h2 = h2.astype(np.int16) diff = abs(h1 - h2) if diff > 90: diff = 180 - diff if diff < PIXEL_HUE_DIFF_THRES: return True return False PIXEL_DIST_MULTIPLIER2 = 1.4 NUM_FILTER_NEAREST_ITTER = 1 def filter_nearest_dist_from2(keypoints1, keypoints2, tree, matches, sift_match, num_neighbors=1): idx1 = sift_match.queryIdx idx2 = sift_match.trainIdx pt1 = keypoints1[idx1].pt pt2 = keypoints2[idx2].pt dist2, ii = tree.query(pt2, k=num_neighbors + 1) for i in range(num_neighbors): if dist2[i+1] == np.inf: continue pt1_2 = keypoints1[matches[ii[i+1]].queryIdx].pt dist1 = np.linalg.norm(np.array(pt1) - np.array(pt1_2)) if dist1 > PIXEL_DIST_MULTIPLIER2 * dist2[i+1]: return False return True PIXEL_DIST_MULTIPLIER1 = 1.8 def filter_nearest_dist_from1(keypoints1, keypoints2, tree, matches, sift_match, num_neighbors=1): idx1 = sift_match.queryIdx idx2 = sift_match.trainIdx pt1 = keypoints1[idx1].pt pt2 = keypoints2[idx2].pt dist1, ii = tree.query(pt1, k=num_neighbors + 1) for i in range(num_neighbors): if dist1[i+1] == np.inf: continue pt2_2 = keypoints2[matches[ii[i+1]].trainIdx].pt dist2 = np.linalg.norm(np.array(pt2) - np.array(pt2_2)) if dist2 > PIXEL_DIST_MULTIPLIER1 * dist1[i+1]: return False return True import math LOCATION_EPS = 1 * 2**0.5 def remove_duplicate_match(keypoints1, keypoints2, matches): matches_filtered = [] def match_in_filtered(match): match_point1 = keypoints1[match.queryIdx].pt match_point2 = keypoints2[match.trainIdx].pt for filtered_match in matches_filtered: cur_point1 = keypoints1[filtered_match.queryIdx].pt cur_point2 = keypoints2[filtered_match.trainIdx].pt dist1 = np.linalg.norm(np.array(match_point1) - np.array(cur_point1)) dist2 = np.linalg.norm(np.array(match_point2) - np.array(cur_point2)) if dist1 == 0.0 or dist2 == 0.0: return True if dist1 < LOCATION_EPS and dist2 < LOCATION_EPS: return True return False for match in matches: if not match_in_filtered(match): matches_filtered.append(match) return matches_filtered # SLOPE_STD_MULTIPLIER = 2.58 # SLOPE_STD_MULTIPLIER = 1.96 SLOPE_STD_MULTIPLIER = 1.65 def remove_slope_outliers(keypoints1, keypoints2, matches): slopes = [] for match in matches: idx1 = match.queryIdx idx2 = match.trainIdx pt1 = keypoints1[idx1].pt pt2 = keypoints2[idx2].pt slope = (pt2[1] - pt1[1]) / (pt2[0] - pt1[0]) slopes.append(slope) slope_mean = np.mean(slopes) slope_std = np.std(slopes) def check_slope(slope): return abs(slope - slope_mean) < SLOPE_STD_MULTIPLIER * slope_std matches_filtered = [matches[idx] for idx, slope in enumerate(slopes) if check_slope(slope)] return matches_filtered SLOPE_ANGLE_THRES = np.pi / 12 def remove_angle_outliers(keypoints1, keypoints2, matches): if len(matches) < 2: return matches, None slopes = [] for match in matches: idx1 = match.queryIdx idx2 = match.trainIdx pt1 = keypoints1[idx1].pt pt2 = keypoints2[idx2].pt sl = - (pt2[1] - pt1[1]) / (pt2[0] - pt1[0] + 384) slopes.append(sl) def check_slope(idx, slope): slopes_tmp = slopes.copy() slopes_tmp.pop(idx) slope_mean_tmp = np.mean(slopes_tmp) base_angle_tmp = math.atan(slope_mean_tmp) return abs(base_angle_tmp - math.atan(slope)) < SLOPE_ANGLE_THRES matches_filtered = [matches[idx] for idx, slope in enumerate(slopes) if check_slope(idx, slope)] return matches_filtered, None class HomographyRANSAC: def __init__(self, threshold=2.0, confidence=0.99, max_iterations=1000): """ RANSAC specifically optimized for homography transformations Args: threshold: Reprojection error threshold in pixels confidence: Desired confidence level (0.99 = 99%) max_iterations: Maximum number of iterations """ self.threshold = threshold self.confidence = confidence self.max_iterations = max_iterations self.min_samples = 4 # Homography needs minimum 4 point correspondences def estimate_homography_transform(self, src_pts, dst_pts): """ Estimate homography transformation from point correspondences Homography transformation: [x'] = H [x] [y'] [y] [w'] [1] where x'/w', y'/w' are the final 2D coordinates Args: src_pts: Source points (Nx2) dst_pts: Destination points (Nx2) Returns: 3x3 homography matrix or None if failed """ if src_pts.shape[0] < 4: return None n_points = src_pts.shape[0] # Build the coefficient matrix A # For each point correspondence, we get 2 equations A = [] for i in range(n_points): x, y = src_pts[i] x_prime, y_prime = dst_pts[i] # First equation: x' = (h11*x + h12*y + h13) / (h31*x + h32*y + h33) # Rearranged: h11*x + h12*y + h13 - h31*x*x' - h32*y*x' - h33*x' = 0 A.append([-x, -y, -1, 0, 0, 0, x*x_prime, y*x_prime, x_prime]) # Second equation: y' = (h21*x + h22*y + h23) / (h31*x + h32*y + h33) # Rearranged: h21*x + h22*y + h23 - h31*x*y' - h32*y*y' - h33*y' = 0 A.append([0, 0, 0, -x, -y, -1, x*y_prime, y*y_prime, y_prime]) A = np.array(A) try: # Solve using SVD (homogeneous least squares) # The solution is the last column of V (corresponding to smallest singular value) U, S, Vt = np.linalg.svd(A) H = Vt[-1].reshape(3, 3) # Normalize so that H[2,2] = 1 H = H / H[2, 2] return H except np.linalg.LinAlgError: return None def apply_homography_transform(self, homography_matrix, points): """ Apply homography transformation to points Args: homography_matrix: 3x3 homography matrix points: Nx2 array of points Returns: Transformed points (Nx2) """ # Convert to homogeneous coordinates homogeneous_pts = np.hstack([points, np.ones((points.shape[0], 1))]) # Apply transformation transformed_homogeneous = (homography_matrix @ homogeneous_pts.T).T # Convert back from homogeneous coordinates (divide by w) # Handle potential division by zero w = transformed_homogeneous[:, 2] w = np.where(np.abs(w) < 1e-8, 1e-8, w) # Avoid division by zero transformed = transformed_homogeneous[:, :2] / w[:, np.newaxis] return transformed def compute_reprojection_error(self, homography_matrix, src_pts, dst_pts): """ Compute reprojection errors for homography transformation Args: homography_matrix: 3x3 homography matrix src_pts: Source points dst_pts: Destination points Returns: Array of reprojection errors (distances) """ # Transform source points transformed_pts = self.apply_homography_transform(homography_matrix, src_pts) # Compute Euclidean distances errors = np.linalg.norm(transformed_pts - dst_pts, axis=1) return errors def adaptive_iterations(self, n_points, current_best_inliers=0): """ Compute adaptive number of iterations based on current best result """ if n_points < self.min_samples: return 0 # Estimate outlier ratio if current_best_inliers > 0: inlier_ratio = current_best_inliers / n_points outlier_ratio = 1 - inlier_ratio else: outlier_ratio = 0.5 # Conservative initial estimate outlier_ratio = max(0.01, min(0.99, outlier_ratio)) # Clamp to reasonable range # Probability that a random sample of 4 points are all inliers prob_all_inliers = (1 - outlier_ratio) ** self.min_samples if prob_all_inliers > 1e-10: # Number of iterations to have confidence% chance of success num_iter = int(np.log(1 - self.confidence) / np.log(1 - prob_all_inliers)) return min(num_iter, self.max_iterations) else: return self.max_iterations def eliminate_outliers(self, src_pts, dst_pts, verbose=False): """ Main RANSAC function for eliminating outliers using homography model Args: src_pts: Source points (Nx2) dst_pts: Destination points (Nx2) verbose: Print progress information Returns: inlier_mask: Boolean mask for inlier correspondences inlier_indices: Indices of inlier correspondences best_homography: Best homography transformation found stats: Dictionary with statistics """ n_points = src_pts.shape[0] if n_points < self.min_samples: empty_mask = np.zeros(n_points, dtype=bool) return empty_mask, [], None, {"iterations": 0, "inliers": 0} best_inliers = [] best_homography = None max_inliers = 0 iteration = 0 if verbose: print(f"Starting RANSAC with {n_points} correspondences") print(f"Threshold: {self.threshold} pixels") while iteration < self.max_iterations: # Adaptive termination if iteration > 100 and iteration % 100 == 0: required_iterations = self.adaptive_iterations(n_points, max_inliers) if iteration >= required_iterations: if verbose: print(f"Early termination at iteration {iteration}") break # Randomly sample 4 correspondences try: sample_indices = random.sample(range(n_points), self.min_samples) except ValueError: break sample_src = src_pts[sample_indices] sample_dst = dst_pts[sample_indices] # Estimate homography transformation homography_matrix = self.estimate_homography_transform(sample_src, sample_dst) if homography_matrix is None: iteration += 1 continue # Compute errors for all points try: errors = self.compute_reprojection_error(homography_matrix, src_pts, dst_pts) except: iteration += 1 continue # Find inliers inlier_indices = np.where(errors < self.threshold)[0] n_inliers = len(inlier_indices) # Update best model if current is better if n_inliers > max_inliers: max_inliers = n_inliers best_inliers = inlier_indices best_homography = homography_matrix.copy() if verbose and iteration % 500 == 0: print(f"Iteration {iteration}: Found {n_inliers} inliers ({n_inliers/n_points*100:.1f}%)") iteration += 1 # Create inlier mask inlier_mask = np.zeros(n_points, dtype=bool) if len(best_inliers) > 0: inlier_mask[best_inliers] = True # Refine the transformation using all inliers if len(best_inliers) >= 4: refined_homography = self.estimate_homography_transform( src_pts[best_inliers], dst_pts[best_inliers] ) if refined_homography is not None: best_homography = refined_homography stats = { "iterations": iteration, "inliers": len(best_inliers), "inlier_ratio": len(best_inliers) / n_points if n_points > 0 else 0, "outlier_elimination_rate": 1 - (len(best_inliers) / n_points) if n_points > 0 else 0 } if verbose: print(f"Final results after {iteration} iterations:") print(f" Inliers: {len(best_inliers)}/{n_points} ({stats['inlier_ratio']*100:.1f}%)") print(f" Outliers eliminated: {stats['outlier_elimination_rate']*100:.1f}%") return inlier_mask, best_inliers, best_homography, stats def filter_opencv_matches(self, matches, kp1, kp2, verbose=False): """ Filter OpenCV matches using homography RANSAC Args: matches: List of cv2.DMatch objects kp1: Keypoints from first image kp2: Keypoints from second image verbose: Print statistics Returns: good_matches: Filtered matches homography_matrix: Estimated homography transformation stats: Statistics dictionary """ if len(matches) < self.min_samples: return [], None, {"error": "Not enough matches"} # Extract point coordinates src_pts = np.float32([kp1[m.queryIdx].pt for m in matches]) dst_pts = np.float32([kp2[m.trainIdx].pt for m in matches]) # Apply RANSAC inlier_mask, inlier_indices, homography_matrix, stats = self.eliminate_outliers( src_pts, dst_pts, verbose=verbose ) # Filter matches based on inliers good_matches = [matches[i] for i in inlier_indices] return good_matches, homography_matrix, stats def filter_homography_ransac(src_pts, dst_pts, threshold, confidence=0.99, max_iterations=1000): """ Convenience function to filter correspondences using homography RANSAC Args: src_pts: Source points (Nx2) dst_pts: Destination points (Nx2) threshold: Reprojection error threshold Returns: detected_inliers: Set of inlier indices """ # Initialize homography RANSAC homography_ransac = HomographyRANSAC( threshold=threshold, confidence=confidence, max_iterations=max_iterations ) # Eliminate outliers inlier_mask, inlier_indices, estimated_homography, stats = homography_ransac.eliminate_outliers( src_pts, dst_pts, verbose=False ) # Evaluate results detected_inliers = set(inlier_indices) return detected_inliers RANSAC_THRESHOLD = 12.0 def sift_match(img1, img2, lowe_ratio=0.75, cross_check=True, verbose=False, legacy_filtering=False): """ Match two images using SIFT and return the matched keypoints. :param img1: First image (numpy array). :param img2: Second image (numpy array). :param ratio: Lowe's ratio test threshold. :param cross_check: Whether to use cross-checking. not used if ratio < 1 :param legacy_filtering: Reproduce the pre-fix behavior where filter_scale was silently skipped (filter_angle applied twice). For comparison only. :return: list of match objects, and two list of keypoints. """ # Initialize SIFT detector sift = cv2.SIFT_create() # Find the keypoints and descriptors with SIFT keypoints1, descriptors1 = sift.detectAndCompute(img1, None) keypoints2, descriptors2 = sift.detectAndCompute(img2, None) if len(keypoints1) < 2 or len(keypoints2) < 2: if verbose: print("Not enough keypoints found in one of the images.(< 2)") return [], keypoints1, keypoints2 if lowe_ratio < 1: # filter matches using lowe's ratio test bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False) matches = bf.knnMatch(descriptors1, descriptors2, k=2) good_matches = [] for m, n in matches: if m.distance < lowe_ratio * n.distance: good_matches.append(m) matches = good_matches if verbose: print(f"Found {len(matches)} matches after Lowe's ratio test") else: # Create a Brute Force Matcher object bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=cross_check) # Match descriptors matches = bf.match(descriptors1, descriptors2) if verbose: print(f"Found {len(matches)} matches after brute-force matching with cross_check={cross_check}") matches = sorted(matches, key=lambda x: x.distance) # filtering matches filtering_functions_1st = [filter_angle, filter_scale] for filter_func in filtering_functions_1st: applied_func = filter_angle if legacy_filtering else filter_func matches = list(filter(partial(applied_func, keypoints1, keypoints2), matches)) if verbose: print(f"Found {len(matches)} matches after {filter_func.__name__}") # filtering hue, saturation, value matches = list(filter(partial(filter_hsv, keypoints1, keypoints2, img1, img2), matches)) if verbose: print(f"Found {len(matches)} matches after filter_hsv") matches = remove_duplicate_match(keypoints1, keypoints2, matches) # detecting outliers if len(matches) >= 15: src_pts = np.float32([keypoints1[m.queryIdx].pt for m in matches]) dst_pts = np.float32([keypoints2[m.trainIdx].pt for m in matches]) detected_inliers = filter_homography_ransac(src_pts, dst_pts, threshold=RANSAC_THRESHOLD) matches = [matches[i] for i in detected_inliers] if verbose: print(f"Found {len(matches)} matches after filter_homography_ransac") elif len(matches) > 1: matches, _ = remove_angle_outliers(keypoints1, keypoints2, matches) if verbose: print(f"Found {len(matches)} matches after remove_angle_outliers") filtering_functions_2nd = [filter_nearest_dist_from2, filter_nearest_dist_from1] for filter_func in filtering_functions_2nd: for i in range(NUM_FILTER_NEAREST_ITTER): if len(matches) < 2: break if filter_func == filter_nearest_dist_from2: locations = np.array([keypoints2[m.trainIdx].pt for m in matches]) else: locations = np.array([keypoints1[m.queryIdx].pt for m in matches]) tree = KDTree(locations, copy_data=True) matches = list(filter(partial(filter_func, keypoints1, keypoints2, tree, matches.copy()), matches)) if verbose: print(f"Found {len(matches)} matches after {filter_func.__name__}, itter {i}") if len(matches) < 2: if verbose: print("Not enough matches after filtering.(< 2)") return [], keypoints1, keypoints2 return matches, keypoints1, keypoints2 ### Ported from diffusers script 2025.10.28 Takemoto### def compute_snr(timesteps, alphas_cumprod): """ Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 """ # Get alphas and sigmas. alpha = get_alphas(timesteps, alphas_cumprod) sigma = get_sigmas(timesteps, alphas_cumprod) # Compute SNR. snr = (alpha / sigma) ** 2 return snr def get_alphas(timesteps, alphas_cumprod): sqrt_alphas_cumprod = alphas_cumprod ** 0.5 # Expand the tensors. # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026 sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float() while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape): sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None] alpha = sqrt_alphas_cumprod.expand(timesteps.shape) return alpha def get_sigmas(timesteps, alphas_cumprod): sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float() while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape): sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None] sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape) return sigma