preprocess_utils.py

from transformers import CLIPTextModel, CLIPTokenizer, logging
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
# suppress partial model loading warning
logging.set_verbosity_error()

import os
from tqdm import tqdm, trange
import torch
import torch.nn as nn
import argparse
from torchvision.io import write_video
from pathlib import Path
from util import *
import torchvision.transforms as T


def get_timesteps(scheduler, num_inference_steps, strength, device):
    # get the original timestep using init_timestep
    init_timestep = min(int(num_inference_steps * strength), num_inference_steps)

    t_start = max(num_inference_steps - init_timestep, 0)
    timesteps = scheduler.timesteps[t_start:]

    return timesteps, num_inference_steps - t_start
    
@torch.no_grad()
def decode_latents(pipe, latents):
    decoded = []
    batch_size = 8
    for b in range(0, latents.shape[0], batch_size):
            latents_batch = 1 / 0.18215 * latents[b:b + batch_size]
            imgs = pipe.vae.decode(latents_batch).sample
            imgs = (imgs / 2 + 0.5).clamp(0, 1)
            decoded.append(imgs)
    return torch.cat(decoded)

@torch.no_grad()
def ddim_inversion(pipe, cond, latent_frames,  batch_size, save_latents=True, timesteps_to_save=None):
    
    timesteps = reversed(pipe.scheduler.timesteps)
    timesteps_to_save = timesteps_to_save if timesteps_to_save is not None else timesteps
    for i, t in enumerate(tqdm(timesteps)):
        for b in range(0, latent_frames.shape[0], batch_size):
            x_batch = latent_frames[b:b + batch_size]
            model_input = x_batch
            cond_batch = cond.repeat(x_batch.shape[0], 1, 1)
            #remove comment from commented block to support controlnet
            # if self.sd_version == 'depth':
            #     depth_maps = torch.cat([self.depth_maps[b: b + batch_size]])
            #     model_input = torch.cat([x_batch, depth_maps],dim=1)

            alpha_prod_t = pipe.scheduler.alphas_cumprod[t]
            alpha_prod_t_prev = (
                pipe.scheduler.alphas_cumprod[timesteps[i - 1]]
                if i > 0 else pipe.scheduler.final_alpha_cumprod
            )

            mu = alpha_prod_t ** 0.5
            mu_prev = alpha_prod_t_prev ** 0.5
            sigma = (1 - alpha_prod_t) ** 0.5
            sigma_prev = (1 - alpha_prod_t_prev) ** 0.5

            
            #remove line below and replace with commented block to support controlnet
            eps = pipe.unet(model_input, t, encoder_hidden_states=cond_batch).sample
            # if self.sd_version != 'ControlNet':
            #     eps = pipe.unet(model_input, t, encoder_hidden_states=cond_batch).sample
            # else:
            #     eps = self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]]))
            
            pred_x0 = (x_batch - sigma_prev * eps) / mu_prev
            latent_frames[b:b + batch_size] = mu * pred_x0 + sigma * eps

    #     if save_latents and t in timesteps_to_save:
    #         torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt'))
    # torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt'))
    return latent_frames    
    
@torch.no_grad()
def ddim_sample(pipe, x, cond, batch_size):
    timesteps = pipe.scheduler.timesteps
    for i, t in enumerate(tqdm(timesteps)):
        for b in range(0, x.shape[0], batch_size):
            x_batch = x[b:b + batch_size]
            model_input = x_batch
            cond_batch = cond.repeat(x_batch.shape[0], 1, 1)
            
            #remove comment from commented block to support controlnet
            # if self.sd_version == 'depth':
            #     depth_maps = torch.cat([self.depth_maps[b: b + batch_size]])
            #     model_input = torch.cat([x_batch, depth_maps],dim=1)

            alpha_prod_t = pipe.scheduler.alphas_cumprod[t]
            alpha_prod_t_prev = (
                pipe.scheduler.alphas_cumprod[timesteps[i + 1]]
                if i < len(timesteps) - 1
                else pipe.scheduler.final_alpha_cumprod
            )
            mu = alpha_prod_t ** 0.5
            sigma = (1 - alpha_prod_t) ** 0.5
            mu_prev = alpha_prod_t_prev ** 0.5
            sigma_prev = (1 - alpha_prod_t_prev) ** 0.5

            #remove line below and replace with commented block to support controlnet
            eps = pipe.unet(model_input, t, encoder_hidden_states=cond_batch).sample
            # if self.sd_version != 'ControlNet':
            #     eps = pipe.unet(model_input, t, encoder_hidden_states=cond_batch).sample
            # else:
            #     eps = self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]]))

            pred_x0 = (x_batch - sigma * eps) / mu
            x[b:b + batch_size] = mu_prev * pred_x0 + sigma_prev * eps
    return x


@torch.no_grad()
def get_text_embeds(pipe, prompt, negative_prompt, batch_size=1, device="cuda"):
    # Tokenize text and get embeddings
    text_input = pipe.tokenizer(prompt, padding='max_length', max_length=pipe.tokenizer.model_max_length,
                                truncation=True, return_tensors='pt')
    text_embeddings = pipe.text_encoder(text_input.input_ids.to(pipe.device))[0]

    # Do the same for unconditional embeddings
    uncond_input = pipe.tokenizer(negative_prompt, padding='max_length', max_length=pipe.tokenizer.model_max_length,
                                  return_tensors='pt')

    uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(pipe.device))[0]

    # Cat for final embeddings
    text_embeddings = torch.cat([uncond_embeddings] * batch_size + [text_embeddings] * batch_size)
    return text_embeddings

@torch.no_grad()
def extract_latents(pipe,
                    num_steps,
                    latent_frames,
                    batch_size,
                    timesteps_to_save,
                    inversion_prompt=''):
    pipe.scheduler.set_timesteps(num_steps)
    cond = get_text_embeds(pipe, inversion_prompt, "", device=pipe.device)[1].unsqueeze(0)
    # latent_frames = self.latents

    inverted_latents = ddim_inversion(pipe, cond,
                                latent_frames,
                                batch_size=batch_size,
                                save_latents=False,
                                timesteps_to_save=timesteps_to_save)
    
    # latent_reconstruction = ddim_sample(pipe, inverted_latents, cond, batch_size=batch_size)

#     rgb_reconstruction = decode_latents(pipe, latent_reconstruction)

#     return rgb_reconstruction
    return inverted_latents
    
@torch.no_grad()
def encode_imgs(pipe, imgs, batch_size=10, deterministic=True):
    imgs = 2 * imgs - 1
    latents = []
    for i in range(0, len(imgs), batch_size):
        posterior = pipe.vae.encode(imgs[i:i + batch_size]).latent_dist
        latent = posterior.mean if deterministic else posterior.sample()
        latents.append(latent * 0.18215)
    latents = torch.cat(latents)
    return latents
    
def get_data(pipe, frames, n_frames):
    """
    converts frames to tensors, saves to device and encodes to obtain latents
    """
    frames = frames[:n_frames]
    if frames[0].size[0] == frames[0].size[1]:
        frames = [frame.convert("RGB").resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
    stacked_tensor_frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(pipe.device)
    # encode to latents
    latents = encode_imgs(pipe, stacked_tensor_frames, deterministic=True).to(torch.float16).to(pipe.device)
    return stacked_tensor_frames, latents