PyTorch implementation of spherical linear interpolation

slerp.py

from torch import FloatTensor, LongTensor, Tensor, Size, lerp, zeros_like
from torch.linalg import norm

# adapted to PyTorch from:
# https://gist.github.com/dvschultz/3af50c40df002da3b751efab1daddf2c
# most of the extra complexity is to support:
# - many-dimensional vectors
# - v0 or v1 with last dim all zeroes, or v0 ~colinear with v1
#   - falls back to lerp()
#   - conditional logic implemented with parallelism rather than Python loops
# - many-dimensional tensor for t
#   - you can ask for batches of slerp outputs by making t more-dimensional than the vectors
#   -   slerp(
#         v0:   torch.Size([2,3]),
#         v1:   torch.Size([2,3]),
#         t:  torch.Size([4,1,1]), 
#       )
#   - this makes it interface-compatible with lerp()
def slerp(v0: FloatTensor, v1: FloatTensor, t: float|FloatTensor, DOT_THRESHOLD=0.9995):
  '''
  Spherical linear interpolation
  Args:
    v0: Starting vector
    v1: Final vector
    t: Float value between 0.0 and 1.0
    DOT_THRESHOLD: Threshold for considering the two vectors as
                            colinear. Not recommended to alter this.
  Returns:
      Interpolation vector between v0 and v1
  '''
  assert v0.shape == v1.shape, "shapes of v0 and v1 must match"

  # Normalize the vectors to get the directions and angles
  v0_norm: FloatTensor = norm(v0, dim=-1)
  v1_norm: FloatTensor = norm(v1, dim=-1)

  v0_normed: FloatTensor = v0 / v0_norm.unsqueeze(-1)
  v1_normed: FloatTensor = v1 / v1_norm.unsqueeze(-1)

  # Dot product with the normalized vectors
  dot: FloatTensor = (v0_normed * v1_normed).sum(-1)
  dot_mag: FloatTensor = dot.abs()

  # if dp is NaN, it's because the v0 or v1 row was filled with 0s
  # If absolute value of dot product is almost 1, vectors are ~colinear, so use lerp
  gotta_lerp: LongTensor = dot_mag.isnan() | (dot_mag > DOT_THRESHOLD)
  can_slerp: LongTensor = ~gotta_lerp

  t_batch_dim_count: int = max(0, t.dim()-v0.dim()) if isinstance(t, Tensor) else 0
  t_batch_dims: Size = t.shape[:t_batch_dim_count] if isinstance(t, Tensor) else Size([])
  out: FloatTensor = zeros_like(v0.expand(*t_batch_dims, *[-1]*v0.dim()))

  # if no elements are lerpable, our vectors become 0-dimensional, preventing broadcasting
  if gotta_lerp.any():
    lerped: FloatTensor = lerp(v0, v1, t)

    out: FloatTensor = lerped.where(gotta_lerp.unsqueeze(-1), out)

  # if no elements are slerpable, our vectors become 0-dimensional, preventing broadcasting
  if can_slerp.any():

    # Calculate initial angle between v0 and v1
    theta_0: FloatTensor = dot.arccos().unsqueeze(-1)
    sin_theta_0: FloatTensor = theta_0.sin()
    # Angle at timestep t
    theta_t: FloatTensor = theta_0 * t
    sin_theta_t: FloatTensor = theta_t.sin()
    # Finish the slerp algorithm
    s0: FloatTensor = (theta_0 - theta_t).sin() / sin_theta_0
    s1: FloatTensor = sin_theta_t / sin_theta_0
    slerped: FloatTensor = s0 * v0 + s1 * v1

    out: FloatTensor = slerped.where(can_slerp.unsqueeze(-1), out)
  
  return out

Example invocation:

from torch import FloatTensor
import torch

device=torch.device('cuda')
dtype=torch.float16

# spherical linear midpoint between [1,0] and [0,1] is [sine(pi/4), sine(pi/4)]
start: FloatTensor = tensor([1,0], dtype=dtype, device=device)
end: FloatTensor = tensor([0,1], dtype=dtype, device=device)
slerp(start, end, 0.5)
# tensor([0.7070, 0.7070], device='cuda:0', dtype=torch.float16)

# many-dimensional data supported; interpolation is performed row-wise
start: FloatTensor = tensor([[1,0], [2,0]], dtype=dtype, device=device)
end: FloatTensor = tensor([[0,1], [0,2]], dtype=dtype, device=device)
slerp(start, end, 0.5)
# tensor([[0.7070, 0.7070],
#         [1.4141, 1.4141]], device='cuda:0', dtype=torch.float16)

# any row where either vector has all zeroes, is computed via fallback to lerp():
start: FloatTensor = tensor([[0,0], [1,0]], dtype=dtype, device=device)
end: FloatTensor = tensor([[1,1], [0,1]], dtype=dtype, device=device)
slerp(start, end, time)
# tensor([[0.5000, 0.5000],
#         [0.7070, 0.7070]], device='cuda:0', dtype=torch.float16)

# any row where both vectors are approx. colinear, is computed via fallback to lerp():
start: FloatTensor = tensor([1,1], dtype=dtype, device=device)
end: FloatTensor = tensor([2,2], dtype=dtype, device=device)
slerp(start, end, 0.5)
# tensor([1.5000, 1.5000], device='cuda:0', dtype=torch.float16)

