Slerp vectors

slerp_diffusion.py

# Install the required libs
# pip install -qq git+https://github.com/huggingface/diffusers.git accelerate tensorboard transformers ftfy gradio
# pip list | grep transformers
# pip install transformers==4.22.2
# pip install -qq "ipywidgets>=7,<8"
# pip install -qq bitsandbytes


import argparse
import itertools
import math
import os
from contextlib import nullcontext
import random

import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch.utils.data import Dataset

import PIL
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDPMScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel
from diffusers.hub_utils import init_git_repo, push_to_hub
from diffusers.optimization import get_scheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm
from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
from tqdm.auto import tqdm
import bitsandbytes as bnb
from torch import autocast
from Ipython.display import display
from pathlib import Path
  
def slerp(t, v0, v1, DOT_THRESHOLD=0.9995):
    """ helper function to spherically interpolate two arrays v1 v2 """

    if not isinstance(v0, np.ndarray):
        inputs_are_torch = True
        input_device = v0.device
        v0 = v0.cpu().numpy()
        v1 = v1.cpu().numpy()

    dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
    if np.abs(dot) > DOT_THRESHOLD:
        v2 = (1 - t) * v0 + t * v1
    else:
        theta_0 = np.arccos(dot)
        sin_theta_0 = np.sin(theta_0)
        theta_t = theta_0 * t
        sin_theta_t = np.sin(theta_t)
        s0 = np.sin(theta_0 - theta_t) / sin_theta_0
        s1 = sin_theta_t / sin_theta_0
        v2 = s0 * v0 + s1 * v1

    if inputs_are_torch:
        v2 = torch.from_numpy(v2).to(input_device)

    return v2

def get_random_latent(n=1):
    generator = torch.Generator(device="cuda")
    seed = generator.seed()
    seeds.append(seed)
    generator = generator.manual_seed(seed)
    
    image_latents = torch.randn(
        (n, pipe.unet.in_channels, height // 8, width // 8),
        generator = generator,
        device = "cuda"
    )
    return image_latents
  
 
secret_value='' # insert your Hugging face token here
pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4', use_auth_token=secret_value, torch_dtype=torch.float16).to("cuda")
pipe.safety_checker = lambda images, **kwargs: (images, False)


prompt="Your prompt"

first_latent=get_random_latent()
second_latent=get_random_latent()

Path('slerp_between_latents').mkdir(exist_ok=True)
for i, t in enumerate(tqdm(np.linspace(0, 1, num_steps))):
    inter_latent = slerp(t, first_latent, second_latent)

    with autocast("cuda"):
        images = pipe([prompt], num_inference_steps=50, guidance_scale=7.5,  latents = inter_latent).images
    images[0].save(f'slerp_between_latents/{i}.png')