I am trying to reproduce this project: https://github.com/feizc/Gradient-Free-Textual-Inversion，But I now have a problem:
Loading pipeline components...: 100%|█████████████| 6/6 [00:14<00:00, 2.43s/it]
convert text inversion: <sks> in id: 49408
load data length: 50
Traceback (most recent call last):
File "/root/autodl-tmp/train_inversion.py", line 298, in <module>
main()
File "/root/autodl-tmp/train_inversion.py", line 291, in main
fitnesses = pipeline.eval(solutions)
File "/root/autodl-tmp/train_inversion.py", line 108, in eval
z = z + self.init_text_inversion
RuntimeError: The size of tensor a (1024) must match the size of tensor b (512) at non-singleton dimension 1
And I found that the shape of z changes when I use different models. For example, when I use 512-base-ema.ckpt, it is [1024], and when I use sd-v1-4.ckpt, it becomes [768]. self_init_text_inversion remains unchanged.
In addition, as far as I know, the project’s initialize_inversion.py uses clip to initialize text and training data. Here I use openai/clip-vit-base-patch16.
1.The following is the initialization code:
(https://github.com/feizc/Gradient-Free-Textual-Inversion/blob/main/initialize_inversion.py)
"""
automatically initialize the textual inversion with CLIP and no-parameter cross-attention
使用 CLIP 和无参数交叉注意力自动初始化文本反转
"""
import torch
import os
import argparse
from PIL import Image
import torch.nn.functional as F
from transformers import CLIPModel, CLIPTokenizer, CLIPProcessor, CLIPTextModel
from utils import imagenet_template, automatic_subjective_classnames
def embedding_generate(model, tokenizer, text_encoder, classnames, templates, device):
"""
pre-caculate the template sentence, token embeddings
预计算模板句子、令牌嵌入
"""
with torch.no_grad():
sentence_weights = []
token_weights = []
token_embedding_table = text_encoder.get_input_embeddings().weight.data
for classname in classnames:
texts = [template(classname) for template in templates] # format with class
texts = tokenizer(texts, padding="max_length", max_length=77, truncation=True, return_tensors="pt") # tokenize
texts = texts['input_ids'].to(device)
class_embeddings = model.get_text_features(texts)
class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
class_embedding /= class_embedding.norm()
sentence_weights.append(class_embedding)
token_ids = tokenizer.encode(classname,add_special_tokens=False)
token_embedding_list = []
for token_id in token_ids:
token_embedding_list.append(token_embedding_table[token_id])
token_weights.append(torch.mean(torch.stack(token_embedding_list), dim=0))
sentence_weights = torch.stack(sentence_weights, dim=1).to(device)
token_weights = torch.stack(token_weights, dim=0).to(device)
return sentence_weights, token_weights
def image_condition_embed_initialize(image_feature_list, sentence_embeddings, token_embeddings):
"""
no-parameter cross-attention: query: image, key: sentence, value: token
无参数交叉注意力：查询：图像，键：句子，值：令牌
"""
inversion_emb_list = []
for image_features in image_feature_list:
cross_attention = image_features @ sentence_embeddings
attention_probs = F.softmax(cross_attention, dim=-1)
inversion_emb = torch.matmul(attention_probs, token_embeddings)
inversion_emb_list.append(inversion_emb)
final_inversion = torch.mean(torch.stack(inversion_emb_list), dim=0)
final_inversion = final_inversion / final_inversion.norm()
return final_inversion
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--save_path", default='./save', type=str)
parser.add_argument("--data_path", default='./coconut_seed_fruit_stage', type=str)
args = parser.parse_args()
save_path = args.save_path
data_path = args.data_path
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
tokenizer = CLIPTokenizer.from_pretrained('/root/.cache/huggingface/patch16b')
model = CLIPModel.from_pretrained('/root/.cache/huggingface/patch16b').to(device)
text_encoder = CLIPTextModel.from_pretrained('/root/.cache/huggingface/patch16b')
processor = CLIPProcessor.from_pretrained('/root/.cache/huggingface/patch16b')
sentence_embeddings, token_embeddings = embedding_generate(model,
tokenizer,
text_encoder,
automatic_subjective_classnames,
imagenet_template,
device)
print('sentence embedding size: ', sentence_embeddings.size(), ' token embedding size: ', token_embeddings.size())
image_feature_list = []
name_list = os.listdir(data_path)
for name in name_list:
image_path = os.path.join(data_path, name)
image = Image.open(image_path)
inputs = processor(images=image, return_tensors="pt").to(device)
image_features = model.get_image_features(**inputs)
image_features = F.normalize(image_features, dim=-1)
image_feature_list.append(image_features)
print('image size: ', len(image_feature_list))
inversion_emb = image_condition_embed_initialize(image_feature_list, sentence_embeddings, token_embeddings)
inversion_emb_dict = {"initialize": inversion_emb.detach().cpu()}
torch.save(inversion_emb_dict, os.path.join(save_path, 'initialize_emb.bin'))
if __name__ == "__main__":
main()
2.The following is the training code:
(https://github.com/feizc/Gradient-Free-Textual-Inversion/blob/main/train_inversion.py)
import cma
import argparse
import torch
import os
import numpy as np
import copy
from sklearn.decomposition import PCA
from diffusers import StableDiffusionPipeline, DDPMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
import torch.nn.functional as F
from utils import TextualInversionDataset
from tqdm import tqdm
class GradientFreePipeline:
def __init__(self, model_path, args, init_text_inversion=None, ):
self.tokenizer = CLIPTokenizer.from_pretrained(
os.path.join(model_path, 'tokenizer')
)
self.text_encoder = CLIPTextModel.from_pretrained(
os.path.join(model_path, 'text_encoder')
)
self.pipe = StableDiffusionPipeline.from_pretrained(
model_path,
text_encoder=self.text_encoder,
tokenizer=self.tokenizer,
).to(args.device)
if args.projection_modeling == 'prior_normal':
self.linear = torch.nn.Linear(args.intrinsic_dim, args.model_dim, bias=False).to(args.device)
embedding = self.text_encoder.get_input_embeddings().weight.clone().cpu()
mu_hat = np.mean(embedding.reshape(-1).detach().cpu().numpy())
std_hat = np.std(embedding.reshape(-1).detach().cpu().numpy())
mu = 0.0
std = args.alpha * std_hat / (np.sqrt(args.intrinsic_dim) * args.sigma)
# incorporate temperature factor
# temp = intrinsic_dim - std_hat * std_hat
# mu = mu_hat / temp
# std = std_hat / np.sqrt(temp)
print('[Embedding] mu: {} | std: {} [RandProj] mu: {} | std: {}'.format(mu_hat, std_hat, mu, std))
for p in self.linear.parameters():
torch.nn.init.normal_(p, mu, std)
elif args.projection_modeling == 'pca':
embedding = self.text_encoder.get_input_embeddings().weight.clone().cpu()
embedding = embedding.detach().cpu().numpy() # (49408, 768)
self.pca_model = PCA(n_components=args.intrinsic_dim)
self.pca_model.fit(embedding)
# Add the placeholder token in tokenizer
num_added_tokens = self.tokenizer.add_tokens(args.placeholder_token)
if num_added_tokens == 0:
raise ValueError(
f"The tokenizer already contains the token {args.placeholder_token}. Please pass a different"
" `placeholder_token` that is not already in the tokenizer."
)
# Convert the initializer_token, placeholder_token to ids
token_ids = self.tokenizer.encode(args.initializer_token, add_special_tokens=False)
initializer_token_id = token_ids[0]
placeholder_token_id = self.tokenizer.convert_tokens_to_ids(args.placeholder_token)
# Resize the token embeddings as we are adding new special tokens to the tokenizer
self.text_encoder.resize_token_embeddings(len(self.tokenizer))
# Initialise the newly added placeholder token with the embeddings of the initializer token
token_embeds = self.text_encoder.get_input_embeddings().weight.data
token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]
print('convert text inversion: ', args.placeholder_token, 'in id: ', str(placeholder_token_id))
self.placeholder_token_id = placeholder_token_id
self.placeholder_token = args.placeholder_token
self.num_call = 0
train_dataset = TextualInversionDataset(
data_root=args.train_data_dir,
tokenizer=self.tokenizer,
size=args.resolution,
placeholder_token=args.placeholder_token,
repeats=args.repeats,
learnable_property=args.learnable_property,
center_crop=args.center_crop,
set="train",
)
self.dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=args.repeats, shuffle=True)
self.batch_size = args.repeats
self.device = args.device
print('load data length: ', len(self.dataloader))
# optimize incremental elements or original inversion
if init_text_inversion is not None:
self.init_text_inversion = init_text_inversion.to(args.device)
else:
self.init_text_inversion = token_embeds[initializer_token_id].to(args.device)
self.args = args
self.best_inversion = None
def eval(self, inversion_embedding):
self.num_call += 1
pe_list = []
if isinstance(inversion_embedding, list): # multiple queries
for pe in inversion_embedding:
if self.args.projection_modeling == 'prior_normal':
z = torch.tensor(pe).type(torch.float32).to(self.device) # z
with torch.no_grad():
z = self.linear(z) # W_p Q
if self.init_text_inversion is not None:
z = z + self.init_text_inversion # W_p Q + p_0
elif self.args.projection_modeling == 'pca':
z = self.pca_model.inverse_transform(pe) # project the original text embedding space
z = torch.tensor(z).type(torch.float32).to(self.device)
print(z.shape)
print(self.init_text_inversion.shape)
if self.init_text_inversion is not None:
z = z + self.init_text_inversion
pe_list.append(z)
elif isinstance(inversion_embedding, np.ndarray): # single query or None
if self.args.projection_modeling == 'prior_normal':
inversion_embedding = torch.tensor(inversion_embedding).type(torch.float32).to(self.device) # z
with torch.no_grad():
inversion_embedding = self.linear(inversion_embedding) # W_p Q
elif self.args.projection_modeling == 'pca':
inversion_embedding = self.pca_model.inverse_transform(inversion_embedding)
inversion_embedding = torch.tensor(inversion_embedding).type(torch.float32).to(self.device)
if self.init_text_inversion is not None:
inversion_embedding = inversion_embedding + self.init_text_inversion # W_p Q + p_0
pe_list.append(inversion_embedding)
else:
raise ValueError(
f'[Inversion Embedding] Only support [list, numpy.ndarray], got `{type(inversion_embedding)}` instead.'
)
loss_list = []
print('begin to calculate loss')
# fixed time step for fair evaluation 公平评估的固定时间步长
noise_scheduler = DDPMScheduler.from_config('./ckpt/scheduler')
timesteps = torch.randint(
0, noise_scheduler.config.num_train_timesteps, (self.batch_size,), device=self.device
).long()
best_loss = 1000
best_inversion = None
for pe in tqdm(pe_list):
token_embeds = self.text_encoder.get_input_embeddings().weight.data
pe.to(self.text_encoder.get_input_embeddings().weight.dtype)
token_embeds[self.placeholder_token_id] = pe
loss = calculate_mse_loss(self.pipe, self.dataloader, self.device, noise_scheduler, timesteps)
if loss < best_loss:
best_loss = loss
best_inversion = pe
loss_list.append(loss)
# update total point
self.best_inversion = best_inversion
return loss_list
def save(self, output_path):
learned_embeds_dict = {self.placeholder_token: self.best_inversion.detach().cpu()}
torch.save(learned_embeds_dict, os.path.join(output_path, "learned_embeds.bin"))
def calculate_mse_loss(image_generator, dataloader, device, noise_scheduler, timesteps):
# print(image_generator.text_encoder.get_input_embeddings().weight.data[49408])
loss_cum = .0
with torch.no_grad():
for batch in dataloader:
# Convert images to latent space
latents = image_generator.vae.encode(batch["pixel_values"].to(device)).latent_dist.sample().detach()
latents = latents * 0.18215
# Sample noise that we'll add to the latents
noise = torch.randn(latents.shape).to(latents.device)
# Sample a random timestep for each image
# Add noise to the latents according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
# Get the text embedding for conditioning
encoder_hidden_states = image_generator.text_encoder(batch["input_ids"].to(device))[0]
# Predict the noise residual
noise_pred = image_generator.unet(noisy_latents, timesteps, encoder_hidden_states).sample
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
loss_cum += loss.item()
return loss_cum / len(dataloader)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--intrinsic_dim", default=256, type=int)
parser.add_argument("--k_shot", default=16, type=int)
parser.add_argument("--batch_size", default=16, type=int)
parser.add_argument("--budget", default=5000, type=int) # number of iterations
parser.add_argument("--popsize", default=20, type=int) # number of candidates
parser.add_argument("--bound", default=0, type=int)
parser.add_argument("--sigma", default=1, type=float)
parser.add_argument("--alpha", default=1, type=float)
parser.add_argument("--print_every", default=50, type=int)
parser.add_argument("--eval_every", default=100, type=int)
parser.add_argument("--alg", default='CMA', type=str) # support other advanced evelution strategy
parser.add_argument("--projection_modeling", default='pca', type=str) # decomposition method {'pca', 'prior_norm'}
parser.add_argument("--model_dim", default=512, type=int) # dim of textual inversion
parser.add_argument("--inversion_initialize", default='./save/initialize_emb.bin', type=str) # dim of textual inversion
parser.add_argument("--seed", default=2023, type=int)
parser.add_argument("--loss_type", default='noise', type=str)
parser.add_argument("--cat_or_add", default='add', type=str)
parser.add_argument("--device", default= torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
parser.add_argument("--parallel", default=False, type=bool, help='Whether to allow parallel evaluation')
parser.add_argument(
"--placeholder_token",
type=str,
default='<sks>',
help="A token to use as a placeholder for the concept.",
)
parser.add_argument(
"--initializer_token",
type=str,
default='coconut_seed_fruit_stage_CT',
help="A token to use as initializer word."
)
parser.add_argument(
"--inference_framework",
default='pt',
type=str,
help='''Which inference framework to use.
Currently supports `pt` and `ort`, standing for pytorch and Microsoft onnxruntime respectively'''
)
parser.add_argument(
"--onnx_model_path",
default=None,
type=str,
help='Path to your onnx model.'
)
parser.add_argument(
"--train_data_dir",
type=str,
default='./coconut_seed_fruit_stage',
help="A folder containing the training data of instance images.",
)
parser.add_argument(
"--learnable_property",
type=str,
default="object",
help="Choose between 'object' and 'style'"
)
parser.add_argument(
"--resolution",
type=int,
default=512,
help=(
"The resolution for input images, all the images in the train/validation dataset will be resized to this"
" resolution"
),
)
parser.add_argument(
"--center_crop", action="store_true", help="Whether to center crop images before resizing to resolution"
)
parser.add_argument("--repeats", type=int, default=5, help="How many times to repeat the training data.")
args = parser.parse_args()
cma_opts = {
'seed': args.seed,
'popsize': args.popsize,
'maxiter': args.budget if args.parallel else args.budget // args.popsize,
'verbose': -1,
}
if args.bound > 0:
cma_opts['bounds'] = [-1 * args.bound, 1 * args.bound]
if args.inversion_initialize is not None:
print('initialize textual inversion')
init_text_inversion = torch.load(args.inversion_initialize, map_location="cpu")["initialize"]
else:
init_text_inversion = None
pipeline = GradientFreePipeline(model_path='./ckpt', args=args, init_text_inversion=init_text_inversion)
es = cma.CMAEvolutionStrategy(args.intrinsic_dim * [0], args.sigma, inopts=cma_opts)
while not es.stop():
solutions = es.ask() # (popsize, intrinsic_dim)
fitnesses = pipeline.eval(solutions)
print(fitnesses) # loss for each point
es.tell(solutions, fitnesses)
pipeline.save('./save')
if __name__ == "__main__":
main()