from abc import ABC, abstractmethod
from contextlib import nullcontext
import time
import os
from functools import partial
from copy import deepcopy
from multiprocessing import Pool
from threading import Lock
from PIL import Image
import numpy as np
import torch
import torch.nn.functional as F
import einops
from transformers import LlamaForCausalLM
import spaces
from vqvae_muse import VQGANModel, get_tokenizer_muse
from torch_vqvae_model import get_tokenizer
def get_torch_float_dtype(dtype):
if dtype in (torch.float16, torch.bfloat16, torch.float32):
return dtype
return {
'float16': torch.float16,
'fp16': torch.float16,
'f16': torch.float16,
'bfloat16': torch.bfloat16,
'bf16': torch.bfloat16,
'float32': torch.float32,
'fp32': torch.float32,
'f32': torch.float32,
}[dtype]
def get_pid():
time.sleep(1)
return os.getpid()
class InferenceModel(ABC):
@abstractmethod
def __call__(input_images, n_new_frames, n_candidates, temperature=1.0, top_p=1.0):
raise NotImplementedError()
class LocalInferenceModel(InferenceModel):
def __init__(self, checkpoint, dtype='float16', torch_device='cuda',
context_frames=16, use_lock=False):
self.checkpoint = checkpoint
self.dtype = dtype
self.torch_device = torch_device
self.context_frames = context_frames
# new tokenizer
self.tokenizer = get_tokenizer_muse()
self.tokenizer.to(self.torch_device)
self.model = LlamaForCausalLM.from_pretrained(
self.checkpoint, torch_dtype=get_torch_float_dtype(self.dtype)
).to(self.torch_device)
print("torch device", self.torch_device)
print("init device", self.model.device)
if use_lock:
self.lock = Lock()
else:
self.lock = nullcontext()
@torch.no_grad()
def compute_perplexity(self, input_images, target_images):
input_images = np.array(input_images)
target_images = np.array(target_images)
assert len(input_images.shape) == 5 and len(target_images.shape) == 5 # [B, S, H, W, C]
assert input_images.shape[0] == target_images.shape[0]
batch_size = input_images.shape[0]
with self.lock:
input_images = torch.tensor(
einops.rearrange(input_images, 'b s h w c -> b s c h w')
).to(self.torch_device)
target_images = torch.tensor(
einops.rearrange(target_images, 'b s h w c -> b s c h w')
).to(self.torch_device)
input_ids = self.tokenizer.tokenize(input_images).view(batch_size, -1)
target_ids = self.tokenizer.tokenize(target_images).view(batch_size, -1)
all_ids = torch.cat([input_ids, target_ids], dim=1)
logits = self.model(all_ids).logits
log_probs = F.log_softmax(logits, dim=-1)
target_ids_onehot = F.one_hot(target_ids, num_classes=logits.shape[-1])
target_log_probs = log_probs[:, input_ids.shape[1] - 1 : -1]
perplexity = torch.exp(
-torch.mean(
torch.sum(target_log_probs * target_ids_onehot, dim=-1),
dim=-1
)
)
return perplexity.detach().cpu().numpy()
@torch.no_grad()
def generate_once(self, input_images, n_new_frames, temperature=1.0, top_p=1.0):
assert type(input_images) == np.ndarray
with self.lock:
input_images = np.array(input_images, dtype=np.float32)
input_images = torch.tensor(
einops.rearrange(input_images, 'b h w c -> b c h w')
).to(self.torch_device)
# not quite sure why i need to redo it here
self.model.to(self.torch_device)
self.tokenizer.to(self.torch_device)
# new tokenizer
_, input_ids = self.tokenizer.encode(input_images)
input_ids = input_ids.view(1, -1)
input_ids = input_ids[:, -(self.context_frames - 1) * 256:]
new_tokens = []
current_context_frames = input_ids.shape[1] // 256
fisrt_generation_left = self.context_frames - current_context_frames
first_new_frames = min(fisrt_generation_left, n_new_frames)
input_ids = self.model.generate(
input_ids=input_ids,
attention_mask=torch.ones_like(input_ids),
pad_token_id=8192,
max_new_tokens=256 * first_new_frames,
do_sample=True,
top_p=top_p,
temperature=temperature,
suppress_tokens=list(range(8192, self.model.vocab_size)),
)
new_tokens.append(input_ids[:, -256 * first_new_frames:])
input_ids = input_ids[:, -(self.context_frames - 1) * 256:]
for _ in range(max(0, n_new_frames - first_new_frames)):
input_ids = self.model.generate(
input_ids=input_ids,
attention_mask=torch.ones_like(input_ids),
pad_token_id=8192,
max_new_tokens=256,
do_sample=True,
top_p=top_p,
temperature=temperature,
suppress_tokens=list(range(8192, self.model.vocab_size)),
)
new_tokens.append(input_ids[:, -256:])
input_ids = input_ids[:, -(self.context_frames - 1) * 256:]
new_tokens = torch.cat(new_tokens, dim=1).view(-1, 256)
new_images = einops.rearrange(
torch.clamp(self.tokenizer.decode_code(new_tokens), 0.0, 1.0),
'b c h w -> b h w c'
).detach().cpu().numpy()
return new_images
@spaces.GPU(duration=180)
def __call__(self, input_images, n_new_frames, n_candidates, temperature=1.0, top_p=1.0):
output = []
for seq in input_images:
output.append(
[self.generate_once(seq, n_new_frames, temperature, top_p)
for _ in range(n_candidates)]
)
return output
class MultiProcessInferenceModel(InferenceModel):
def __init__(self, checkpoint, torch_devices=None, dtype='float16',
context_frames=16, use_lock=False, perplexity_batch_size=2):
if torch_devices is None or torch_devices == '':
torch_devices = [f'cuda:{i}' for i in range(torch.cuda.device_count())]
self.torch_devices = torch_devices
self.n_processes = len(torch_devices)
print(f'Using {self.n_processes} processes for inference')
self.worker_pool = Pool(self.n_processes)
self.worker_pids = self.worker_pool.starmap(get_pid, [tuple() for _ in range(self.n_processes)])
self.device_map = {
pid: torch_device
for pid, torch_device in zip(self.worker_pids, self.torch_devices)
}
self.worker_pool.starmap(
self.initialize_worker,
[(self.device_map, checkpoint, dtype, context_frames) for _ in range(self.n_processes)]
)
self.perplexity_batch_size = perplexity_batch_size
if use_lock:
self.lock = Lock()
else:
self.lock = nullcontext()
@staticmethod
def initialize_worker(device_map, checkpoint, dtype, context_frames):
global _current_process_backend
torch_device = device_map[os.getpid()]
_current_process_backend = LocalInferenceModel(
checkpoint, dtype, torch_device, context_frames
)
@staticmethod
def generate_once(input_images, n_new_frames, temperature=1.0, top_p=1.0):
return _current_process_backend.generate_once(input_images, n_new_frames, temperature, top_p)
@staticmethod
def compute_perplexity_once(input_images, target_images):
return _current_process_backend.compute_perplexity(input_images, target_images)
def compute_perplexity(self, input_images, target_images):
with self.lock:
map_args = []
for i in range(0, len(input_images), self.perplexity_batch_size):
map_args.append((
input_images[i : i + self.perplexity_batch_size],
target_images[i : i + self.perplexity_batch_size]
))
outputs = self.worker_pool.starmap(self.compute_perplexity_once, map_args)
return np.concatenate(outputs, axis=0)
def __call__(self, input_images, n_new_frames, n_candidates, temperature=1.0, top_p=1.0):
with self.lock:
map_args = []
for seq in input_images:
for _ in range(n_candidates):
map_args.append((seq, n_new_frames, temperature, top_p))
outputs = self.worker_pool.starmap(self.generate_once, map_args)
reshaped_output = []
index = 0
for _ in range(len(input_images)):
candidates = []
for _ in range(n_candidates):
candidates.append(outputs[index])
index += 1
reshaped_output.append(candidates)
return reshaped_output