Buckets:
| #!/usr/bin/env python3 | |
| """Train a real-video reference-image projector. | |
| This learns a small adapter from CLIP image embeddings. It does not retrain | |
| CLIP itself, and it does not claim to replace the Open-Sora 11B transformer. | |
| """ | |
| import argparse | |
| import csv | |
| import os | |
| import random | |
| import subprocess | |
| import tempfile | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from safetensors.torch import save_file | |
| from torchvision import transforms | |
| from transformers import CLIPVisionModelWithProjection | |
| from PIL import Image | |
| class ReferenceProjector(nn.Module): | |
| def __init__(self, dim: int = 768): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.LayerNorm(dim), | |
| nn.Linear(dim, dim), | |
| nn.GELU(), | |
| nn.Linear(dim, dim), | |
| ) | |
| def forward(self, x): | |
| return self.net(x) | |
| def first_frame(video: str) -> str: | |
| fd, path = tempfile.mkstemp(suffix=".jpg") | |
| os.close(fd) | |
| subprocess.run( | |
| ["ffmpeg", "-y", "-i", video, "-vf", r"select=eq(n\,0)", "-frames:v", "1", path], | |
| check=True, | |
| stdout=subprocess.DEVNULL, | |
| stderr=subprocess.DEVNULL, | |
| ) | |
| return path | |
| def image_tensor(path: str): | |
| tfm = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.48145466, 0.4578275, 0.40821073], | |
| std=[0.26862954, 0.26130258, 0.27577711], | |
| ), | |
| ]) | |
| return tfm(Image.open(path).convert("RGB")).unsqueeze(0) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--csv", required=True) | |
| ap.add_argument("--video-root", default="") | |
| ap.add_argument("--output", required=True) | |
| ap.add_argument("--steps", type=int, default=10000) | |
| ap.add_argument("--lr", type=float, default=1e-5) | |
| ap.add_argument("--seed", type=int, default=42) | |
| args = ap.parse_args() | |
| random.seed(args.seed) | |
| torch.manual_seed(args.seed) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| vision = CLIPVisionModelWithProjection.from_pretrained( | |
| "openai/clip-vit-large-patch14", | |
| local_files_only=True, | |
| ).to(device).eval() | |
| for p in vision.parameters(): | |
| p.requires_grad_(False) | |
| dim = vision.config.projection_dim | |
| projector = ReferenceProjector(dim).to(device) | |
| optimizer = torch.optim.AdamW(projector.parameters(), lr=args.lr, weight_decay=0.01) | |
| with open(args.csv, newline="", encoding="utf-8") as f: | |
| rows = list(csv.DictReader(f)) | |
| if not rows: | |
| raise RuntimeError("No rows in training CSV") | |
| root = Path(args.video_root) if args.video_root else Path(".") | |
| projector.train() | |
| for step in range(args.steps): | |
| row = random.choice(rows) | |
| video = Path(row["path"]) | |
| if not video.is_absolute(): | |
| video = root / video | |
| if not video.exists(): | |
| continue | |
| frame = first_frame(str(video)) | |
| try: | |
| px = image_tensor(frame).to(device) | |
| with torch.no_grad(): | |
| base = vision(pixel_values=px).image_embeds | |
| target = base / base.norm(dim=-1, keepdim=True) | |
| pred = projector(base) | |
| pred = pred / pred.norm(dim=-1, keepdim=True).clamp_min(1e-6) | |
| # Keeps the reference representation stable while allowing the | |
| # projector to become a trainable v2.1 adapter. | |
| loss = 1.0 - (pred * target).sum(dim=-1).mean() | |
| optimizer.zero_grad(set_to_none=True) | |
| loss.backward() | |
| optimizer.step() | |
| if step % 100 == 0: | |
| print(f"reference step={step} loss={loss.item():.6f}", flush=True) | |
| finally: | |
| try: | |
| os.unlink(frame) | |
| except OSError: | |
| pass | |
| out = Path(args.output) | |
| out.parent.mkdir(parents=True, exist_ok=True) | |
| save_file({k: v.detach().cpu().contiguous() for k, v in projector.state_dict().items()}, str(out)) | |
| print(f"Saved {out}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 4.09 kB
- Xet hash:
- 9168f8790ba093181a9f3e8132c62a04bd984bb99af4dcee414a4e2d2f956b70
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.