mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-09 00:09:55 +00:00
Justin John
365 lines
16 KiB
Python
365 lines
16 KiB
Python
import glob
|
|
import os
|
|
import numpy as np
|
|
import cv2
|
|
from pathlib import Path
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision.transforms as T
|
|
import argparse
|
|
from PIL import Image
|
|
import yaml
|
|
from tqdm import tqdm
|
|
from transformers import logging
|
|
from diffusers import DDIMScheduler, StableDiffusionPipeline
|
|
|
|
from tokenflow_utils import *
|
|
from utils import save_video, seed_everything
|
|
|
|
# suppress partial model loading warning
|
|
logging.set_verbosity_error()
|
|
|
|
VAE_BATCH_SIZE = 10
|
|
|
|
class TokenFlow(nn.Module):
|
|
def __init__(self, config,
|
|
pipe,
|
|
frames=None,
|
|
# latents = None,
|
|
inverted_latents = None):
|
|
super().__init__()
|
|
self.config = config
|
|
self.device = config["device"]
|
|
self.to = torch.float16 if self.device == 'cuda' else torch.float32
|
|
|
|
sd_version = config["sd_version"]
|
|
self.sd_version = sd_version
|
|
if sd_version == '2.1':
|
|
model_key = "stabilityai/stable-diffusion-2-1-base"
|
|
elif sd_version == '2.0':
|
|
model_key = "stabilityai/stable-diffusion-2-base"
|
|
elif sd_version == '1.5':
|
|
model_key = "runwayml/stable-diffusion-v1-5"
|
|
elif sd_version == 'depth':
|
|
model_key = "stabilityai/stable-diffusion-2-depth"
|
|
else:
|
|
raise ValueError(f'Stable-diffusion version {sd_version} not supported.')
|
|
|
|
# Create SD models
|
|
print('Loading SD model')
|
|
|
|
# pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")
|
|
# pipe.enable_xformers_memory_efficient_attention()
|
|
|
|
self.vae = pipe.vae
|
|
self.tokenizer = pipe.tokenizer
|
|
self.text_encoder = pipe.text_encoder
|
|
self.unet = pipe.unet
|
|
|
|
self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
|
|
self.scheduler.set_timesteps(config["n_timesteps"], device=self.device)
|
|
print('SD model loaded')
|
|
|
|
# data
|
|
self.frames, self.inverted_latents = frames, inverted_latents
|
|
self.latents_path = self.get_latents_path()
|
|
|
|
# load frames
|
|
self.paths, self.frames, self.latents, self.eps = self.get_data()
|
|
|
|
if self.sd_version == 'depth':
|
|
self.depth_maps = self.prepare_depth_maps()
|
|
|
|
self.text_embeds = self.get_text_embeds(config["prompt"], config["negative_prompt"])
|
|
# pnp_inversion_prompt = self.get_pnp_inversion_prompt()
|
|
self.pnp_guidance_embeds = self.get_text_embeds(config["pnp_inversion_prompt"], config["pnp_inversion_prompt"]).chunk(2)[0]
|
|
|
|
@torch.no_grad()
|
|
def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'):
|
|
depth_maps = []
|
|
midas = torch.hub.load("intel-isl/MiDaS", model_type)
|
|
midas.to(self.device)
|
|
midas.eval()
|
|
|
|
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
|
|
|
|
if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
|
|
transform = midas_transforms.dpt_transform
|
|
else:
|
|
transform = midas_transforms.small_transform
|
|
|
|
for i in range(len(self.paths)):
|
|
img = cv2.imread(self.paths[i])
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
latent_h = img.shape[0] // 8
|
|
latent_w = img.shape[1] // 8
|
|
|
|
input_batch = transform(img).to(self.device)
|
|
prediction = midas(input_batch)
|
|
|
|
depth_map = torch.nn.functional.interpolate(
|
|
prediction.unsqueeze(1),
|
|
size=(latent_h, latent_w),
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
)
|
|
depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
|
|
depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
|
|
depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
|
|
depth_maps.append(depth_map)
|
|
|
|
return torch.cat(depth_maps).to(to).to(self.device)
|
|
|
|
def get_pnp_inversion_prompt(self):
|
|
inv_prompts_path = os.path.join(str(Path(self.latents_path).parent), 'inversion_prompt.txt')
|
|
# read inversion prompt
|
|
with open(inv_prompts_path, 'r') as f:
|
|
inv_prompt = f.read()
|
|
return inv_prompt
|
|
|
|
def get_latents_path(self):
|
|
read_from_files = self.frames is None
|
|
# read_from_files = True
|
|
if read_from_files:
|
|
latents_path = os.path.join(self.config["latents_path"], f'sd_{self.config["sd_version"]}',
|
|
Path(self.config["data_path"]).stem, f'steps_{self.config["n_inversion_steps"]}')
|
|
latents_path = [x for x in glob.glob(f'{latents_path}/*') if '.' not in Path(x).name]
|
|
n_frames = [int([x for x in latents_path[i].split('/') if 'nframes' in x][0].split('_')[1]) for i in range(len(latents_path))]
|
|
print("n_frames", n_frames)
|
|
latents_path = latents_path[np.argmax(n_frames)]
|
|
print("latents_path", latents_path)
|
|
self.config["n_frames"] = min(max(n_frames), self.config["n_frames"])
|
|
|
|
else:
|
|
n_frames = self.frames.shape[0]
|
|
self.config["n_frames"] = min(n_frames, self.config["n_frames"])
|
|
|
|
if self.config["n_frames"] % self.config["batch_size"] != 0:
|
|
# make n_frames divisible by batch_size
|
|
self.config["n_frames"] = self.config["n_frames"] - (self.config["n_frames"] % self.config["batch_size"])
|
|
print("Number of frames: ", self.config["n_frames"])
|
|
if read_from_files:
|
|
print("YOOOOOOO", os.path.join(latents_path, 'latents'))
|
|
return os.path.join(latents_path, 'latents')
|
|
else:
|
|
return None
|
|
|
|
@torch.no_grad()
|
|
def get_text_embeds(self, prompt, negative_prompt, batch_size=1):
|
|
# Tokenize text and get embeddings
|
|
text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
|
|
truncation=True, return_tensors='pt')
|
|
text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
|
|
|
|
# Do the same for unconditional embeddings
|
|
uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length,
|
|
return_tensors='pt')
|
|
|
|
uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.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 encode_imgs(self, imgs, batch_size=VAE_BATCH_SIZE, deterministic=False):
|
|
imgs = 2 * imgs - 1
|
|
latents = []
|
|
for i in range(0, len(imgs), batch_size):
|
|
posterior = self.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
|
|
|
|
@torch.no_grad()
|
|
def decode_latents(self, latents, batch_size=VAE_BATCH_SIZE):
|
|
latents = 1 / 0.18215 * latents
|
|
imgs = []
|
|
for i in range(0, len(latents), batch_size):
|
|
imgs.append(self.vae.decode(latents[i:i + batch_size]).sample)
|
|
imgs = torch.cat(imgs)
|
|
imgs = (imgs / 2 + 0.5).clamp(0, 1)
|
|
return imgs
|
|
|
|
|
|
def get_data(self):
|
|
read_from_files = self.frames is None
|
|
# read_from_files = True
|
|
if read_from_files:
|
|
# load frames
|
|
paths = [os.path.join(self.config["data_path"], "%05d.jpg" % idx) for idx in
|
|
range(self.config["n_frames"])]
|
|
if not os.path.exists(paths[0]):
|
|
paths = [os.path.join(self.config["data_path"], "%05d.png" % idx) for idx in
|
|
range(self.config["n_frames"])]
|
|
frames = [Image.open(paths[idx]).convert('RGB') for idx in range(self.config["n_frames"])]
|
|
if frames[0].size[0] == frames[0].size[1]:
|
|
frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames]
|
|
frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(self.to).to(self.device)
|
|
save_video(frames, f'{self.config["output_path"]}/input_fps10.mp4', fps=10)
|
|
save_video(frames, f'{self.config["output_path"]}/input_fps20.mp4', fps=20)
|
|
save_video(frames, f'{self.config["output_path"]}/input_fps30.mp4', fps=30)
|
|
else:
|
|
frames = self.frames
|
|
# encode to latents
|
|
latents = self.encode_imgs(frames, deterministic=True).to(self.to).to(self.device)
|
|
# get noise
|
|
eps = self.get_ddim_eps(latents, range(self.config["n_frames"])).to(self.to).to(self.device)
|
|
if not read_from_files:
|
|
return None, frames, latents, eps
|
|
return paths, frames, latents, eps
|
|
|
|
def get_ddim_eps(self, latent, indices):
|
|
read_from_files = self.inverted_latents is None
|
|
# read_from_files = True
|
|
if read_from_files:
|
|
noisest = max([int(x.split('_')[-1].split('.')[0]) for x in glob.glob(os.path.join(self.latents_path, f'noisy_latents_*.pt'))])
|
|
print("noisets:", noisest)
|
|
print("indecies:", indices)
|
|
latents_path = os.path.join(self.latents_path, f'noisy_latents_{noisest}.pt')
|
|
noisy_latent = torch.load(latents_path)[indices].to(self.device)
|
|
|
|
# path = os.path.join('test_latents', f'noisy_latents_{noisest}.pt')
|
|
# f_noisy_latent = torch.load(path)[indices].to(self.device)
|
|
# print(f_noisy_latent==noisy_latent)
|
|
else:
|
|
noisest = max([int(key.split("_")[-1]) for key in self.inverted_latents.keys()])
|
|
print("noisets:", noisest)
|
|
print("indecies:", indices)
|
|
noisy_latent = self.inverted_latents[f'noisy_latents_{noisest}'][indices]
|
|
|
|
alpha_prod_T = self.scheduler.alphas_cumprod[noisest]
|
|
mu_T, sigma_T = alpha_prod_T ** 0.5, (1 - alpha_prod_T) ** 0.5
|
|
eps = (noisy_latent - mu_T * latent) / sigma_T
|
|
return eps
|
|
|
|
@torch.no_grad()
|
|
def denoise_step(self, x, t, indices):
|
|
# register the time step and features in pnp injection modules
|
|
read_files = self.inverted_latents is None
|
|
|
|
if read_files:
|
|
source_latents = load_source_latents_t(t, self.latents_path)[indices]
|
|
|
|
else:
|
|
source_latents = self.inverted_latents[f'noisy_latents_{t}'][indices]
|
|
|
|
latent_model_input = torch.cat([source_latents] + ([x] * 2))
|
|
if self.sd_version == 'depth':
|
|
latent_model_input = torch.cat([latent_model_input, torch.cat([self.depth_maps[indices]] * 3)], dim=1)
|
|
|
|
register_time(self, t.item())
|
|
|
|
# compute text embeddings
|
|
text_embed_input = torch.cat([self.pnp_guidance_embeds.repeat(len(indices), 1, 1),
|
|
torch.repeat_interleave(self.text_embeds, len(indices), dim=0)])
|
|
|
|
# apply the denoising network
|
|
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input)['sample']
|
|
|
|
# perform guidance
|
|
_, noise_pred_uncond, noise_pred_cond = noise_pred.chunk(3)
|
|
noise_pred = noise_pred_uncond + self.config["guidance_scale"] * (noise_pred_cond - noise_pred_uncond)
|
|
|
|
# compute the denoising step with the reference model
|
|
denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
|
|
return denoised_latent
|
|
|
|
@torch.autocast(dtype=torch.float16, device_type='cuda')
|
|
def batched_denoise_step(self, x, t, indices):
|
|
batch_size = self.config["batch_size"]
|
|
denoised_latents = []
|
|
pivotal_idx = torch.randint(batch_size, (len(x)//batch_size,)) + torch.arange(0,len(x),batch_size)
|
|
|
|
register_pivotal(self, True)
|
|
self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx])
|
|
register_pivotal(self, False)
|
|
for i, b in enumerate(range(0, len(x), batch_size)):
|
|
register_batch_idx(self, i)
|
|
denoised_latents.append(self.denoise_step(x[b:b + batch_size], t, indices[b:b + batch_size]))
|
|
denoised_latents = torch.cat(denoised_latents)
|
|
return denoised_latents
|
|
|
|
def init_method(self, conv_injection_t, qk_injection_t):
|
|
self.qk_injection_timesteps = self.scheduler.timesteps[:qk_injection_t] if qk_injection_t >= 0 else []
|
|
self.conv_injection_timesteps = self.scheduler.timesteps[:conv_injection_t] if conv_injection_t >= 0 else []
|
|
register_extended_attention_pnp(self, self.qk_injection_timesteps)
|
|
register_conv_injection(self, self.conv_injection_timesteps)
|
|
set_tokenflow(self.unet)
|
|
|
|
def save_vae_recon(self):
|
|
os.makedirs(f'{self.config["output_path"]}/vae_recon', exist_ok=True)
|
|
decoded = self.decode_latents(self.latents)
|
|
for i in range(len(decoded)):
|
|
T.ToPILImage()(decoded[i]).save(f'{self.config["output_path"]}/vae_recon/%05d.png' % i)
|
|
save_video(decoded, f'{self.config["output_path"]}/vae_recon_10.mp4', fps=10)
|
|
save_video(decoded, f'{self.config["output_path"]}/vae_recon_20.mp4', fps=20)
|
|
save_video(decoded, f'{self.config["output_path"]}/vae_recon_30.mp4', fps=30)
|
|
|
|
def edit_video(self):
|
|
save_files = self.inverted_latents is None # if we're in the original non-demo setting
|
|
if save_files:
|
|
os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
|
|
self.save_vae_recon()
|
|
# self.save_vae_recon()
|
|
pnp_f_t = int(self.config["n_timesteps"] * self.config["pnp_f_t"])
|
|
pnp_attn_t = int(self.config["n_timesteps"] * self.config["pnp_attn_t"])
|
|
|
|
self.init_method(conv_injection_t=pnp_f_t, qk_injection_t=pnp_attn_t)
|
|
|
|
noisy_latents = self.scheduler.add_noise(self.latents, self.eps, self.scheduler.timesteps[0])
|
|
edited_frames = self.sample_loop(noisy_latents, torch.arange(self.config["n_frames"]))
|
|
|
|
if save_files:
|
|
save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_10.mp4')
|
|
save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_20.mp4', fps=20)
|
|
save_video(edited_frames, f'{self.config["output_path"]}/tokenflow_PnP_fps_30.mp4', fps=30)
|
|
print('Done!')
|
|
else:
|
|
return edited_frames
|
|
|
|
def sample_loop(self, x, indices):
|
|
save_files = self.inverted_latents is None # if we're in the original non-demo setting
|
|
# save_files = True
|
|
if save_files:
|
|
os.makedirs(f'{self.config["output_path"]}/img_ode', exist_ok=True)
|
|
for i, t in enumerate(tqdm(self.scheduler.timesteps, desc="Sampling")):
|
|
x = self.batched_denoise_step(x, t, indices)
|
|
|
|
decoded_latents = self.decode_latents(x)
|
|
if save_files:
|
|
for i in range(len(decoded_latents)):
|
|
T.ToPILImage()(decoded_latents[i]).save(f'{self.config["output_path"]}/img_ode/%05d.png' % i)
|
|
|
|
return decoded_latents
|
|
|
|
|
|
# def run(config):
|
|
# seed_everything(config["seed"])
|
|
# print(config)
|
|
# editor = TokenFlow(config)
|
|
# editor.edit_video()
|
|
|
|
|
|
# if __name__ == '__main__':
|
|
# parser = argparse.ArgumentParser()
|
|
# parser.add_argument('--config_path', type=str, default='configs/config_pnp.yaml')
|
|
# opt = parser.parse_args()
|
|
# with open(opt.config_path, "r") as f:
|
|
# config = yaml.safe_load(f)
|
|
# config["output_path"] = os.path.join(config["output_path"] + f'_pnp_SD_{config["sd_version"]}',
|
|
# Path(config["data_path"]).stem,
|
|
# config["prompt"][:240],
|
|
# f'attn_{config["pnp_attn_t"]}_f_{config["pnp_f_t"]}',
|
|
# f'batch_size_{str(config["batch_size"])}',
|
|
# str(config["n_timesteps"]),
|
|
# )
|
|
# os.makedirs(config["output_path"], exist_ok=True)
|
|
# print(config["data_path"])
|
|
# assert os.path.exists(config["data_path"]), "Data path does not exist"
|
|
# with open(os.path.join(config["output_path"], "config.yaml"), "w") as f:
|
|
# yaml.dump(config, f)
|
|
# run(config)
|