import torch
from torch import nn
from einops import rearrange
import numpy as np
from typing import List
from models.id_embedding.helpers import get_rep_pos, shift_tensor_dim0
from models.id_embedding.meta_net import StyleVectorizer
from models.celeb_embeddings import _get_celeb_embeddings_basis
from functools import partial
import torch.nn.functional as F
import torch.nn as nn
import torch.nn.init as init
DEFAULT_PLACEHOLDER_TOKEN = ["*"]
PROGRESSIVE_SCALE = 2000
def get_clip_token_for_string(tokenizer, string):
batch_encoding = tokenizer(string, return_length=True, padding=True, truncation=True, return_overflowing_tokens=False, return_tensors="pt")
tokens = batch_encoding["input_ids"]
return tokens
def get_embedding_for_clip_token(embedder, token):
return embedder(token.unsqueeze(0))
class EmbeddingManagerId_adain(nn.Module):
def __init__(
self,
tokenizer,
text_encoder,
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"),
experiment_name = "normal_GAN",
num_embeds_per_token: int = 2,
loss_type: str = None,
mlp_depth: int = 2,
token_dim: int = 1024,
input_dim: int = 1024,
**kwargs
):
super().__init__()
self.device = device
self.num_es = num_embeds_per_token
self.get_token_for_string = partial(get_clip_token_for_string, tokenizer)
self.get_embedding_for_tkn = partial(get_embedding_for_clip_token, text_encoder.text_model.embeddings)
self.token_dim = token_dim
''' 1. Placeholder mapping dicts '''
self.placeholder_token = self.get_token_for_string("*")[0][1]
if experiment_name == "normal_GAN":
self.celeb_embeddings_mean, self.celeb_embeddings_std = _get_celeb_embeddings_basis(tokenizer, text_encoder, "datasets_face/good_names.txt")
elif experiment_name == "man_GAN":
self.celeb_embeddings_mean, self.celeb_embeddings_std = _get_celeb_embeddings_basis(tokenizer, text_encoder, "datasets_face/good_names_man.txt")
elif experiment_name == "woman_GAN":
self.celeb_embeddings_mean, self.celeb_embeddings_std = _get_celeb_embeddings_basis(tokenizer, text_encoder, "datasets_face/good_names_woman.txt")
else:
print("Hello, please notice this ^_^")
assert 0
print("now experiment_name:", experiment_name)
self.celeb_embeddings_mean = self.celeb_embeddings_mean.to(device)
self.celeb_embeddings_std = self.celeb_embeddings_std.to(device)
self.name_projection_layer = StyleVectorizer(input_dim, self.token_dim * self.num_es, depth=mlp_depth, lr_mul=0.1)
self.embedding_discriminator = Embedding_discriminator(self.token_dim * self.num_es, dropout_rate = 0.2)
self.adain_mode = 0
def forward(
self,
tokenized_text,
embedded_text,
name_batch,
random_embeddings = None,
timesteps = None,
):
if tokenized_text is not None:
batch_size, n, device = *tokenized_text.shape, tokenized_text.device
other_return_dict = {}
if random_embeddings is not None:
mlp_output_embedding = self.name_projection_layer(random_embeddings)
total_embedding = mlp_output_embedding.view(mlp_output_embedding.shape[0], 2, 1024)
if self.adain_mode == 0:
adained_total_embedding = total_embedding * self.celeb_embeddings_std + self.celeb_embeddings_mean
else:
adained_total_embedding = total_embedding
other_return_dict["total_embedding"] = total_embedding
other_return_dict["adained_total_embedding"] = adained_total_embedding
if name_batch is not None:
if isinstance(name_batch, list):
name_tokens = self.get_token_for_string(name_batch)[:, 1:3]
name_embeddings = self.get_embedding_for_tkn(name_tokens.to(random_embeddings.device))[0]
other_return_dict["name_embeddings"] = name_embeddings
else:
assert 0
if tokenized_text is not None:
placeholder_pos = get_rep_pos(tokenized_text,
[self.placeholder_token])
placeholder_pos = np.array(placeholder_pos)
if len(placeholder_pos) != 0:
batch_size = adained_total_embedding.shape[0]
end_index = min(batch_size, placeholder_pos.shape[0])
embedded_text[placeholder_pos[:, 0], placeholder_pos[:, 1]] = adained_total_embedding[:end_index,0,:]
embedded_text[placeholder_pos[:, 0], placeholder_pos[:, 1] + 1] = adained_total_embedding[:end_index,1,:]
return embedded_text, other_return_dict
def load(self, ckpt_path):
ckpt = torch.load(ckpt_path, map_location='cuda')
if ckpt.get("name_projection_layer") is not None:
self.name_projection_layer = ckpt.get("name_projection_layer").float()
print('[Embedding Manager] weights loaded.')
def save(self, ckpt_path):
save_dict = {}
save_dict["name_projection_layer"] = self.name_projection_layer
torch.save(save_dict, ckpt_path)
def trainable_projection_parameters(self):
trainable_list = []
trainable_list.extend(list(self.name_projection_layer.parameters()))
return trainable_list
class Embedding_discriminator(nn.Module):
def __init__(self, input_size, dropout_rate):
super(Embedding_discriminator, self).__init__()
self.input_size = input_size
self.fc1 = nn.Linear(input_size, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 1)
self.LayerNorm1 = nn.LayerNorm(512)
self.LayerNorm2 = nn.LayerNorm(256)
self.leaky_relu = nn.LeakyReLU(0.2)
self.dropout_rate = dropout_rate
if self.dropout_rate > 0:
self.dropout1 = nn.Dropout(dropout_rate)
self.dropout2 = nn.Dropout(dropout_rate)
def forward(self, input):
x = input.view(-1, self.input_size)
if self.dropout_rate > 0:
x = self.leaky_relu(self.dropout1(self.fc1(x)))
else:
x = self.leaky_relu(self.fc1(x))
if self.dropout_rate > 0:
x = self.leaky_relu(self.dropout2(self.fc2(x)))
else:
x = self.leaky_relu(self.fc2(x))
x = self.fc3(x)
return x
def save(self, ckpt_path):
save_dict = {}
save_dict["fc1"] = self.fc1
save_dict["fc2"] = self.fc2
save_dict["fc3"] = self.fc3
save_dict["LayerNorm1"] = self.LayerNorm1
save_dict["LayerNorm2"] = self.LayerNorm2
save_dict["leaky_relu"] = self.leaky_relu
save_dict["dropout1"] = self.dropout1
save_dict["dropout2"] = self.dropout2
torch.save(save_dict, ckpt_path)
def load(self, ckpt_path):
ckpt = torch.load(ckpt_path, map_location='cuda')
if ckpt.get("first_name_proj_layer") is not None:
self.fc1 = ckpt.get("fc1").float()
self.fc2 = ckpt.get("fc2").float()
self.fc3 = ckpt.get("fc3").float()
self.LayerNorm1 = ckpt.get("LayerNorm1").float()
self.LayerNorm2 = ckpt.get("LayerNorm2").float()
self.leaky_relu = ckpt.get("leaky_relu").float()
self.dropout1 = ckpt.get("dropout1").float()
self.dropout2 = ckpt.get("dropout2").float()
print('[Embedding D] weights loaded.')