| """Precompute frozen SigLIP2 features for Dense UV Parser fixed views. |
| |
| The direct trainer has no render augmentation, so recomputing the same frozen |
| vision tower output in every epoch wastes most of the semantic-encoder time. |
| The optional spatial cache stores full-width FP16 patch features after removing |
| only deterministic letterbox padding. Trainable projections remain in the |
| downstream model. |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| from contextlib import nullcontext |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
|
|
| TOOLKIT_ROOT = Path(__file__).resolve().parents[1] |
| WORKSPACE_ROOT = TOOLKIT_ROOT.parent |
| if str(WORKSPACE_ROOT) not in sys.path: |
| sys.path.insert(0, str(WORKSPACE_ROOT)) |
|
|
| from SkingToolkit.renderer import DifferentiableRenderer |
| from SkingToolkit.dense_uv_parser.semantic_backbone import ( |
| SigLIP2VisionBackbone, |
| ) |
| from SkingToolkit.dense_uv_parser.semantic_cache import ( |
| SIGLIP_CACHE_VERSION, |
| SigLIPGlobalCache, |
| ) |
| from SkingToolkit.dense_uv_parser.runtime import get_device |
| from SkingToolkit.dense_uv_parser.skin_dataset import SkinUVDataset |
| from SkingToolkit.dense_uv_parser.utils import parse_views |
|
|
| try: |
| from tqdm import tqdm |
| except ImportError: |
| tqdm = None |
|
|
|
|
| def autocast_context(device, precision): |
| if precision == "no" or device.type == "cpu": |
| return nullcontext() |
| dtype = torch.float16 if precision == "fp16" else torch.bfloat16 |
| return torch.autocast(device_type=device.type, dtype=dtype) |
|
|
|
|
| def cache_is_reusable( |
| cache_dir, |
| dataset, |
| views, |
| siglip_model, |
| require_spatial=False, |
| ): |
| try: |
| cache = SigLIPGlobalCache( |
| cache_dir, |
| expected_views=views, |
| expected_model=siglip_model, |
| expected_data_dir=dataset.data_dir, |
| require_spatial=require_spatial, |
| ) |
| except (FileNotFoundError, ValueError, KeyError, json.JSONDecodeError): |
| return False |
| return list(cache.filename_to_index) == [path.name for path in dataset.skin_paths] |
|
|
|
|
| def build_arg_parser(): |
| parser = argparse.ArgumentParser( |
| description="Cache fixed-view SigLIP2 global and optional spatial features." |
| ) |
| parser.add_argument("--data_dir", default="../skins") |
| parser.add_argument("--cache_dir", required=True) |
| parser.add_argument("--mappings_dir", required=True) |
| parser.add_argument( |
| "--views", |
| default="walk_front_both_layer_ortho,walk_back_both_layer_ortho", |
| ) |
| parser.add_argument("--siglip_model", default="google/siglip2-base-patch16-224") |
| parser.add_argument("--siglip_local_files_only", action="store_true") |
| parser.add_argument("--max_samples", type=int, default=None) |
| parser.add_argument("--batch_size", type=int, default=32) |
| parser.add_argument("--num_workers", type=int, default=16) |
| parser.add_argument("--prefetch_factor", type=int, default=4) |
| parser.add_argument("--device", default="auto") |
| parser.add_argument("--mixed_precision", choices=["no", "fp16", "bf16"], default="bf16") |
| parser.add_argument( |
| "--spatial", |
| action="store_true", |
| help="Also cache raw cropped patch features as FP16 mmap data.", |
| ) |
| parser.add_argument("--force", action="store_true") |
| return parser |
|
|
|
|
| @torch.inference_mode() |
| def main(): |
| args = build_arg_parser().parse_args() |
| if args.batch_size < 1: |
| raise ValueError("--batch_size must be positive.") |
| views = parse_views(args.views) |
| if len(views) < 2: |
| raise ValueError("At least front and back cache views are required.") |
| dataset = SkinUVDataset(args.data_dir, max_samples=args.max_samples) |
| cache_dir = Path(args.cache_dir) |
| if not args.force and cache_is_reusable( |
| cache_dir, |
| dataset, |
| views, |
| args.siglip_model, |
| require_spatial=args.spatial, |
| ): |
| cache_kind = "global + spatial" if args.spatial else "global" |
| print(f"Reusing complete SigLIP {cache_kind} cache: {cache_dir}") |
| return |
|
|
| device = get_device(args.device) |
| renderer = DifferentiableRenderer(mappings_dir=args.mappings_dir).to(device).eval() |
| missing_views = [view for view in views if view not in renderer.views] |
| if missing_views: |
| raise ValueError(f"Renderer cache mappings are missing views {missing_views}.") |
| backbone = SigLIP2VisionBackbone( |
| model_name=args.siglip_model, |
| token_channels=128, |
| local_files_only=args.siglip_local_files_only, |
| ).to(device).eval() |
|
|
| loader_kwargs = { |
| "batch_size": args.batch_size, |
| "shuffle": False, |
| "num_workers": args.num_workers, |
| "pin_memory": device.type == "cuda", |
| } |
| if args.num_workers > 0: |
| loader_kwargs["persistent_workers"] = True |
| loader_kwargs["prefetch_factor"] = args.prefetch_factor |
| loader = DataLoader(dataset, **loader_kwargs) |
|
|
| cache_dir.mkdir(parents=True, exist_ok=True) |
| temporary_embeddings = cache_dir / ".embeddings.npy.tmp" |
| temporary_spatial_embeddings = cache_dir / ".spatial_embeddings.npy.tmp" |
| temporary_metadata = cache_dir / ".metadata.json.tmp" |
| embeddings = np.lib.format.open_memmap( |
| temporary_embeddings, |
| mode="w+", |
| dtype=np.float16, |
| shape=(len(dataset), len(views), backbone.raw_feature_dim), |
| ) |
| spatial_embeddings = None |
| spatial_shape = None |
| iterator = tqdm(loader, desc="cache SigLIP2", leave=False) if tqdm else loader |
| offset = 0 |
| for batch in iterator: |
| uv = batch["uv"].to(device, non_blocking=True) |
| renders = torch.stack( |
| [renderer.forward_view(uv, view) for view in views], dim=1 |
| ) |
| batch_size, view_count, _, height, width = renders.shape |
| images = renders[:, :, :3].reshape(batch_size * view_count, 3, height, width) |
| if images.is_cuda: |
| images = images.contiguous(memory_format=torch.channels_last) |
| with autocast_context(device, args.mixed_precision): |
| encoded = ( |
| backbone.encode_dense(images) |
| if args.spatial |
| else backbone.encode_global(images) |
| ) |
| raw_global = encoded["raw_global"].reshape( |
| batch_size, view_count, backbone.raw_feature_dim |
| ) |
| embeddings[offset : offset + batch_size] = ( |
| raw_global.detach().to(device="cpu", dtype=torch.float16).numpy() |
| ) |
| if args.spatial: |
| raw_spatial = encoded["raw_spatial"] |
| current_spatial_shape = tuple(raw_spatial.shape[-2:]) |
| if spatial_embeddings is None: |
| spatial_shape = current_spatial_shape |
| spatial_embeddings = np.lib.format.open_memmap( |
| temporary_spatial_embeddings, |
| mode="w+", |
| dtype=np.float16, |
| shape=( |
| len(dataset), |
| len(views), |
| backbone.raw_spatial_feature_dim, |
| *spatial_shape, |
| ), |
| ) |
| spatial_size_gib = ( |
| np.prod(spatial_embeddings.shape) |
| * np.dtype(np.float16).itemsize |
| / (1024**3) |
| ) |
| print( |
| "SigLIP2 spatial cache: " |
| f"shape={spatial_embeddings.shape}, " |
| f"estimated_size={spatial_size_gib:.1f} GiB." |
| ) |
| elif current_spatial_shape != spatial_shape: |
| raise ValueError( |
| "SigLIP spatial shape changed within a fixed-view cache: " |
| f"expected {spatial_shape}, got {current_spatial_shape}." |
| ) |
| raw_spatial = raw_spatial.reshape( |
| batch_size, |
| view_count, |
| backbone.raw_spatial_feature_dim, |
| *spatial_shape, |
| ) |
| spatial_embeddings[offset : offset + batch_size] = ( |
| raw_spatial.detach().to(device="cpu", dtype=torch.float16).numpy() |
| ) |
| offset += batch_size |
| embeddings.flush() |
| del embeddings |
| if spatial_embeddings is not None: |
| spatial_embeddings.flush() |
| del spatial_embeddings |
| if offset != len(dataset): |
| raise RuntimeError(f"Cached {offset} samples, expected {len(dataset)}.") |
|
|
| metadata = { |
| "version": SIGLIP_CACHE_VERSION, |
| "data_dir": str(Path(args.data_dir).resolve()), |
| "filenames": [path.name for path in dataset.skin_paths], |
| "views": views, |
| "siglip_model": args.siglip_model, |
| "feature_dim": backbone.raw_feature_dim, |
| "dtype": "float16", |
| } |
| if args.spatial: |
| if spatial_shape is None: |
| raise RuntimeError("Spatial SigLIP cache was requested but no features were written.") |
| metadata.update( |
| { |
| "spatial_feature_dim": backbone.raw_spatial_feature_dim, |
| "spatial_height": spatial_shape[0], |
| "spatial_width": spatial_shape[1], |
| "spatial_dtype": "float16", |
| "spatial_layout": "N,V,C,H,W", |
| "spatial_letterbox_padding_removed": True, |
| } |
| ) |
| temporary_metadata.write_text( |
| json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8" |
| ) |
| os.replace(temporary_embeddings, cache_dir / "embeddings.npy") |
| if args.spatial: |
| os.replace( |
| temporary_spatial_embeddings, |
| cache_dir / "spatial_embeddings.npy", |
| ) |
| os.replace(temporary_metadata, cache_dir / "metadata.json") |
| cache_kind = "global + spatial" if args.spatial else "global" |
| print( |
| f"Cached {cache_kind} features for {len(dataset)} skins x " |
| f"{len(views)} views to {cache_dir} " |
| f"({backbone.raw_feature_dim} features/view)." |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