# we can compute slerp at a batch of timepoints simultaneously,
# by providing a time tensor that is more-dimensional than start/end
start: FloatTensor = tensor([1,0], dtype=dtype, device=device)
end: FloatTensor = tensor([0,1], dtype=dtype, device=device)
time: FloatTensor = tensor([[0.25], [0.5]], dtype=torch.float16, device=device)
slerp(start, end, time)
# tensor([[0.9238, 0.3826],
#         [0.7070, 0.7070]], device='cuda:0', dtype=torch.float16)

hi @Birch-san ,

thanks for the snippet.

I found it when searching for ways to make my idea work. Also kudos for your technical blog, I see you have quite some hands-on experience with SD topics.

I'd like to implement two ideas, using FLUX's VAE :

• encode an image, sample from the distributions vector, decode using the VAE and hopefully get a variation of the initial image (one that is different to a human eye while retaining the same quality)
• encode two distinct images, compute an affine combination of their latents (or slerp: that's where your snippet came in), decode and hopefully get an in-between image. (btw I found no visible difference in output image between using linear or spherical interpolation)

These would be cheap ways, resource-wise, to generate images as they enable to skip the reverse diffusion.

None of these two worked well enough for me using the FLUX (dev or schnell) VAE:

• sampling yielded an identical image, because very close in the latent space I guess. Further perturbing the sample with more stds only brought artefacts (stds are very small relative to the means)
• While I did manage to get in-between images (see below), the result has artefacts, even when decoded with the VAE (thus to the full resolution)

Is this related to the fact that the KL loss weight of this VAE is relatively low, giving priority to reconstruction and therefore the interpolation might hit a no man's land in the latent space?
I'd like to understand more whether this is achievable or not and if so how.

A big thanks for your insight,

Clement

############################# results and snippet ######################################

below is my attempt at interpolating between a portrait without smile and one with smile: I do get a half smile, but when zooming in artefacts show up.

The artefacts on the interpolated image, looking like blur (here at the decoded resolution ie no upscale)

My code to get the interpolation above, using the diffusers API:
(code for first idea -sampling- is essentially doing latent_dist.sample() )

(Note that in my actual use case, the latents will come from within the forward pass ie from the call to the pipeline, as opposed to artifically obtained by encoding an image)

import torch
from diffusers import AutoencoderKL
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import pil_to_tensor


smiling = <path_img1>
not_smiling = <path_img2>


vae: AutoencoderKL = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae"
)
with torch.no_grad():
    # vae:AutoencoderKL = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix")

    vae.to(dtype=torch.float32)  # otherwise it produces NaNs, even madebyollin's VAE
    vae.to(device="cuda")

    assert vae.device == torch.device("cuda:0")
    assert vae.dtype == torch.float32

    ## 1) smile image
    # make image as tensor
    img = Image.open(smiling).resize((1024, 1024)).convert("RGB")
    img_tensor = pil_to_tensor(img).unsqueeze(0) / 255.0
    img_tensor = img_tensor.to(vae.device)
    img_tensor = img_tensor.to(vae.dtype)

    # get the inferred latent distribution
    latent_dist: DiagonalGaussianDistribution = vae.encode(
        img_tensor, return_dict=False
    )[0]
    print(
        f"{latent_dist.mean.shape=} {latent_dist.std.shape=} {latent_dist.mean.mean()=} {latent_dist.std.mean()=}"
    )

    assert not latent_dist.mean.isnan().any()
    assert not latent_dist.std.isnan().any()
    assert latent_dist.deterministic is False

    latent_smiling = (
        latent_dist.mode()
    )  # .squeeze(0)  # sample from the latent distribution

    ## 2) no smile image
    # make image as tensor
    img = Image.open(not_smiling).resize((1024, 1024)).convert("RGB")
    img_tensor = pil_to_tensor(img).detach().cpu().unsqueeze(0) / 255.0
    img_tensor = img_tensor.to(vae.device)
    img_tensor = img_tensor.to(vae.dtype)

    # get the inferred latent distribution
    latent_dist: DiagonalGaussianDistribution = vae.encode(
        img_tensor, return_dict=False
    )[0]
    print(
        f"{latent_dist.mean.shape=} {latent_dist.std.shape=} {latent_dist.mean.mean()=} {latent_dist.std.mean()=}"
    )

    assert not latent_dist.mean.isnan().any()
    assert not latent_dist.std.isnan().any()
    assert latent_dist.deterministic is False

    latent_not_smiling = (
        latent_dist.mode()
    )  # .squeeze(0)  # sample from the latent distribution

    assert (latent_not_smiling != latent_smiling).any()

    # avg then decode
    t = 0.5
    avg_latent = (
        t * latent_smiling + (1 - t) * latent_not_smiling
    )  # linear interpolation
    # avg_latent = slerp(latent_smiling, latent_not_smiling, t)
    print(
        f"{avg_latent.shape=} {avg_latent.mean()=} {avg_latent.std()=} {avg_latent.min()=} {avg_latent.max()=}"
    )
    image = vae.decode(avg_latent).sample.squeeze(0).cpu().detach()
    image = transforms.ToPILImage()(image)
    image.save("interpolated_.png")