import csv import dataclasses import subprocess from copy import deepcopy import itertools from concurrent.futures import ThreadPoolExecutor import pathlib from typing import List import diffusers import transformers import safetensors.torch import torch.utils.data from tqdm import tqdm from datetime import datetime import random import os import time from torch.utils.tensorboard import SummaryWriter torch.manual_seed(0) random.seed(0) LATENTS_OUTPUT_DIR = pathlib.Path("latents") CAPTIONS_OUTPUT_DIR = pathlib.Path("captions2") DANBOORU_ARTISTS_PATH = pathlib.Path("danbooru_artist.csv") E621_ARTISTS_PATH = pathlib.Path("e621_artist.csv") LOCK_FILE = "safetensors.lock" device = torch.device("cuda") dtype = torch.float16 train_logger = SummaryWriter(f"logs/pony_scoreless_{datetime.now().strftime('%Y%m%d_%H%M%S')}") def accumulate_grads(): batch_size = 1 epochs = 1 tokenizer = create_tokenizer(device) model_a = diffusers.StableDiffusionXLPipeline.from_single_file( "NoobAI-XL-v1.1.safetensors", torch_dtype=dtype, ) delattr(model_a, "vae") model_a.unet.to(device=device) # model_a.unet.enable_xformers_memory_efficient_attention() model_a.unet.enable_gradient_checkpointing() model_a.text_encoder.to(device=device) model_a.text_encoder.gradient_checkpointing_enable() model_a.text_encoder_2.to(device=device) model_a.text_encoder_2.gradient_checkpointing_enable() model_a.text_encoder_combined = CombinedCLIPTextEncoder(model_a.text_encoder, model_a.text_encoder_2, batch_size) model_b = diffusers.StableDiffusionXLPipeline.from_single_file( "animagine-xl-4.0.safetensors", torch_dtype=dtype, ) delattr(model_b, "vae") model_b.unet.to(device=device) # model_b.unet.enable_xformers_memory_efficient_attention() model_b.unet.enable_gradient_checkpointing() model_b.text_encoder.to(device=device) model_b.text_encoder.gradient_checkpointing_enable() model_b.text_encoder_2.to(device=device) model_b.text_encoder_2.gradient_checkpointing_enable() model_b.text_encoder_combined = CombinedCLIPTextEncoder(model_b.text_encoder, model_b.text_encoder_2, batch_size) model_a.unet.eval() model_a.text_encoder.eval() model_a.text_encoder_2.eval() model_b.unet.eval() model_b.text_encoder.eval() model_b.text_encoder_2.eval() # shared_stats = {} # stats_lock = threading.Lock() # # Two barriers for synchronization between two threads. # grad_barrier1 = threading.Barrier(2) # grad_barrier2 = threading.Barrier(2) # def scaling_hook_factory(key, branch_id, target_scale=1.0): # nonlocal shared_stats, stats_lock, grad_barrier1, grad_barrier2 # def scaling_hook(_module, _grad_input, grad_output): # """ # A full-backward hook that: # 1. Computes, for each non-None tensor in grad_output, its maximum absolute value. # We store these in a dictionary (keyed by output index). # 2. Waits once until both threads have stored their local max values. # 3. Computes, for each output index, the global maximum from both models. # 4. Waits a second time to ensure synchronization before clearing the shared stats. # 5. Scales each non-None output tensor independently using its computed scaling factor. # Outputs that are None are passed through unchanged. # """ # # Step 1: Compute and store local maximums per output index. # print(f"backprop for {key}") # local_maxes = {} # for i, g in enumerate(grad_output): # if g is not None: # local_maxes[i] = g.detach().abs().max().cpu().item() # with stats_lock: # shared_stats[f"{key}_{branch_id}"] = local_maxes # # Step 2: Wait until both threads have stored their values. # grad_barrier1.wait() # # Step 3: Compute the global maximum for each output index. # with stats_lock: # stats_a = shared_stats.get(f"{key}_a", {}) # stats_b = shared_stats.get(f"{key}_b", {}) # # Build a dictionary for global max per output index. # global_maxes = {} # for i in local_maxes.keys(): # assert i in stats_a and i in stats_b, key # global_maxes[i] = max(stats_a[i], stats_b[i]) # # Step 4: Wait again to ensure both threads have computed the global values. # barrier_val = grad_barrier2.wait() # # Let only one thread clear the shared stats. # if barrier_val == 0: # with stats_lock: # shared_stats.pop(f"{key}_a") # shared_stats.pop(f"{key}_b") # # Step 5: For each output tensor, compute a scaling factor and apply it. # scaled_outputs = [] # for i, g in enumerate(grad_output): # if g is not None: # global_max = global_maxes[i] # # Compute scaling factor only if global_max is positive and below target_scale. # if 0 < global_max < target_scale: # g = g * (target_scale / global_max) # scaled_outputs.append(g) # else: # scaled_outputs.append(None) # return tuple(scaled_outputs) # return scaling_hook # for model, branch_id in zip((model_a, model_b), ("a", "b")): # for k, v in get_modules(model): # if k.endswith("transformer_blocks") or k.endswith("encoder.layers"): # for i, module in enumerate(v): # module.register_full_backward_hook(scaling_hook_factory(f"{k}.{i}", branch_id)) scheduler = create_scheduler(device) data_loader = get_data_loader(tokenizer, batch_size) total_steps = 0 log_scalars_a = {} log_scalars_b = {} log_scalars_sync = {} n1 = torch.tensor(-1, device=device, dtype=torch.long) ldexp_offset = torch.tensor(20, device=device, dtype=torch.long) def create_hook(param, k, log_scalars): param.grad = torch.zeros_like(param) log_scalars[k] = ldexp_offset.clone() def hook(grad): nonlocal param, log_scalars, k while True: new_grad = param.grad + grad.abs().ldexp(log_scalars[k]) if not new_grad.isfinite().all(): # overflow log_scalars[k] -= 1 param.grad.ldexp_(n1) else: break param.grad.copy_(new_grad) return param.grad return hook for model, log_scalars in ((model_a, log_scalars_a), (model_b, log_scalars_b)): for k, v in get_params(model): v.register_hook(create_hook(v, k, log_scalars)) # for model, path in ((model_a, "grads_a.safetensors"), (model_b, "grads_b.safetensors")): # with safetensors.safe_open(path, "pt") as f: # for k, v in get_params(model): # if k in f.keys(): # v.grad = f.get_tensor(k).to(v) noisy_latents = timesteps = time_ids = None def get_pred(args): nonlocal noisy_latents, timesteps, time_ids model, tokens = args txt = model.text_encoder_combined(tokens[0]) return model.unet( noisy_latents, timesteps, encoder_hidden_states=txt["conds"], added_cond_kwargs={ "text_embeds": txt["pooled"], "time_ids": time_ids, }, ).sample params = list(v for k, v in itertools.chain(get_params(model_a), get_params(model_b))) with ThreadPoolExecutor(max_workers=2) as worker: for epoch_i in range(epochs): for step_i, (latent_infos, tokens_a, tokens_b, post_ids) in enumerate(tqdm(data_loader)): latents = torch.cat([latent_info["latent"] for latent_info in latent_infos], dim=0).to(device=device, dtype=dtype) crop_hw = torch.stack([latent_info["crop_hw"] for latent_info in latent_infos]).to(device=device) orig_hw = torch.stack([latent_info["orig_hw"] for latent_info in latent_infos]).to(device=device) noise, noisy_latents, timesteps = get_noise_noisy_latents_and_timesteps(scheduler, latents) time_ids = get_add_time_ids(orig_hw, crop_hw) # if step_i < 1000: # total_steps += batch_size # continue pred_a, pred_b = worker.map(get_pred, ((model_a, tokens_a), (model_b, tokens_b))) mse = torch.nn.functional.mse_loss(pred_a, pred_b, reduction="none").flatten(start_dim=1).mean(dim=-1) loss = (mse / mse.detach()).mean() train_logger.add_scalar("grads/loss", loss.item(), total_steps) train_logger.add_scalar("grads/loss_raw", mse.mean().item(), total_steps) train_logger.add_scalar("grads/timestep", timesteps[0].item(), total_steps) torch.autograd.grad(loss, params, retain_graph=False, allow_unused=True) # calls backward hooks for (k, v_a), (k_b, v_b) in zip(get_params(model_a), get_params(model_b)): assert k == k_b if v_a.grad is not None and v_b.grad is not None: while log_scalars_a[k] > log_scalars_b[k]: log_scalars_a[k] -= 1 v_a.grad.ldexp_(n1) while log_scalars_b[k] > log_scalars_a[k]: log_scalars_b[k] -= 1 v_b.grad.ldexp_(n1) log_scalars_sync[k] = log_scalars_a[k] if (step_i + 1) % 10 == 0: train_logger.add_scalar("grads/max_a", max(v.grad.max().item() for k, v in get_params(model_a) if v.grad is not None), total_steps) train_logger.add_scalar("grads/max_b", max(v.grad.max().item() for k, v in get_params(model_b) if v.grad is not None), total_steps) if (step_i + 1) % 1000 == 0: save_grads(model_a, "grads_a.safetensors", first=True) safetensors.torch.save_file(log_scalars_sync, "log_scalars.safetensors") save_grads(model_b, "grads_b.safetensors", last=True) total_steps += batch_size def get_modules(model): return itertools.chain( prefix_iter(model.unet.named_modules(), "unet."), prefix_iter(model.text_encoder.named_modules(), "text_encoder."), prefix_iter(model.text_encoder_2.named_modules(), "text_encoder_2."), ) def get_params(model): return itertools.chain( prefix_iter(model.unet.named_parameters(), "unet."), prefix_iter(model.text_encoder.named_parameters(), "text_encoder."), prefix_iter(model.text_encoder_2.named_parameters(), "text_encoder_2."), ) def prefix_iter(item_iter, prefix): return ((prefix + k, v) for k, v in item_iter) def save_grads(model, path, first=False, last=False): if first: wait_for_lock_removal() safetensors.torch.save_file( {k: v.grad.cpu().contiguous() for k, v in get_params(model) if v.grad is not None}, path, ) if last: # Create a lock file to signal that new checkpoints have been saved with open(LOCK_FILE, "w") as f: f.write("pending download") print("Checkpoint pair saved, lock file created.") def wait_for_lock_removal(poll_interval=5): """Wait until the lock file is removed by the local download script.""" while os.path.exists(LOCK_FILE): time.sleep(poll_interval) def create_scheduler(device: torch.device): scheduler = diffusers.DDPMScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False, ) inv_snr = ((1-scheduler.alphas_cumprod) / scheduler.alphas_cumprod).to(device) scheduler.inv_snr = inv_snr scheduler.inv_snr_weights = inv_snr / inv_snr.sum() return scheduler def debiased_loss_scaling(timesteps, noise_scheduler): return noise_scheduler.inv_snr[timesteps] def get_noise_noisy_latents_and_timesteps(scheduler, latents): batch_size = latents.shape[0] noise = torch.randn_like(latents, device=latents.device) timesteps = torch.multinomial(scheduler.inv_snr_weights, batch_size) noisy_latents = scheduler.add_noise(latents, noise, timesteps) return noise, noisy_latents, timesteps def get_add_time_ids(original_size, crops_coords_top_left): add_time_ids = torch.cat([ original_size, crops_coords_top_left, torch.tensor([[1024]*2], device=original_size.device).expand(len(original_size), -1), ], dim=1) return add_time_ids def get_data_loader(tokenizer, batch_size: int): return torch.utils.data.DataLoader( PromptDataset(tokenizer), batch_size=batch_size, shuffle=True, collate_fn=lambda x: zip(*x), ) @dataclasses.dataclass class ArtistScore: artist_tag: str count: int class PromptDataset(torch.utils.data.Dataset): def __init__(self, tokenizer): self.tokenizer = tokenizer self.latent_paths = list(LATENTS_OUTPUT_DIR.iterdir()) with open(DANBOORU_ARTISTS_PATH, "r", encoding='utf-8') as f: reader = csv.DictReader(f) self.b_artists = [ArtistScore(r["trigger"], int(r["count"])) for r in reader if r["artist"] != "banned_artist"] self.b_artists.sort(key=lambda t: t.count, reverse=True) self.b_artist_scores = torch.tensor(list(map(lambda t: t.count, self.b_artists)), device=device, dtype=torch.float32) self.b_artist_scores /= self.b_artist_scores.sum() with open(E621_ARTISTS_PATH, "r", encoding='utf-8') as f: reader = csv.DictReader(f,) self.a_artists = self.b_artists + [ArtistScore(r["trigger"], int(r["count"])) for r in reader if r["artist"] not in ["conditional_dnp", "avoid_posting", "unknown_artist", "third-party_edit", "sound_warning", "anonymous_artist"]] self.a_artists.sort(key=lambda t: t.count, reverse=True) self.a_artist_scores = torch.tensor(list(map(lambda t: t.count, self.a_artists)), device=device, dtype=torch.float32) self.a_artist_scores /= self.a_artist_scores.sum() self.a_prefix = "masterpiece, best quality, newest, absurdres, highres, safe, " self.b_suffix = ", masterpiece, high score, great score, absurdres" def __len__(self): return len(self.latent_paths) def __getitem__(self, item): post_id = self.latent_paths[item].stem latent = safetensors.torch.load_file(LATENTS_OUTPUT_DIR / f"{post_id}.safetensors", device=str(device)) caption = (CAPTIONS_OUTPUT_DIR / f"{post_id}.txt").read_text() caption_a = self.a_prefix + caption caption_b = caption + self.b_suffix if item % 2 == 0: artist_a = self.a_artists[torch.multinomial(self.a_artist_scores, 1).item()] caption_a = artist_a.artist_tag + ", " + caption_a else: artist_b = self.b_artists[torch.multinomial(self.b_artist_scores, 1).item()] caption_b = artist_b.artist_tag + ", " + caption_b tokens_a = self.tokenizer.chunk_tokens(self.tokenizer([caption_a.replace("),", ") ,")])) tokens_b = self.tokenizer.chunk_tokens(self.tokenizer([caption_b.replace("),", ") ,")])) return latent, tokens_a, tokens_b, post_id class CombinedCLIPTextEncoder(torch.nn.Module): def __init__(self, clip_l, clip_g, batch_size): super().__init__() assert batch_size == 1 self.clip_l = clip_l self.clip_g = clip_g def forward(self, tokens): tokens_clip_l = tokens["clip_l"].copy() del tokens_clip_l["prompt_starts"] tokens_clip_g = tokens["clip_g"].copy() clip_g_starts = tokens_clip_g.pop("prompt_starts") clip_l_encoded = self.clip_l(**tokens_clip_l, output_hidden_states=True, return_dict=True) clip_g_encoded = self.clip_g(**tokens_clip_g, output_hidden_states=True, return_dict=True) combined_encoded = torch.cat([clip_l_encoded["hidden_states"][-2], clip_g_encoded["hidden_states"][-2]], dim=-1) combined_encoded_reshape = combined_encoded.reshape(1, -1, 2048) return { "conds": combined_encoded_reshape, "pooled": clip_g_encoded.text_embeds[clip_g_starts], } def create_tokenizer(device: torch.device): tokenizer_l = transformers.CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") tokenizer_g = transformers.CLIPTokenizer.from_pretrained("laion/CLIP-ViT-g-14-laion2B-s34B-b88K") return CombinedCLIPTokenizer(tokenizer_l, tokenizer_g, device) class CombinedCLIPTokenizer(torch.nn.Module): comma_token = 267 def __init__(self, tokenizer_l, tokenizer_g, output_device: torch.device): super().__init__() self.tokenizer_l = tokenizer_l self.tokenizer_g = tokenizer_g self.output_device = output_device def forward(self, prompts: List[str]) -> dict: tokens_l = self.tokenizer_l(prompts, add_special_tokens=False) return { "clip_l": tokens_l, "clip_g": deepcopy(tokens_l), } def chunk_tokens(self, tokens: dict): return { "clip_l": self._chunk_tokens_impl(self.tokenizer_l, tokens["clip_l"]), "clip_g": self._chunk_tokens_impl(self.tokenizer_g, tokens["clip_g"]), } def _chunk_tokens_impl(self, tokenizer, tokens: dict): input_ids = [] attention_masks = [] chunk_counts = [] for prompt, mask in zip(tokens["input_ids"], tokens["attention_mask"]): last_comma = 0 current_chunk = [] chunks = [] chunks_attn = [] def next_chunk(): nonlocal current_chunk current_chunk = [tokenizer.bos_token_id] + current_chunk + [tokenizer.eos_token_id] num_tokens = len(current_chunk) current_chunk.extend([tokenizer.pad_token_id] * (77 - num_tokens)) chunks.append(current_chunk) current_chunk = [] chunks_attn.append([1] * num_tokens + [0] * (77 - num_tokens)) for token_i, token in enumerate(prompt): is_last_token = token_i == len(prompt) - 1 seq_suffix = prompt[last_comma:token_i + int(is_last_token)] if token == self.comma_token or is_last_token: if len(current_chunk) + len(seq_suffix) > 77 - 2: # leave space for bos and eos next_chunk() seq_suffix = prompt[last_comma+1:token_i + int(is_last_token)] # remove leading comma # can always append, sequences without commas will never be longer than 77 tokens current_chunk.extend(seq_suffix) last_comma = token_i if current_chunk or not chunks: next_chunk() chunk_counts.append(len(chunks)) input_ids.extend(chunks) attention_masks.extend(chunks_attn) return { "input_ids": torch.tensor(input_ids, device=self.output_device), "attention_mask": torch.tensor(attention_masks, device=self.output_device), "prompt_starts": torch.tensor([0] + chunk_counts[:-1], device=self.output_device).cumsum(dim=0), } def shutdown_machine(): """Shutdown the machine. Adjust the command as necessary for your environment.""" wait_for_lock_removal() print("All checkpoints have been downloaded. Shutting down the machine.") try: subprocess.run("runpodctl stop pod $RUNPOD_POD_ID", shell=True, check=True) except Exception as e: print(f"Error shutting down: {e}") if __name__ == "__main__": accumulate_grads() shutdown_machine()