""" LoRA fine-tuning for PixelDiT using precomputed image+caption embeddings. Flow matching loss: predict velocity (noise - x), noisy image via shifted schedule. Only LoRA weights update — base model is fully frozen. Usage: # Precompute first: python scripts/precompute_lora_data.py --images /data/my_images --out /data/lora_cache # Train: python scripts/train_lora.py --data /data/lora_cache --out lora_out/ --epochs 100 # Inference with trained LoRA: from peft import PeftModel from pixeldit.modeling_pixeldit_hf import PixelDiTModel model = PixelDiTModel.from_pretrained("madtune/pixeldit-diffusers", subfolder="transformer") model = PeftModel.from_pretrained(model, "lora_out/") """ import argparse import json import os import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset from tqdm import tqdm # ---- Flow schedule (matching NVIDIA's training config) ---------------------- _T = 1000 _TXT_MAX = 300 _GEMMA_ID = "Efficient-Large-Model/gemma-2-2b-it" _SELECT_IDX = [0] + list(range(-(_TXT_MAX - 1), 0)) def _build_flow_schedule(flow_shift: float): # sigmas[t] ≈ shift * (t/1000) / (1 + (shift-1) * (t/1000)) # matches FlowMatchEulerDiscreteScheduler with the same shift value betas = np.linspace(1.0, 0.001, _T, dtype=np.float64) sigmas_raw = 1.0 - betas sigmas = flow_shift * sigmas_raw / (1 + (flow_shift - 1) * sigmas_raw) alphas = 1.0 - sigmas return torch.from_numpy(sigmas).float(), torch.from_numpy(alphas).float() def q_sample(x, t, noise, alphas, sigmas): a = alphas[t].view(-1, 1, 1, 1) s = sigmas[t].view(-1, 1, 1, 1) return a * x + s * noise # ---- Dataset ---------------------------------------------------------------- class LoraDataset(Dataset): def __init__(self, data_dir): meta_path = os.path.join(data_dir, "meta.json") self.meta = {} if os.path.exists(meta_path): with open(meta_path, "r", encoding="utf-8") as f: self.meta = json.load(f) print( f"[dataset] encoder={self.meta.get('encoder', 'unknown')} " f"n={self.meta.get('n_samples')} " f"emb_dim={self.meta.get('emb_dim')} " f"trigger={self.meta.get('trigger') or 'none'}" ) self.imgs = np.load(os.path.join(data_dir, "lora_images.npy"), mmap_mode="r") self.embs = np.load(os.path.join(data_dir, "lora_embs.npy"), mmap_mode="r") self.masks = np.load(os.path.join(data_dir, "lora_masks.npy"), mmap_mode="r") captions_path = os.path.join(data_dir, "captions.json") self.captions = None if os.path.exists(captions_path): with open(captions_path, "r", encoding="utf-8") as f: self.captions = json.load(f) if len(self.captions) != len(self.imgs): raise RuntimeError( f"{captions_path} has {len(self.captions)} captions for {len(self.imgs)} images" ) assert len(self.imgs) == len(self.embs) == len(self.masks) # sanity check: catch all-zeros from a failed precompute run sample = self.embs[:min(4, len(self.embs))] if np.count_nonzero(sample) == 0: raise RuntimeError( f"lora_embs.npy in {data_dir} are ALL ZEROS — " "re-run precompute_lora_data.py to regenerate." ) def __len__(self): return len(self.imgs) @property def img_size(self): return self.meta.get("img_size", 512) def __getitem__(self, idx): img = torch.from_numpy(self.imgs[idx].astype(np.float32)) # [3, H, H] mask = torch.from_numpy(self.masks[idx].astype(np.float32)) # [300] return idx, img, mask # return idx so we can look up learnable embeddings # ---- Learnable Embeddings --------------------------------------------------- class LearnableEmbeddings(nn.Module): """Wrapper to make embeddings learnable parameters during training.""" def __init__(self, embeddings: torch.Tensor): super().__init__() # embeddings: [N, 300, 2304] float32 self.embs = nn.Parameter(embeddings) def forward(self, indices: torch.Tensor): """Return embeddings for given indices. indices: [B]""" return self.embs[indices] # [B, 300, 2304] def encode_gemma_batch(text_encoder, tokenizer, captions, device): toks = tokenizer( captions, max_length=_TXT_MAX, padding="max_length", truncation=True, return_tensors="pt", ).to(device) emb = text_encoder( input_ids=toks.input_ids, attention_mask=toks.attention_mask, ).last_hidden_state return emb[:, _SELECT_IDX, :] # ---- Main ------------------------------------------------------------------- def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", required=True, help="precomputed cache dir") ap.add_argument("--out", default="lora_out/", help="output dir for LoRA weights") ap.add_argument("--model", default="madtune/pixeldit-diffusers") ap.add_argument("--epochs", type=int, default=100) ap.add_argument("--batch", type=int, default=2) ap.add_argument("--accum", type=int, default=4, help="gradient accumulation steps") ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--lora_r", type=int, default=16) ap.add_argument("--lora_alpha", type=int, default=16) ap.add_argument("--train_text_encoder", action="store_true", help="train a Gemma LoRA together with the PixelDiT transformer LoRA") ap.add_argument("--text_encoder_model", default=_GEMMA_ID, help="Gemma model id/path") ap.add_argument("--text_device", default=None, help="device for Gemma LoRA training; default follows --device") ap.add_argument("--text_lora_r", type=int, default=8) ap.add_argument("--text_lora_alpha", type=int, default=8) ap.add_argument("--text_lora_targets", default="q_proj,k_proj,v_proj,o_proj", help="comma-separated Gemma module names for PEFT LoRA") ap.add_argument("--text_lr", type=float, default=None, help="Gemma LoRA LR; default is lr * 0.25") ap.add_argument("--cfg_drop", type=float, default=0.1, help="CFG dropout probability") ap.add_argument("--device", default="cuda:0") ap.add_argument("--save_every", type=int, default=10) ap.add_argument("--flow_shift", type=float, default=4.0, help="flow schedule shift — must match inference scheduler (default 4.0 for 1024px)") ap.add_argument("--grad_ckpt", action="store_true", help="enable gradient checkpointing to reduce VRAM at the cost of speed") ap.add_argument("--timestep_logit_std", type=float, default=1.0, help="std of logit-normal timestep sampling (higher = more uniform; 0 = pure midpoint)") ap.add_argument("--loss_weighting", default="sigma_sqrt", choices=["sigma_sqrt", "none"], help="sigma_sqrt upweights low-noise steps (identity/detail); none = uniform (Flux default: sigma_sqrt)") args = ap.parse_args() os.makedirs(args.out, exist_ok=True) device = torch.device(args.device) text_device = torch.device(args.text_device if args.text_device else args.device) # 1. Load model + inject LoRA print("Loading PixelDiTModel...") try: from peft import get_peft_model, LoraConfig except ImportError: print("peft not installed — run: pip install peft") sys.exit(1) from diffusers.pipelines.pixeldit import PixelDiTModel model = PixelDiTModel.from_pretrained(args.model, subfolder="transformer") lora_cfg = LoraConfig( r = args.lora_r, lora_alpha = args.lora_alpha, target_modules = ["qkv_x", "qkv_y", "proj_x", "proj_y"], lora_dropout = 0.05, bias = "none", ) model = get_peft_model(model, lora_cfg) model.print_trainable_parameters() if args.grad_ckpt: model.enable_input_require_grads() model.gradient_checkpointing_enable() print("[+] Gradient checkpointing enabled") model = model.to(device).train() print(f"[devices] PixelDiT={device} Gemma={text_device if args.train_text_encoder else 'off'}") text_encoder = None tokenizer = None # null embedding for CFG dropout — zeros null_emb = torch.zeros(1, 300, 2304, device=device) null_mask = torch.zeros(1, 300, device=device) # 2. Dataset + loader dataset = LoraDataset(args.data) img_size = dataset.img_size loader = DataLoader(dataset, batch_size=args.batch, shuffle=True, num_workers=2, pin_memory=True, drop_last=True) print(f"Dataset: {len(dataset)} samples img_size={img_size} batch={args.batch} accum={args.accum} steps/epoch={len(loader)}") # 2b. Text conditioning path learnable_embs = None opt_groups = [{"params": [p for p in model.parameters() if p.requires_grad]}] if args.train_text_encoder: if dataset.meta.get("encoder") not in (None, "gemma"): raise RuntimeError("--train_text_encoder requires a Gemma precompute cache") if dataset.captions is None: raise RuntimeError( "--train_text_encoder requires captions.json. " "Re-run scripts/precompute_lora_data.py with this patched version." ) print("Loading Gemma text encoder + LoRA...") from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_model) tokenizer.padding_side = "right" text_encoder = ( AutoModelForCausalLM.from_pretrained( args.text_encoder_model, torch_dtype=torch.bfloat16 ) .get_decoder() ) if hasattr(text_encoder, "config"): text_encoder.config.use_cache = False text_lora_cfg = LoraConfig( r=args.text_lora_r, lora_alpha=args.text_lora_alpha, target_modules=[m.strip() for m in args.text_lora_targets.split(",") if m.strip()], lora_dropout=0.05, bias="none", ) text_encoder = get_peft_model(text_encoder, text_lora_cfg) if args.grad_ckpt and hasattr(text_encoder, "gradient_checkpointing_enable"): text_encoder.gradient_checkpointing_enable() text_encoder.enable_input_require_grads() text_encoder = text_encoder.to(text_device).train() text_encoder.print_trainable_parameters() opt_groups.append({ "params": [p for p in text_encoder.parameters() if p.requires_grad], "lr": args.text_lr if args.text_lr is not None else args.lr * 0.25, }) else: print("Loading embeddings as learnable parameters...") all_embs = torch.from_numpy(dataset.embs[:].astype(np.float32)) learnable_embs = LearnableEmbeddings(all_embs).to(device) print(f" {learnable_embs.embs.numel()} embedding params") opt_groups.append({"params": learnable_embs.parameters(), "lr": args.lr * 0.1}) # 3. Optimizer opt = torch.optim.AdamW(opt_groups, lr=args.lr, weight_decay=1e-2) total_steps = args.epochs * len(loader) // args.accum sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps, eta_min=args.lr * 0.1) sigmas_cpu, alphas_cpu = _build_flow_schedule(args.flow_shift) alphas = alphas_cpu.to(device) sigmas = sigmas_cpu.to(device) print(f"Flow shift: {args.flow_shift}") best_loss = float("inf") step = 0 def save_lora(path, loss): os.makedirs(path, exist_ok=True) transformer_dir = os.path.join(path, "transformer") model.save_pretrained(transformer_dir) if text_encoder is not None: text_encoder.save_pretrained(os.path.join(path, "text_encoder")) if learnable_embs is not None: torch.save(learnable_embs.state_dict(), os.path.join(path, "learnable_embs.pt")) meta = { "data_dir": args.data, "model": args.model, "img_size": img_size, "epochs": args.epochs, "batch": args.batch, "accum": args.accum, "lr": args.lr, "lora_r": args.lora_r, "lora_alpha": args.lora_alpha, "cfg_drop": args.cfg_drop, "flow_shift": args.flow_shift, "timestep_logit_std": args.timestep_logit_std, "loss_weighting": args.loss_weighting, "loss": loss, "has_learnable_embeddings": learnable_embs is not None, "has_text_encoder_lora": text_encoder is not None, "text_encoder_model": args.text_encoder_model if text_encoder is not None else None, "text_device": str(text_device) if text_encoder is not None else None, "text_lora_r": args.text_lora_r if text_encoder is not None else None, "text_lora_alpha": args.text_lora_alpha if text_encoder is not None else None, "text_lora_targets": args.text_lora_targets if text_encoder is not None else None, "text_lr": args.text_lr if args.text_lr is not None else args.lr * 0.25, "precompute": dataset.meta, } with open(os.path.join(path, "training_meta.json"), "w", encoding="utf-8") as f: json.dump(meta, f, indent=2) for epoch in range(args.epochs): total_loss, n = 0.0, 0 bar = tqdm(loader, desc=f"epoch {epoch+1}/{args.epochs}") opt.zero_grad() for i, (indices, imgs, _masks) in enumerate(bar): indices = indices.to(device) # [B] batch indices imgs = imgs.to(device) # [B, 3, H, H] B = imgs.shape[0] if text_encoder is not None: batch_captions = [dataset.captions[int(idx)] for idx in indices.detach().cpu().tolist()] with torch.autocast(device_type=text_device.type, dtype=torch.bfloat16): embs = encode_gemma_batch(text_encoder, tokenizer, batch_captions, text_device) embs = embs.to(device) else: embs = learnable_embs(indices) # CFG dropout: replace some embeddings with the null (empty) embedding if args.cfg_drop > 0: drop = (torch.rand(B, device=device) < args.cfg_drop).view(B, 1, 1) embs = torch.where(drop, null_emb.expand(B, -1, -1).to(embs.dtype), embs) # Logit-normal timestep sampling — concentrates gradient signal # around mid-noise where the model does the most meaningful work. noise = torch.randn_like(imgs) u = torch.sigmoid(torch.randn(B, device=device) * args.timestep_logit_std) t = (u * _T).long().clamp(0, _T - 1) x_t = q_sample(imgs, t, noise, alphas, sigmas) target = noise - imgs # velocity: direction from data → noise with torch.autocast(device_type="cuda", dtype=torch.bfloat16): pred = model(x_t.bfloat16(), t, embs.bfloat16()) # per-sample MSE, then apply sigma weighting before reducing loss_per = F.mse_loss(pred.float(), target, reduction="none").mean(dim=(1, 2, 3)) # sigma_sqrt: weight = 1/sigma² — upweights low-noise steps where # identity and fine detail are learned; matches Flux Dev training. if args.loss_weighting == "sigma_sqrt": sig = sigmas[t].clamp(min=1e-3) # [B] w = (1.0 / sig ** 2).to(loss_per.device) loss = (w * loss_per).mean() else: loss = loss_per.mean() (loss / args.accum).backward() total_loss += loss.item() n += 1 if (i + 1) % args.accum == 0: all_params = [p for p in model.parameters() if p.requires_grad] if text_encoder is not None: all_params += [p for p in text_encoder.parameters() if p.requires_grad] if learnable_embs is not None: all_params += list(learnable_embs.parameters()) nn.utils.clip_grad_norm_(all_params, 1.0) opt.step() sched.step() opt.zero_grad() step += 1 bar.set_postfix(loss=f"{total_loss/n:.4f}", lr=f"{sched.get_last_lr()[0]:.2e}") # apply gradients from any tail batches that didn't fill a full accum window if n % args.accum != 0: all_params = [p for p in model.parameters() if p.requires_grad] if text_encoder is not None: all_params += [p for p in text_encoder.parameters() if p.requires_grad] if learnable_embs is not None: all_params += list(learnable_embs.parameters()) nn.utils.clip_grad_norm_(all_params, 1.0) opt.step() sched.step() opt.zero_grad() step += 1 avg = total_loss / n print(f" epoch {epoch+1} loss={avg:.4f}") is_last = (epoch + 1) == args.epochs is_save = (epoch + 1) % args.save_every == 0 is_best = avg < best_loss if is_best: best_loss = avg save_lora(os.path.join(args.out, "best"), avg) print(f" best saved → {args.out}/best/") if is_save or is_best or is_last: ckpt_name = f"ckpt_epoch_{epoch+1:03d}" save_lora(os.path.join(args.out, ckpt_name), avg) print(f" checkpoint → {args.out}/{ckpt_name}/") print(f"\nDone. Best loss: {best_loss:.4f}") print(f"LoRA checkpoints saved to {args.out}/") print("\nEach checkpoint contains:") print(" transformer/adapter_model.safetensors - PixelDiT transformer LoRA") if args.train_text_encoder: print(" text_encoder/adapter_model.safetensors - Gemma text-encoder LoRA") else: print(" learnable_embs.pt - fine-tuned per-image embeddings [N, 300, 2304]") print(" training_meta.json - metadata") print("\nTo use in inference:") print(f" python generate.py --prompt 'corrychase, your prompt here' --lora {args.out}/best --device {args.device}") if __name__ == "__main__": main()