deforum-stable-diffusion/helpers/conditioning.py

import torch
import torch.nn as nn
from torch.nn import functional as F
import clip
from torchvision.transforms import Normalize as Normalize
from torchvision.utils import make_grid
import numpy as np
from IPython import display
from sklearn.cluster import KMeans
import torchvision.transforms.functional as TF

###
# Loss functions
###


## CLIP -----------------------------------------

class MakeCutouts(nn.Module):
 def __init__(self, cut_size, cutn, cut_pow=1.):
 super().__init__()
 self.cut_size = cut_size
 self.cutn = cutn
 self.cut_pow = cut_pow

 def forward(self, input):
 sideY, sideX = input.shape[2:4]
 max_size = min(sideX, sideY)
 min_size = min(sideX, sideY, self.cut_size)
 cutouts = []
 for _ in range(self.cutn):
 size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)
 offsetx = torch.randint(0, sideX - size + 1, ())
 offsety = torch.randint(0, sideY - size + 1, ())
 cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]
 cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))
 return torch.cat(cutouts)


def spherical_dist_loss(x, y):
 x = F.normalize(x, dim=-1)
 y = F.normalize(y, dim=-1)
 return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)

def make_clip_loss_fn(root, args):
 clip_size = root.clip_model.visual.input_resolution # for openslip: clip_model.visual.image_size

 def parse_prompt(prompt):
 if prompt.startswith('http://') or prompt.startswith('https://'):
 vals = prompt.rsplit(':', 2)
 vals = [vals[0] + ':' + vals[1], *vals[2:]]
 else:
 vals = prompt.rsplit(':', 1)
 vals = vals + ['', '1'][len(vals):]
 return vals[0], float(vals[1])

 def parse_clip_prompts(clip_prompt):
 target_embeds, weights = [], []
 for prompt in clip_prompt:
 txt, weight = parse_prompt(prompt)
 target_embeds.append(root.clip_model.encode_text(clip.tokenize(txt).to(root.device)).float())
 weights.append(weight)
 target_embeds = torch.cat(target_embeds)
 weights = torch.tensor(weights, device=root.device)
 if weights.sum().abs() < 1e-3:
 raise RuntimeError('Clip prompt weights must not sum to 0.')
 weights /= weights.sum().abs()
 return target_embeds, weights

 normalize = Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
 std=[0.26862954, 0.26130258, 0.27577711])

 make_cutouts = MakeCutouts(clip_size, args.cutn, args.cut_pow)
 target_embeds, weights = parse_clip_prompts(args.clip_prompt)

 def clip_loss_fn(x, sigma, **kwargs):
 nonlocal target_embeds, weights, make_cutouts, normalize
 clip_in = normalize(make_cutouts(x.add(1).div(2)))
 image_embeds = root.clip_model.encode_image(clip_in).float()
 dists = spherical_dist_loss(image_embeds[:, None], target_embeds[None])
 dists = dists.view([args.cutn, 1, -1])
 losses = dists.mul(weights).sum(2).mean(0)
 return losses.sum()

 return clip_loss_fn

def make_aesthetics_loss_fn(root,args):
 clip_size = root.clip_model.visual.input_resolution # for openslip: clip_model.visual.image_size

 def aesthetics_cond_fn(x, sigma, **kwargs):
 clip_in = F.interpolate(x, (clip_size, clip_size))
 image_embeds = root.clip_model.encode_image(clip_in).float()
 losses = (10 - root.aesthetics_model(image_embeds)[0])
 return losses.sum()

 return aesthetics_cond_fn

## end CLIP -----------------------------------------

# blue loss from @johnowhitaker's tutorial on Grokking Stable Diffusion
def blue_loss_fn(x, sigma, **kwargs):
 # How far are the blue channel values to 0.9:
 error = torch.abs(x[:,-1, :, :] - 0.9).mean() 
 return error

# MSE loss from init
def make_mse_loss(target):
 def mse_loss(x, sigma, **kwargs):
 return (x - target).square().mean()
 return mse_loss

# MSE loss from init
def exposure_loss(target):
 def exposure_loss_fn(x, sigma, **kwargs):
 error = torch.abs(x-target).mean()
 return error
 return exposure_loss_fn

def mean_loss_fn(x, sigma, **kwargs):
 error = torch.abs(x).mean() 
 return error

def var_loss_fn(x, sigma, **kwargs):
 error = x.var()
 return error

def get_color_palette(root, n_colors, target, verbose=False):
 def display_color_palette(color_list):
 # Expand to 64x64 grid of single color pixels
 images = color_list.unsqueeze(2).repeat(1,1,64).unsqueeze(3).repeat(1,1,1,64)
 images = images.double().cpu().add(1).div(2).clamp(0, 1)
 images = torch.tensor(np.array(images))
 grid = make_grid(images, 8).cpu()
 display.display(TF.to_pil_image(grid))
 return

 # Create color palette
 kmeans = KMeans(n_clusters=n_colors, random_state=0).fit(torch.flatten(target[0],1,2).T.cpu().numpy())
 color_list = torch.Tensor(kmeans.cluster_centers_).to(root.device)
 if verbose:
 display_color_palette(color_list)
 # Get ratio of each color class in the target image
 color_indexes, color_counts = np.unique(kmeans.labels_, return_counts=True)
 # color_list = color_list[color_indexes]
 return color_list, color_counts

