scripts/train_stanno_on_embeddings.py

"""
Train a STANNO autoencoder on pre-computed CLIP embeddings.

Run generate_clip_embeddings.py first to produce the .npy file, then run
this script to train and save the STANNO.  The resulting .pkl file can be
loaded into ComfyUI via the STANNOLoad node.

Usage:
    python scripts/train_stanno_on_embeddings.py \
        --embeddings style_embeddings.npy \
        --out         stanno_clip_style.pkl \
        [--hidden     256]   \
        [--epochs     300]   \
        [--lr         0.005] \
        [--trainer    fixed]

The input/output dimension is inferred automatically from the embedding file
(typically 768 for SD 1.5 / ViT-L-14 CLIP).
"""

from __future__ import annotations
import argparse
import pickle
import sys
from pathlib import Path

import numpy as np


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Train a STANNO autoencoder on CLIP embeddings")
    p.add_argument("--embeddings", required=True, help="Path to .npy file of shape (N, dim)")
    p.add_argument("--out", required=True, help="Output .pkl path for the trained STANNO")
    p.add_argument("--hidden", type=int, default=256,
                   help="Hidden layer width (default 256 → [dim, 256, dim])")
    p.add_argument("--extra-hidden", type=int, default=0,
                   help="Add a second hidden layer of this width (0 = disabled)")
    p.add_argument("--epochs", type=int, default=300, help="Training epochs (default 300)")
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--lr", type=float, default=0.005, help="Learning rate")
    p.add_argument("--trainer", default="fixed",
                   choices=["fixed", "local_rule", "evolutionary"])
    return p.parse_args()


def main() -> None:
    args = parse_args()

    embeddings_path = Path(args.embeddings)
    if not embeddings_path.is_file():
        print(f"File not found: {embeddings_path}")
        sys.exit(1)

    embeddings = np.load(str(embeddings_path)).astype(np.float32)
    n, dim = embeddings.shape
    print(f"Loaded {n} embeddings of dim={dim} from {embeddings_path}")

    # Build layers list
    layers = [dim, args.hidden]
    if args.extra_hidden > 0:
        layers.append(args.extra_hidden)
    layers.append(dim)
    print(f"Architecture: {layers}")

    # Import STANNO (add repo root to path if needed)
    repo_root = str(Path(__file__).parent.parent)
    if repo_root not in sys.path:
        sys.path.insert(0, repo_root)

    from stanno.config.schema import STANNOConfig
    from stanno.core.stanno import STANNO

    config = STANNOConfig(
        layers=layers,
        trainer_type=args.trainer,
        learning_rate=args.lr,
    )
    stanno = STANNO(config)

    report_every = max(1, args.epochs // 10)

    def log_cb(epoch: int, loss: float) -> None:
        if (epoch + 1) % report_every == 0:
            print(f"  epoch {epoch + 1:5d} / {args.epochs}  loss={loss:.5f}")

    print(f"\nTraining STANNO ({args.trainer}) for {args.epochs} epochs …")
    stanno.fit(
        embeddings,
        embeddings,
        epochs=args.epochs,
        batch_size=args.batch_size,
        callback=log_cb,
    )

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(str(out_path), "wb") as f:
        pickle.dump(stanno, f)

    # Quick sanity check
    preds = stanno.predict(embeddings[:8])
    mse = float(np.mean((preds - embeddings[:8]) ** 2))
    print(f"\nFinal MSE on first 8 samples: {mse:.5f}")
    print(f"Saved trained STANNO → {out_path}")
    print("\nNext steps:")
    print("  1. Load in ComfyUI: STANNO Loader node → model_path =", out_path)
    print("  2. For Dream Conditioning: connect to 'STANNO Dream Conditioning' node")
    print("  3. For Dynamic LoRA:       connect to 'STANNO Dynamic LoRA' node")


if __name__ == "__main__":
    main()