""" 2_train_lora.py — STAGE 2: train a FLUX.1-dev LoRA from the CSV ================================================================= Reads dataset.csv (image_name, caption) produced by Stage 1, locates each image in --images, and trains a LoRA adapter on FLUX.1-dev. Outputs a .safetensors file usable in Stage 3, diffusers, or ComfyUI. This is the CORRECTED trainer. Your main.py / mainv2.py / hugging.py were built on Stable Diffusion's UNet with a fake loss and never learned. The fixes: 1. FLUX is a transformer (DiT), not a UNet -> FluxPipeline + .transformer 2. LoRA actually attaches to real FLUX attn layers (not target_modules=[]) 3. Uses FLUX's own CLIP + T5 encoders (no 512->768 nn.Linear projection hack) 4. REAL flow-matching loss: predict velocity (noise - x0), not outputs.mean() 5. Saves .safetensors via save_lora_weights (not an unloadable .pth) Requires FLUX.1-dev (gated): run `huggingface-cli login` and accept the license. GPU: 24GB recommended (16GB needs extra offloading). Run: python 2_train_lora.py --csv dataset.csv --images ./train \ --trigger gr4v_style --out ./flux-lora-out --steps 1000 Install: pip install -U "diffusers>=0.32" transformers peft accelerate \ bitsandbytes safetensors pillow pandas tqdm numpy """ import argparse from pathlib import Path import numpy as np import pandas as pd import torch import torch.nn.functional as F from PIL import Image from tqdm import tqdm from diffusers import FluxPipeline from peft import LoraConfig, get_peft_model_state_dict def load_pipeline_with_lora(base_model: str, rank: int, weight_dtype, device): pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=weight_dtype) pipe.to(device) # Freeze the base model — we only train the LoRA. pipe.transformer.requires_grad_(False) pipe.vae.requires_grad_(False) pipe.text_encoder.requires_grad_(False) pipe.text_encoder_2.requires_grad_(False) # Attach LoRA to FLUX's real attention projections. lora_config = LoraConfig( r=rank, lora_alpha=rank, init_lora_weights="gaussian", target_modules=["to_q", "to_k", "to_v", "to_out.0"], ) pipe.transformer.add_adapter(lora_config) pipe.transformer.enable_gradient_checkpointing() trainable = [p for p in pipe.transformer.parameters() if p.requires_grad] n = sum(p.numel() for p in trainable) print(f"Trainable LoRA params: {n/1e6:.2f}M") if n == 0: raise SystemExit("LoRA attached to nothing — aborting. Check target_modules.") return pipe, trainable @torch.no_grad() def build_cache(pipe, rows, image_dir, resolution, device, weight_dtype): """Encode each image->latent and caption->embeddings once, up front.""" cache = [] vae_sf = pipe.vae.config.scaling_factor vae_shift = getattr(pipe.vae.config, "shift_factor", 0.0) image_dir = Path(image_dir) for _, row in tqdm(list(rows.iterrows()), desc="Encoding"): img_path = image_dir / str(row["image_name"]) if not img_path.exists(): print(f" missing image, skipping: {img_path}") continue caption = str(row["caption"]) # image -> latent img = Image.open(img_path).convert("RGB").resize((resolution, resolution)) arr = np.asarray(img, dtype=np.float32) / 127.5 - 1.0 # [-1,1] HxWxC px = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(device, weight_dtype) latent = pipe.vae.encode(px).latent_dist.sample() latent = (latent - vae_shift) * vae_sf # caption -> embeddings (FLUX's own CLIP + T5) prompt_embeds, pooled_embeds, text_ids = pipe.encode_prompt( prompt=caption, prompt_2=caption, max_sequence_length=512, device=device, num_images_per_prompt=1, ) cache.append({ "latent": latent.cpu(), "prompt_embeds": prompt_embeds.cpu(), "pooled_embeds": pooled_embeds.cpu(), "text_ids": text_ids.cpu(), }) if not cache: raise SystemExit("No images encoded — check that CSV names match files in --images.") print(f"Cached {len(cache)} samples.") return cache def train(pipe, trainable, cache, args, device, weight_dtype): transformer = pipe.transformer transformer.train() optimizer = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=1e-4) vae_scale = 2 ** (len(pipe.vae.config.block_out_channels) - 1) step = 0 pbar = tqdm(total=args.steps, desc="Training") while step < args.steps: for i, item in enumerate(cache): latents = item["latent"].to(device, weight_dtype) prompt_embeds = item["prompt_embeds"].to(device, weight_dtype) pooled = item["pooled_embeds"].to(device, weight_dtype) text_ids = item["text_ids"].to(device, weight_dtype) bsz = latents.shape[0] # --- flow matching --- noise = torch.randn_like(latents) t = torch.sigmoid(torch.randn(bsz, device=device)).to(weight_dtype) t_ = t.view(bsz, 1, 1, 1) noisy = (1.0 - t_) * latents + t_ * noise # x_t target = noise - latents # velocity h, w = noisy.shape[2], noisy.shape[3] packed = pipe._pack_latents(noisy, bsz, noisy.shape[1], h, w) img_ids = pipe._prepare_latent_image_ids(bsz, h // 2, w // 2, device, weight_dtype) guidance = torch.full((bsz,), args.guidance, device=device, dtype=weight_dtype) model_pred = transformer( hidden_states=packed, timestep=t, guidance=guidance, pooled_projections=pooled, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=img_ids, return_dict=False, )[0] model_pred = pipe._unpack_latents(model_pred, h * vae_scale, w * vae_scale, vae_scale) loss = F.mse_loss(model_pred.float(), target.float()) / args.grad_accum loss.backward() if (i + 1) % args.grad_accum == 0: torch.nn.utils.clip_grad_norm_(trainable, 1.0) optimizer.step() optimizer.zero_grad() step += 1 pbar.update(1) pbar.set_postfix(loss=f"{loss.item()*args.grad_accum:.4f}") if step % args.save_every == 0 or step >= args.steps: save_lora(transformer, args.out, args.trigger, step) if step >= args.steps: break pbar.close() def save_lora(transformer, out_dir, trigger, step): out = Path(out_dir) / f"{trigger}_step{step}" out.mkdir(parents=True, exist_ok=True) lora_sd = get_peft_model_state_dict(transformer) FluxPipeline.save_lora_weights(str(out), lora_sd, safe_serialization=True) print(f"\n saved LoRA -> {out}/") def main(): ap = argparse.ArgumentParser() ap.add_argument("--csv", default="dataset.csv") ap.add_argument("--images", default="./train") ap.add_argument("--out", default="./flux-lora-out") ap.add_argument("--trigger", default="gr4v_style") ap.add_argument("--base", default="black-forest-labs/FLUX.1-dev") ap.add_argument("--resolution", type=int, default=1024) ap.add_argument("--rank", type=int, default=16) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--steps", type=int, default=1000) ap.add_argument("--grad_accum", type=int, default=4) ap.add_argument("--save_every", type=int, default=250) ap.add_argument("--guidance", type=float, default=1.0) args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" weight_dtype = torch.bfloat16 torch.manual_seed(42) rows = pd.read_csv(args.csv) assert {"image_name", "caption"}.issubset(rows.columns), \ "CSV must have image_name,caption columns (run 1_build_dataset.py)" print("Loading FLUX + attaching LoRA...") pipe, trainable = load_pipeline_with_lora(args.base, args.rank, weight_dtype, device) print("Encoding dataset...") cache = build_cache(pipe, rows, args.images, args.resolution, device, weight_dtype) print("Training...") train(pipe, trainable, cache, args, device, weight_dtype) print(f"Done. LoRA written under {args.out}/") if __name__ == "__main__": main()