Xander updates

lora_diffusion/cli_lora_pti.py

@@ -46,6 +46,8 @@
  prepare_clip_model_sets,
  evaluate_pipe,
  UNET_EXTENDED_TARGET_REPLACE,
+ parse_safeloras_embeds,
+ apply_learned_embed_in_clip,
 )
 
 def preview_training_batch(train_dataloader, mode, n_imgs = 40):
@@ -67,6 +69,52 @@ def preview_training_batch(train_dataloader, mode, n_imgs = 40):
  print(f"\nSaved {imgs_saved} preview training imgs to {outdir}")
  return
 
+def sim_matrix(a, b, eps=1e-8):
+ """
+ added eps for numerical stability
+ """
+ a_n, b_n = a.norm(dim=1)[:, None], b.norm(dim=1)[:, None]
+ a_norm = a / torch.max(a_n, eps * torch.ones_like(a_n))
+ b_norm = b / torch.max(b_n, eps * torch.ones_like(b_n))
+ sim_mt = torch.mm(a_norm, b_norm.transpose(0, 1))
+ return sim_mt
+
+
+def compute_pairwise_distances(x,y):
+ # compute the L2 distance of each row in x to each row in y (both are torch tensors)
+ # x is a torch tensor of shape (m, d)
+ # y is a torch tensor of shape (n, d)
+ # returns a torch tensor of shape (m, n)
+
+ n = y.shape[0]
+ m = x.shape[0]
+ d = x.shape[1]
+
+ x = x.unsqueeze(1).expand(m, n, d)
+ y = y.unsqueeze(0).expand(m, n, d)
+
+ return torch.pow(x - y, 2).sum(2)
+
+
+def print_most_similar_tokens(tokenizer, optimized_token, text_encoder, n=10):
+ with torch.no_grad():
+ # get all the token embeddings:
+ token_embeds = text_encoder.get_input_embeddings().weight.data
+
+ # Compute the cosine-similarity between the optimized tokens and all the other tokens
+ similarity = sim_matrix(optimized_token.unsqueeze(0), token_embeds).squeeze()
+ similarity = similarity.detach().cpu().numpy()
+
+ distances = compute_pairwise_distances(optimized_token.unsqueeze(0), token_embeds).squeeze()
+ distances = distances.detach().cpu().numpy()
+
+ # print similarity for the most similar tokens:
+ most_similar_tokens = np.argsort(similarity)[::-1]
+
+ print(f"{tokenizer.decode(most_similar_tokens[0])} --> mean: {optimized_token.mean().item():.3f}, std: {optimized_token.std().item():.3f}, norm: {optimized_token.norm():.4f}")
+ for token_id in most_similar_tokens[1:n+1]:
+ print(f"sim of {similarity[token_id]:.3f} & L2 of {distances[token_id]:.3f} with \"{tokenizer.decode(token_id)}\"")
+
 
 def get_models(
  pretrained_model_name_or_path,
@@ -139,19 +187,21 @@ def get_models(
  pretrained_vae_name_or_path or pretrained_model_name_or_path,
  subfolder=None if pretrained_vae_name_or_path else "vae",
  revision=None if pretrained_vae_name_or_path else revision,
+ local_files_only = True,
  )
  unet = UNet2DConditionModel.from_pretrained(
  pretrained_model_name_or_path,
  subfolder="unet",
  revision=revision,
+ local_files_only = True,
  )
 
  return (
  text_encoder.to(device),
  vae.to(device),
  unet.to(device),
  tokenizer,
- placeholder_token_ids,
+ placeholder_token_ids
  )
 
 
@@ -477,12 +527,13 @@ def train_inversion(
 
  if global_step % accum_iter == 0:
  # print gradient of text encoder embedding
- print(
- text_encoder.get_input_embeddings()
- .weight.grad[index_updates, :]
- .norm(dim=-1)
- .mean()
- )
+ if 0:
+ print(
+ text_encoder.get_input_embeddings()
+ .weight.grad[index_updates, :]
+ .norm(dim=-1)
+ .mean()
+ )
  optimizer.step()
  optimizer.zero_grad()
 
@@ -517,8 +568,10 @@ def train_inversion(
  index_no_updates
  ] = orig_embeds_params[index_no_updates]
 
- for i, t in enumerate(optimizing_embeds):
- print(f"token {i} --> mean: {t.mean().item():.3f}, std: {t.std().item():.3f}, norm: {t.norm():.4f}")
+ if global_step % 50 == 0:
+ print("------------------------------")
+ for i, t in enumerate(optimizing_embeds):
+ print_most_similar_tokens(tokenizer, t, text_encoder)
 
  global_step += 1
  progress_bar.update(1)
@@ -537,7 +590,7 @@ def train_inversion(
  placeholder_token_ids=placeholder_token_ids,
  placeholder_tokens=placeholder_tokens,
  save_path=os.path.join(
- save_path, f"step_inv_{global_step}.safetensors"
+ save_path, f"step_inv_{global_step:04d}.safetensors"
  ),
  save_lora=False,
  )
@@ -583,7 +636,7 @@ def train_inversion(
  return
 
 import matplotlib.pyplot as plt
-def plot_loss_curve(losses, name, moving_avg=20):
+def plot_loss_curve(losses, name, moving_avg=5):
  losses = np.array(losses)
  losses = np.convolve(losses, np.ones(moving_avg)/moving_avg, mode='valid')
  plt.plot(losses)
@@ -654,7 +707,7 @@ def perform_tuning(
  vae,
  text_encoder,
  scheduler,
- optimized_embeddings = text_encoder.get_input_embeddings().weight[:, :],
+ optimized_embeddings = text_encoder.get_input_embeddings().weight[~index_no_updates, :], 
  train_inpainting=train_inpainting,
  t_mutliplier=0.8,
  mixed_precision=True,
@@ -683,6 +736,12 @@ def perform_tuning(
  index_no_updates
  ] = orig_embeds_params[index_no_updates]
 
+ if global_step % 100 == 0:
+ optimizing_embeds = text_encoder.get_input_embeddings().weight[~index_no_updates]
+ print("------------------------------")
+ for i, t in enumerate(optimizing_embeds):
+ print_most_similar_tokens(tokenizer, t, text_encoder)
+
 
  global_step += 1
 
@@ -696,7 +755,7 @@ def perform_tuning(
  placeholder_token_ids=placeholder_token_ids,
  placeholder_tokens=placeholder_tokens,
  save_path=os.path.join(
- save_path, f"step_{global_step}.safetensors"
+ save_path, f"step_{global_step:04d}.safetensors"
  ),
  target_replace_module_text=lora_clip_target_modules,
  target_replace_module_unet=lora_unet_target_modules,
@@ -706,16 +765,15 @@ def perform_tuning(
  .mean()
  .item()
  )
-
  print("LORA Unet Moved", moved)
+
  moved = (
  torch.tensor(
  list(itertools.chain(*inspect_lora(text_encoder).values()))
  )
  .mean()
  .item()
  )
-
  print("LORA CLIP Moved", moved)
 
  if log_wandb:
@@ -778,6 +836,7 @@ def train(
  placeholder_tokens: str = "",
  placeholder_token_at_data: Optional[str] = None,
  initializer_tokens: Optional[str] = None,
+ load_pretrained_inversion_embeddings_path: Optional[str] = None,
  seed: int = 42,
  resolution: int = 512,
  color_jitter: bool = True,
@@ -788,7 +847,8 @@ def train(
  save_steps: int = 100,
  gradient_accumulation_steps: int = 4,
  gradient_checkpointing: bool = False,
- lora_rank: int = 4,
+ lora_rank_unet: int = 4,
+ lora_rank_text_encoder: int = 4,
  lora_unet_target_modules={"CrossAttention", "Attention", "GEGLU"},
  lora_clip_target_modules={"CLIPAttention"},
  lora_dropout_p: float = 0.0,
@@ -825,6 +885,10 @@ def train(
  script_start_time = time.time()
  torch.manual_seed(seed)
 
+ if use_template == "person" and not use_face_segmentation_condition:
+ print("### WARNING ### : Using person template without face segmentation condition")
+ print("When training people, it is highly recommended to use face segmentation condition!!")
+
  # Get a dict with all the arguments:
  args_dict = locals()
 
@@ -841,7 +905,7 @@ def train(
 
  if output_dir is not None:
  os.makedirs(output_dir, exist_ok=True)
- # print(placeholder_tokens, initializer_tokens)
+
  if len(placeholder_tokens) == 0:
  placeholder_tokens = []
  print("PTI : Placeholder Tokens not given, using null token")
@@ -874,6 +938,7 @@ def train(
 
  print("PTI : Placeholder Tokens", placeholder_tokens)
  print("PTI : Initializer Tokens", initializer_tokens)
+ print("PTI : Token Map: ", token_map)
 
  # get the models
  text_encoder, vae, unet, tokenizer, placeholder_token_ids = get_models(
@@ -886,7 +951,8 @@ def train(
  )
 
  noise_scheduler = DDPMScheduler.from_config(
- pretrained_model_name_or_path, subfolder="scheduler"
+ pretrained_model_name_or_path, subfolder="scheduler", 
+ local_files_only = True,
  )
 
  if gradient_checkpointing:
@@ -925,8 +991,6 @@ def train(
  train_inpainting=train_inpainting,
  )
 
- train_dataset.blur_amount = 200
-
  if train_inpainting:
  assert not cached_latents, "Cached latents not supported for inpainting"
 
@@ -963,7 +1027,7 @@ def train(
  vae = None
 
  # STEP 1 : Perform Inversion
- if perform_inversion and not cached_latents:
+ if perform_inversion and not cached_latents and (load_pretrained_inversion_embeddings_path is None):
  preview_training_batch(train_dataloader, "inversion")
 
  print("PTI : Performing Inversion")
@@ -1014,34 +1078,44 @@ def train(
  del ti_optimizer
  print("############### Inversion Done ###############")
 
+ elif load_pretrained_inversion_embeddings_path is not None:
+
+ print("PTI : Loading pretrained inversion embeddings..")
+ from safetensors.torch import safe_open
+ # Load the pretrained embeddings from the lora file:
+ safeloras = safe_open(load_pretrained_inversion_embeddings_path, framework="pt", device="cpu")
+ #monkeypatch_or_replace_safeloras(pipe, safeloras)
+ tok_dict = parse_safeloras_embeds(safeloras)
+ apply_learned_embed_in_clip(
+ tok_dict,
+ text_encoder,
+ tokenizer,
+ idempotent=True,
+ )
+
  # Next perform Tuning with LoRA:
  if not use_extended_lora:
  unet_lora_params, _ = inject_trainable_lora(
  unet,
- r=lora_rank,
+ r=lora_rank_unet,
  target_replace_module=lora_unet_target_modules,
  dropout_p=lora_dropout_p,
  scale=lora_scale,
  )
  print("PTI : not use_extended_lora...")
+ print("PTI : Will replace modules: ", lora_unet_target_modules)
  else:
  print("PTI : USING EXTENDED UNET!!!")
  lora_unet_target_modules = (
  lora_unet_target_modules | UNET_EXTENDED_TARGET_REPLACE
  )
  print("PTI : Will replace modules: ", lora_unet_target_modules)
  unet_lora_params, _ = inject_trainable_lora_extended(
- unet, r=lora_rank, target_replace_module=lora_unet_target_modules
+ unet, r=lora_rank_unet, target_replace_module=lora_unet_target_modules
  )
 
- n_optimizable_unet_params = sum(
- [el.numel() for el in itertools.chain(*unet_lora_params)]
- )
- print("PTI : n_optimizable_unet_params: ", n_optimizable_unet_params)
-
- print(f"PTI : has {len(unet_lora_params)} lora")
- print("PTI : Before training:")
- inspect_lora(unet)
+ #n_optimizable_unet_params = sum([el.numel() for el in itertools.chain(*unet_lora_params)])
+ #print("PTI : Number of optimizable UNET parameters: ", n_optimizable_unet_params)
 
  params_to_optimize = [
  {"params": itertools.chain(*unet_lora_params), "lr": unet_lr},
@@ -1073,15 +1147,15 @@ def train(
  text_encoder_lora_params, _ = inject_trainable_lora(
  text_encoder,
  target_replace_module=lora_clip_target_modules,
- r=lora_rank,
+ r=lora_rank_text_encoder,
  )
  params_to_optimize += [
- {
- "params": itertools.chain(*text_encoder_lora_params),
- "lr": text_encoder_lr,
- }
+ {"params": itertools.chain(*text_encoder_lora_params),
+ "lr": text_encoder_lr}
  ]
- inspect_lora(text_encoder)
+
+ #n_optimizable_text_Encoder_params = sum( [el.numel() for el in itertools.chain(*text_encoder_lora_params)])
+ #print("PTI : Number of optimizable text-encoder parameters: ", n_optimizable_text_Encoder_params)
 
  lora_optimizers = optim.AdamW(params_to_optimize, weight_decay=weight_decay_lora)
 
@@ -1090,8 +1164,6 @@ def train(
  print("Training text encoder!")
  text_encoder.train()
 
- train_dataset.blur_amount = 70
-
  lr_scheduler_lora = get_scheduler(
  lr_scheduler_lora,
  optimizer=lora_optimizers,
@@ -1101,6 +1173,22 @@ def train(
  if not cached_latents: 
  preview_training_batch(train_dataloader, "tuning")
 
+ #print("PTI : n_optimizable_unet_params: ", n_optimizable_unet_params)
+ print(f"PTI : has {len(unet_lora_params)} lora")
+ print("PTI : Before training:")
+
+ moved = (
+ torch.tensor(list(itertools.chain(*inspect_lora(unet).values())))
+ .mean().item())
+ print(f"LORA Unet Moved {moved:.6f}")
+
+
+ moved = (
+ torch.tensor(
+ list(itertools.chain(*inspect_lora(text_encoder).values()))
+ ).mean().item())
+ print(f"LORA CLIP Moved {moved:.6f}")
+
  perform_tuning(
  unet,
  vae,
@@ -1132,6 +1220,8 @@ def train(
  training_time = time.time() - script_start_time
  print(f"Training time: {training_time/60:.1f} minutes")
  args_dict["training_time_s"] = int(training_time)
+ args_dict["n_epochs"] = math.ceil(max_train_steps_tuning / len(train_dataloader.dataset))
+ args_dict["n_training_imgs"] = len(train_dataloader.dataset)
 
  # Save the args_dict to the output directory as a json file:
  with open(os.path.join(output_dir, "lora_training_args.json"), "w") as f: