Create trainer.py

trainer.py

@@ -0,0 +1,413 @@
+# ─────────────────────────────────────────────────────────────────────────────
+# Cell 2 — JSON-conditioned sd15-flow-lune finetune  (225-token chunked CLIP)
+#
+# Adapts the Flux-teacher lune trainer to condition on structured JSON instead
+# of natural-language prompts. Warm-starts from the checkpoint-00018765 UNet,
+# trains on AbstractPhil/synthetic-object-relations-json, pushes to a NEW repo.
+#
+# Run TWICE — flip RUN — to produce the A/B couple:
+#   RUN = "prompt"  → conditions on json_prompt      → sd15-flow-lune-json-prompt
+#   RUN = "vit"     → conditions on vit_json_prompt  → sd15-flow-lune-json-vit
+#
+# CLIP CONDITIONING — 225-token kohya-style chunked encoding:
+#   Cell 1b found 94% of the ViT-derived JSON exceeds the 75-token single-CLIP-
+#   chunk budget (rich captions → 90-140 token JSON). So both runs encode the
+#   JSON as 3 chunks of 75 content tokens: tokenize to 227, split into 3×77
+#   (each re-wrapped with BOS/EOS), encode each through frozen CLIP, stitch the
+#   hidden states → a 227-length conditioning sequence. The UNet cross-attention
+#   takes any length, so no UNet change. Short prompt-JSON simply fills one
+#   chunk and pads the rest — both runs use identical machinery, so the only
+#   variable in the A/B is the conditioning text.
+#
+# Prereqs: Cell 1b must have added vit_json_prompt (done). This cell reads the
+# dataset's parquet shards directly, so it does NOT depend on the repo README
+# metadata. Full finetune, flow-matching timesteps untouched (full 0-1000),
+# base_lr 1e-5, save_optimizer=False (ships only the UNet).
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+#  RUN SELECTOR — flip this for each of the two finetunes
+# ═════════════════════════════════════════════════════════════════════════════
+
+RUN = "vit"        # "prompt" → json_prompt      "vit" → vit_json_prompt
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 1. INSTALL
+# ═════════════════════════════════════════════════════════════════════════════
+
+import subprocess, sys, os
+
+def _pip(*pkgs):
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *pkgs])
+
+print("Installing dependencies…")
+_pip("-U", "diffusers>=0.30", "transformers>=4.50", "accelerate>=1.0",
+     "datasets>=4.0", "huggingface_hub>=0.25", "tensorboard")
+print("  done.")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 2. AUTH
+# ═════════════════════════════════════════════════════════════════════════════
+
+def _load_hf_token():
+    if os.environ.get("HF_TOKEN"):
+        return "env"
+    try:
+        from google.colab import userdata
+        tok = userdata.get("HF_TOKEN")
+        if tok:
+            os.environ["HF_TOKEN"] = tok
+            os.environ["HUGGING_FACE_HUB_TOKEN"] = tok
+            return "secrets"
+    except Exception:
+        pass
+    return None
+
+print(f"HF token: {_load_hf_token() or 'not set'}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 3. IMPORTS
+# ═════════════════════════════════════════════════════════════════════════════
+
+import json
+import datetime
+from dataclasses import dataclass, asdict
+from tqdm.auto import tqdm
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader
+import datasets
+from diffusers import UNet2DConditionModel, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer
+from huggingface_hub import HfApi, create_repo
+
+assert torch.cuda.is_available(), "No GPU. Switch Colab runtime."
+print(f"GPU: {torch.cuda.get_device_name(0)}  "
+      f"({torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB)")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 4. CONFIG
+# ═════════════════════════════════════════════════════════════════════════════
+
+#                   condition column   length column        target repo
+_RUN_MAP = {
+    "prompt": ("json_prompt",     "json_token_len",     "AbstractPhil/sd15-flow-lune-json-prompt"),
+    "vit":    ("vit_json_prompt", "vit_json_token_len", "AbstractPhil/sd15-flow-lune-json-vit"),
+}
+assert RUN in _RUN_MAP, f"RUN must be one of {list(_RUN_MAP)}, got {RUN!r}"
+_COND_COL, _LEN_COL, _HF_REPO = _RUN_MAP[RUN]
+
+
+@dataclass
+class TrainConfig:
+    # --- set by the RUN selector ---
+    condition_column: str = _COND_COL
+    len_column: str       = _LEN_COL
+    hf_repo_id: str       = _HF_REPO
+
+    # --- sources ---
+    dataset_name: str   = "AbstractPhil/synthetic-object-relations-json"
+    sd_base: str        = "stable-diffusion-v1-5/stable-diffusion-v1-5"   # VAE + CLIP
+    unet_repo: str      = "AbstractPhil/sd15-flow-lune-flux"
+    unet_subfolder: str = "flux_t2_6_pose_t4_6_port_t1_4/checkpoint-00018765/unet"
+
+    output_dir: str      = "./outputs"
+    max_clip_tokens: int = 225           # 3 chunks × 75 — rows above this are filtered out
+
+    seed: int        = 42
+    batch_size: int  = 8
+    base_lr: float   = 1e-5              # modality shift — above the 2e-6 continued-run rate
+    shift: float     = 2.0
+    dropout: float   = 0.1               # CFG conditioning dropout
+
+    num_train_epochs: int    = 4         # prototype: short, just see if it behaves
+    warmup_epochs: int       = 1
+    checkpointing_steps: int = 1000
+    num_workers: int         = 0         # collate does VAE/CLIP on GPU — must be 0
+
+    vae_scale: float     = 0.18215
+    save_optimizer: bool = False         # prototype: ship only the UNet, skip the ~7 GB .pt
+    upload_to_hub: bool  = True
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 5. 225-TOKEN CHUNKED CLIP ENCODING  (kohya-style)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def encode_clip_225(prompts, tokenizer, text_encoder, device):
+    """Encode text as 3 × 75-token chunks → one 227-length hidden-state sequence.
+
+    A CLIP window is BOS + 75 content + EOS. We tokenize to 227 (= 3×75 + 2),
+    split into three 75-token bodies, re-wrap each with the original BOS/EOS
+    into a valid 77-token chunk, encode every chunk through the frozen CLIP,
+    then concatenate the hidden states — keeping only the first BOS and last
+    EOS. The UNet cross-attention consumes the 227-length result directly.
+    """
+    chunk_len = tokenizer.model_max_length          # 77
+    body_len  = chunk_len - 2                        # 75
+    n_chunks  = 3
+    max_tok   = body_len * n_chunks                  # 225
+
+    ids = tokenizer(prompts, padding="max_length", max_length=max_tok + 2,
+                    truncation=True, return_tensors="pt").input_ids        # [B, 227]
+    bos, eos = ids[:, :1], ids[:, -1:]
+    chunks = []
+    for k in range(n_chunks):
+        s = 1 + k * body_len
+        chunks.append(torch.cat([bos, ids[:, s:s + body_len], eos], dim=1))  # [B, 77]
+    ids = torch.stack(chunks, dim=1).reshape(-1, chunk_len).to(device)        # [B*3, 77]
+
+    hs = text_encoder(ids)[0]                                                 # [B*3, 77, 768]
+    hs = hs.reshape(len(prompts), n_chunks * chunk_len, -1)                    # [B, 231, 768]
+
+    out = [hs[:, :1]]                                                         # first BOS
+    for k in range(n_chunks):
+        s = k * chunk_len + 1
+        out.append(hs[:, s:s + body_len])                                     # 75 content tokens
+    out.append(hs[:, -1:])                                                    # last EOS
+    return torch.cat(out, dim=1)                                              # [B, 227, 768]
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 6. LOAD UNET — straight from the diffusers checkpoint folder
+# ═════════════════════════════════════════════════════════════════════════════
+
+def load_unet(config, device="cuda"):
+    print(f"\nLoading UNet from {config.unet_repo}/{config.unet_subfolder}…")
+    unet = UNet2DConditionModel.from_pretrained(
+        config.unet_repo, subfolder=config.unet_subfolder, torch_dtype=torch.float32)
+    print(f"  ✓ UNet loaded ({sum(p.numel() for p in unet.parameters()) / 1e6:.0f}M params)")
+    return unet.to(device)
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 7. TRAIN
+# ═════════════════════════════════════════════════════════════════════════════
+
+def train(config):
+    device = "cuda"
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.manual_seed(config.seed)
+    torch.cuda.manual_seed(config.seed)
+
+    date_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    real_output_dir = os.path.join(config.output_dir, f"{RUN}_{date_time}")
+    os.makedirs(real_output_dir, exist_ok=True)
+    t_writer = SummaryWriter(log_dir=real_output_dir, flush_secs=60)
+
+    hf_api = None
+    if config.upload_to_hub:
+        try:
+            hf_api = HfApi()
+            create_repo(repo_id=config.hf_repo_id, repo_type="model",
+                        exist_ok=True, private=False)
+            print(f"✓ HF repo ready: {config.hf_repo_id}")
+        except Exception as e:
+            print(f"⚠ hub upload disabled: {e}")
+            config.upload_to_hub = False
+
+    config_path = os.path.join(real_output_dir, "config.json")
+    with open(config_path, "w") as f:
+        json.dump({"run": RUN, **asdict(config)}, f, indent=2)
+    if config.upload_to_hub:
+        hf_api.upload_file(path_or_fileobj=config_path, path_in_repo="config.json",
+                           repo_id=config.hf_repo_id, repo_type="model")
+
+    # ── VAE + CLIP (frozen) ──────────────────────────────────────────────────
+    print(f"\nLoading VAE + CLIP from {config.sd_base}…")
+    vae = AutoencoderKL.from_pretrained(config.sd_base, subfolder="vae",
+                                        torch_dtype=torch.float32).to(device)
+    vae.requires_grad_(False); vae.eval()
+    tokenizer = CLIPTokenizer.from_pretrained(config.sd_base, subfolder="tokenizer")
+    text_encoder = CLIPTextModel.from_pretrained(
+        config.sd_base, subfolder="text_encoder", torch_dtype=torch.float32).to(device)
+    text_encoder.requires_grad_(False); text_encoder.eval()
+    print("✓ VAE + CLIP loaded (frozen)")
+    print(f"  conditioning: 3×{tokenizer.model_max_length - 2}-token chunks "
+          f"→ 227-length sequence (kohya-style)")
+
+    # ── dataset: load the parquet shards directly ────────────────────────────
+    # Cell 1b added vit_json_* columns to the parquet but its push left the repo
+    # README's dataset_info stale (11 cols vs 13 in the parquet), so
+    # load_dataset(repo) CastErrors casting parquet → README schema. The parquet
+    # builder reads the parquet's own (correct) schema and never sees the README.
+    print(f"\nLoading dataset {config.dataset_name} (parquet-direct)…")
+    _pq = sorted(f for f in HfApi().list_repo_files(config.dataset_name, repo_type="dataset")
+                 if f.endswith(".parquet"))
+    if not _pq:
+        raise RuntimeError(f"no parquet files found in {config.dataset_name}")
+    ds = datasets.load_dataset(
+        "parquet",
+        data_files={"train": [f"hf://datasets/{config.dataset_name}/{f}" for f in _pq]},
+        split="train")
+    ds = ds.cast_column("image", datasets.Image())
+    col, lcol = config.condition_column, config.len_column
+    if col not in ds.column_names or lcol not in ds.column_names:
+        raise RuntimeError(
+            f"dataset is missing '{col}'/'{lcol}'. Run Cell 1b then Cell 1c first.")
+    n_before = len(ds)
+    ds = ds.filter(lambda ex: isinstance(ex[col], str) and ex[col].strip() != ""
+                              and 0 < ex[lcol] <= config.max_clip_tokens)
+    ds = ds.select_columns(["id", "image", col])
+    print(f"✓ {len(ds)}/{n_before} rows usable for '{col}' (JSON ≤ {config.max_clip_tokens} tok)")
+    if len(ds) == 0:
+        raise RuntimeError(f"no usable rows for '{col}'.")
+
+    steps_per_epoch = len(ds) // config.batch_size
+    total_steps     = steps_per_epoch * config.num_train_epochs
+    warmup_steps    = max(steps_per_epoch * config.warmup_epochs, 1)
+    print(f"\nSchedule: {steps_per_epoch} steps/epoch × {config.num_train_epochs} "
+          f"epochs = {total_steps} steps  (warmup {warmup_steps})")
+
+    @torch.no_grad()
+    def collate_fn(examples):
+        """Encode images (VAE) and JSON conditioning text (chunked CLIP)."""
+        images, prompts, ids = [], [], []
+        for ex in examples:
+            img = ex["image"].convert("RGB")
+            img = torch.tensor(np.array(img)).permute(2, 0, 1).float() / 255.0
+            images.append(img * 2.0 - 1.0)              # [-1, 1]
+            prompts.append(ex[col])
+            ids.append(ex["id"])
+        images  = torch.stack(images).to(device)
+        latents = vae.encode(images).latent_dist.sample() * config.vae_scale
+        ehs = encode_clip_225(prompts, tokenizer, text_encoder, device)   # [B, 227, 768]
+        return latents.cpu(), ehs.cpu(), ids, prompts
+
+    train_loader = DataLoader(ds, batch_size=config.batch_size, shuffle=True,
+                              collate_fn=collate_fn, num_workers=config.num_workers,
+                              pin_memory=True)
+
+    # ── UNet + fresh optimizer ───────────────────────────────────────────────
+    unet = load_unet(config, device)
+    unet.requires_grad_(True); unet.train()
+    optimizer = torch.optim.AdamW(unet.parameters(), lr=config.base_lr,
+                                  betas=(0.9, 0.999), weight_decay=0.01, eps=1e-8)
+
+    def lr_scale(step):
+        return step / warmup_steps if step < warmup_steps else 1.0
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_scale)
+    print(f"✓ fresh AdamW, lr {config.base_lr:.1e}, linear warmup")
+
+    global_step = 0
+    train_logs  = {"train_step": [], "train_loss": [], "train_timestep": [],
+                   "trained_images": []}
+
+    def get_prediction(batch, log_to=None):
+        latents, ehs, ids, prompts = batch
+        latents = latents.to(dtype=torch.float32, device=device)
+        ehs     = ehs.to(dtype=torch.float32, device=device)
+        bsz = latents.shape[0]
+
+        # CFG conditioning dropout — zero some embeddings
+        drop = torch.rand(bsz, device=device) < config.dropout
+        ehs = ehs.clone(); ehs[drop] = 0
+
+        # shifted flow-matching timesteps, full 0–1000 range (no masking)
+        sigmas = torch.rand(bsz, device=device)
+        sigmas = (config.shift * sigmas) / (1 + (config.shift - 1) * sigmas)
+        timesteps = sigmas * 1000
+        sigmas = sigmas[:, None, None, None]
+
+        noise  = torch.randn_like(latents)
+        noisy  = noise * sigmas + latents * (1 - sigmas)
+        target = noise - latents                              # velocity
+
+        pred = unet(noisy, timesteps, ehs, return_dict=False)[0]
+        loss = F.mse_loss(pred, target, reduction="none")
+        loss = loss.mean(dim=list(range(1, len(loss.shape))))
+
+        if log_to is not None:
+            for i in range(bsz):
+                log_to["train_step"].append(global_step)
+                log_to["train_loss"].append(loss[i].item())
+                log_to["train_timestep"].append(timesteps[i].item())
+                log_to["trained_images"].append(
+                    {"step": global_step, "id": ids[i], "prompt": prompts[i]})
+        return loss.mean()
+
+    def plot_logs(d):
+        plt.figure(figsize=(10, 6))
+        plt.scatter(d["train_timestep"], d["train_loss"], s=3, c=d["train_step"],
+                    marker=".", cmap="cool")
+        plt.xlabel("timestep"); plt.ylabel("loss"); plt.yscale("log")
+        plt.colorbar(label="step")
+
+    def save_checkpoint(step, epoch):
+        ckpt = os.path.join(real_output_dir, f"checkpoint-{step:08}")
+        os.makedirs(ckpt, exist_ok=True)
+        unet.save_pretrained(os.path.join(ckpt, "unet"), safe_serialization=True)
+        meta = {"step": step, "epoch": epoch, "run": RUN, "condition_column": col,
+                "trained_images": train_logs["trained_images"]}
+        meta_path = os.path.join(ckpt, "trained_images.json")
+        with open(meta_path, "w") as f:
+            json.dump(meta, f, indent=2)
+        if config.save_optimizer:
+            torch.save({"cfg": asdict(config), "student": unet.state_dict(),
+                        "opt": optimizer.state_dict(), "scheduler": scheduler.state_dict(),
+                        "gstep": step, "epoch": epoch},
+                       os.path.join(ckpt, f"sd15_flow_lune_json_e{epoch}_s{step}.pt"))
+        print(f"✓ checkpoint {step} (epoch {epoch})")
+        if config.upload_to_hub and hf_api is not None:
+            try:
+                hf_api.upload_folder(folder_path=os.path.join(ckpt, "unet"),
+                                     path_in_repo=f"checkpoint-{step:08}/unet",
+                                     repo_id=config.hf_repo_id, repo_type="model")
+                hf_api.upload_file(path_or_fileobj=meta_path,
+                                   path_in_repo=f"checkpoint-{step:08}/trained_images.json",
+                                   repo_id=config.hf_repo_id, repo_type="model")
+                print(f"  ✓ uploaded to {config.hf_repo_id}")
+            except Exception as e:
+                print(f"  ⚠ upload failed: {e}")
+
+    # ── training loop ────────────────────────────────────────────────────────
+    print(f"\nStarting training — RUN='{RUN}', conditioning on '{col}'\n")
+    progress = tqdm(total=total_steps)
+    for epoch in range(config.num_train_epochs):
+        for batch in train_loader:
+            if global_step >= total_steps:
+                break
+            loss = get_prediction(batch, log_to=train_logs)
+            t_writer.add_scalar("train/loss", loss.item(), global_step)
+            t_writer.add_scalar("train/lr", scheduler.get_last_lr()[0], global_step)
+
+            loss.backward()
+            gn = torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)
+            t_writer.add_scalar("train/grad_norm", gn.item(), global_step)
+            optimizer.step(); scheduler.step(); optimizer.zero_grad()
+
+            progress.update(1)
+            progress.set_postfix({"epoch": epoch, "loss": f"{loss.item():.4f}",
+                                  "lr": f"{scheduler.get_last_lr()[0]:.2e}"})
+            global_step += 1
+
+            if global_step % 100 == 0:
+                plot_logs(train_logs)
+                t_writer.add_figure("train_loss", plt.gcf(), global_step)
+                plt.close()
+            if global_step % config.checkpointing_steps == 0:
+                save_checkpoint(global_step, epoch)
+
+        save_checkpoint(global_step, epoch)        # end-of-epoch
+
+    progress.close()
+    print(f"\n✅ Training complete — RUN='{RUN}'")
+    print(f"   → https://huggingface.co/{config.hf_repo_id}")
+    other = "vit" if RUN == "prompt" else "prompt"
+    print(f"   Now set RUN = \"{other}\" and re-run this cell for the other half "
+          f"of the couple.")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 8. RUN
+# ═════════════════════════════════════════════════════════════════════════════
+
+train(TrainConfig())