def make_rgb_color_match_loss(root, target, n_colors, ignore_sat_weight=None, img_shape=None, device='cuda:0'):
 """
 target (tensor): Image sample (values from -1 to 1) to extract the color palette
 n_colors (int): Number of colors in the color palette
 ignore_sat_weight (None or number>0): Scale to ignore color saturation in color comparison
 img_shape (None or (int, int)): shape (width, height) of sample that the conditioning gradient is applied to, 
 if None then calculate target color distribution during gradient calculation 
 rather than once at the beginning
 """
 assert n_colors > 0, "Must use at least one color with color match loss"

 def adjust_saturation(sample, saturation_factor):
 # as in torchvision.transforms.functional.adjust_saturation, but for tensors with values from -1,1
 return blend(sample, TF.rgb_to_grayscale(sample), saturation_factor)

 def blend(img1, img2, ratio):
 return (ratio * img1 + (1.0 - ratio) * img2).clamp(-1, 1).to(img1.dtype)

 def color_distance_distributions(n_colors, img_shape, color_list, color_counts, n_images=1):
 # Get the target color distance distributions
 # Ensure color counts total the amout of pixels in the image
 n_pixels = img_shape[0]*img_shape[1]
 color_counts = (color_counts * n_pixels / sum(color_counts)).astype(int)

 # Make color distances for each color, sorted by distance
 color_distributions = torch.zeros((n_colors, n_images, n_pixels), device=device)
 for i_image in range(n_images):
 for ic,color0 in enumerate(color_list):
 i_dist = 0
 for jc,color1 in enumerate(color_list):
 color_dist = torch.linalg.norm(color0 - color1)
 color_distributions[ic, i_image, i_dist:i_dist+color_counts[jc]] = color_dist
 i_dist += color_counts[jc]
 color_distributions, _ = torch.sort(color_distributions,dim=2)
 return color_distributions

 color_list, color_counts = get_color_palette(root, n_colors, target)
 color_distributions = None
 if img_shape is not None:
 color_distributions = color_distance_distributions(n_colors, img_shape, color_list, color_counts)

 def rgb_color_ratio_loss(x, sigma, **kwargs):
 nonlocal color_distributions
 all_color_norm_distances = torch.ones(len(color_list), x.shape[0], x.shape[2], x.shape[3]).to(device) * 6.0 # distance to color won't be more than max norm1 distance between -1 and 1 in 3 color dimensions

 for ic,color in enumerate(color_list):
 # Make a tensor of entirely one color
 color = color[None,:,None].repeat(1,1,x.shape[2]).unsqueeze(3).repeat(1,1,1,x.shape[3])
 # Get the color distances
 if ignore_sat_weight is None:
 # Simple color distance
 color_distances = torch.linalg.norm(x - color, dim=1)
 else:
 # Color distance if the colors were saturated
 # This is to make color comparison ignore shadows and highlights, for example
 color_distances = torch.linalg.norm(adjust_saturation(x, ignore_sat_weight) - color, dim=1)

 all_color_norm_distances[ic] = color_distances
 all_color_norm_distances = torch.flatten(all_color_norm_distances,start_dim=2)

 if color_distributions is None:
 color_distributions = color_distance_distributions(n_colors, 
 (x.shape[2], x.shape[3]), 
 color_list, 
 color_counts, 
 n_images=x.shape[0])

 # Sort the color distances so we can compare them as if they were a cumulative distribution function
 all_color_norm_distances, _ = torch.sort(all_color_norm_distances,dim=2)

 color_norm_distribution_diff = all_color_norm_distances - color_distributions

 return color_norm_distribution_diff.square().mean()

 return rgb_color_ratio_loss


###
# Thresholding functions for grad
###
def threshold_by(threshold, threshold_type, clamp_schedule):

 def dynamic_thresholding(vals, sigma):
 # Dynamic thresholding from Imagen paper (May 2022)
 s = np.percentile(np.abs(vals.cpu()), threshold, axis=tuple(range(1,vals.ndim)))
 s = np.max(np.append(s,1.0))
 vals = torch.clamp(vals, -1*s, s)
 vals = torch.FloatTensor.div(vals, s)
 return vals

 def static_thresholding(vals, sigma):
 vals = torch.clamp(vals, -1*threshold, threshold)
 return vals

 def mean_thresholding(vals, sigma): # Thresholding that appears in Jax and Disco
 magnitude = vals.square().mean(axis=(1,2,3),keepdims=True).sqrt()
 vals = vals * torch.where(magnitude > threshold, threshold / magnitude, 1.0)
 return vals

 def scheduling(vals, sigma):
 clamp_val = clamp_schedule[sigma.item()]
 magnitude = vals.square().mean().sqrt()
 vals = vals * magnitude.clamp(max=clamp_val) / magnitude
 #print(clamp_val)
 return vals

 if threshold_type == 'dynamic':
 return dynamic_thresholding
 elif threshold_type == 'static':
 return static_thresholding
 elif threshold_type == 'mean':
 return mean_thresholding
 elif threshold_type == 'schedule':
 return scheduling
 else:
 raise Exception(f"Thresholding type {threshold_type} not supported")