File size: 3,041 Bytes
8f0d906 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """
Generate CLIP image embeddings for a folder of reference images.
These embeddings are then used to train a STANNO as a style autoencoder,
which can be loaded into the ComfyUI STANNODreamCond or STANNODynamicLoRA
nodes for conditioning/weight-patch injection.
Usage:
python scripts/generate_clip_embeddings.py \
--dir my_style_images/ \
--out style_embeddings.npy \
[--model ViT-L-14] [--pretrained openai]
Requirements:
pip install open-clip-torch Pillow
Outputs:
A .npy file of shape (N, 768) — one 768-dim CLIP embedding per image.
Compatible with SD 1.5 CLIP-L text encoder embedding space.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import numpy as np
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="Generate CLIP embeddings for a folder of images")
p.add_argument("--dir", required=True, help="Folder of input images (png, jpg, webp)")
p.add_argument("--out", required=True, help="Output .npy path")
p.add_argument("--model", default="ViT-L-14", help="OpenCLIP model name")
p.add_argument("--pretrained", default="openai", help="OpenCLIP pretrained weights")
p.add_argument("--batch", type=int, default=16, help="Batch size for encoding")
p.add_argument("--device", default="cuda", help="Device: cuda or cpu")
return p.parse_args()
def main() -> None:
args = parse_args()
try:
import torch
import open_clip
from PIL import Image
except ImportError as e:
print(f"Missing dependency: {e}")
print("Install with: pip install open-clip-torch Pillow")
sys.exit(1)
image_paths = sorted(
p for ext in ("*.png", "*.jpg", "*.jpeg", "*.webp")
for p in Path(args.dir).glob(ext)
)
if not image_paths:
print(f"No images found in {args.dir}")
sys.exit(1)
print(f"Found {len(image_paths)} images in {args.dir}")
model, _, preprocess = open_clip.create_model_and_transforms(
args.model, pretrained=args.pretrained
)
model.eval().to(args.device)
all_embeddings: list[np.ndarray] = []
for i in range(0, len(image_paths), args.batch):
batch_paths = image_paths[i : i + args.batch]
imgs = torch.stack(
[preprocess(Image.open(str(p)).convert("RGB")) for p in batch_paths]
).to(args.device)
with torch.no_grad():
feats = model.encode_image(imgs)
all_embeddings.append(feats.cpu().numpy())
print(f" Encoded {min(i + args.batch, len(image_paths))}/{len(image_paths)}")
embeddings = np.concatenate(all_embeddings, axis=0).astype(np.float32)
out_path = Path(args.out)
out_path.parent.mkdir(parents=True, exist_ok=True)
np.save(str(out_path), embeddings)
print(f"\nSaved {embeddings.shape} embeddings → {out_path}")
print(f"Use this file with train_stanno_on_embeddings.py or STANNOTrainImages (ComfyUI).")
if __name__ == "__main__":
main()
|