Code example - Interpolation between prompts for SDXL

[asomoza]

Hi, this is a code example for doing interpolation between prompts with SDXL since the online examples are only for SD 1.5.

I did this because of #7568 and I always had in mind to understand and test how to do this.

Basically what I did is use the pipeline functions instead of using the models (tokenizers and text encoders) to skip doing it for each one and also I used a lighting lora instead of the turbo model because I like it more and to show that you can use it with any model you want.

Both of these models don't need to use the classifier free guidance so I didn't use the interpolation on the negative prompts but if you want to use them it should be easy to add.

Here's the result between these prompts:

prompt one = "A cinematic shot of a baby racoon wearing an intricate italian priest robe."
prompt_two = "A cinematic shot of a baby rabbit wearing an tuxedo."

Had to resize the original image because it was too big.

Here's the code:

import numpy as np
import torch
from tqdm import tqdm

from diffusers import AutoencoderKL, EulerDiscreteScheduler, StableDiffusionXLPipeline
from diffusers.utils.torch_utils import randn_tensor

# based on https://huggingface.co/learn/cookbook/stable_diffusion_interpolation#example-3-interpolation-between-multiple-prompts 
# and https://gist.github.com/karpathy/00103b0037c5aaea32fe1da1af553355
def slerp(v0, v1, num, t0=0, t1=1):
    # convert them to full precision if not the dot product overflows
    v0 = v0.detach().cpu().numpy().astype(np.float32)
    v1 = v1.detach().cpu().numpy().astype(np.float32)

    def interpolation(t, v0, v1, DOT_THRESHOLD=0.9995):
        """helper function to spherically interpolate two arrays v1 v2"""
        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
        return v2

    t = np.linspace(t0, t1, num)

    v3 = torch.tensor(np.array([interpolation(t[i], v0, v1) for i in range(num)]), dtype=torch.float16)

    return v3


vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "RunDiffusion/Juggernaut-XL-v9", torch_dtype=torch.float16, variant="fp16", vae=vae
).to("cuda")
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config, timestep_spacing="trailing")
pipeline.load_lora_weights(
    "ByteDance/SDXL-Lightning", weight_name="sdxl_lightning_4step_lora.safetensors", adapter_name="lighting"
)
pipeline.set_progress_bar_config(disable=True)

height = 1024
width = 1024
num_interpolation_steps = 30
seed = 2139965163
generator = torch.Generator(device="cpu").manual_seed(seed)

# add all the prompts you want here
prompts = [
    "A cinematic shot of a baby racoon wearing an intricate italian priest robe.",
    "A cinematic shot of a baby rabbit wearing an tuxedo.",
]
batch_size = len(prompts)

prompts_embeds, _negative_prompts_embeds, pooled_prompts_embeds, _negative_pooled_prompts_embeds = (
    pipeline.encode_prompt(prompts, do_classifier_free_guidance=False)
)

shape = (1, pipeline.unet.config.in_channels, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor)
latents = randn_tensor(shape, generator=generator, device=prompts_embeds.device, dtype=prompts_embeds.dtype)

interpolated_prompt_embeds = []
interpolated_pooled_prompt_embeds = []

for i in range(batch_size - 1):
    interpolated_prompt_embeds.append(slerp(prompts_embeds[i], prompts_embeds[i + 1], num_interpolation_steps))
    interpolated_pooled_prompt_embeds.append(
        slerp(pooled_prompts_embeds[i], pooled_prompts_embeds[i + 1], num_interpolation_steps)
    )

interpolated_prompt_embeds = torch.cat(interpolated_prompt_embeds, dim=0).to(prompts_embeds.device)
interpolated_pooled_prompt_embeds = torch.cat(interpolated_pooled_prompt_embeds, dim=0).to(prompts_embeds.device)

images = []

for prompt_embeds, pooled_prompt_embeds in tqdm(
    zip(
        interpolated_prompt_embeds,
        interpolated_pooled_prompt_embeds,
    ),
    total=len(interpolated_prompt_embeds),
):
    images.append(
        pipeline(
            height=height,
            width=width,
            latents=latents,
            prompt_embeds=prompt_embeds[None, ...],
            pooled_prompt_embeds=pooled_prompt_embeds[None, ...],
            guidance_scale=1.0,
            num_inference_steps=4,
            generator=generator,
        ).images[0]
    )

images[0].save(
    "animation.gif",
    save_all=True,
    append_images=images[1:],
    optimize=False,
    duration=num_interpolation_steps * 3,
    loop=0,
)

[sayakpaul]

Another one: https://github.com/huggingface/cookbook/blob/main/notebooks/en/stable_diffusion_interpolation.ipynb