from __future__ import annotations import argparse import asyncio import json from pathlib import Path from lagunavision.backbones.factory import available_backbones, build_backbone from lagunavision.data.hf_materialize import DEFAULT_HF_DATASETS, materialize_hf_dataset, parse_dataset_requests from lagunavision.data.llava import materialize_llava_hf, materialize_llava_json from lagunavision.data.general_recipe import ( PILOT_300K_RECIPE, materialize_general_recipe, recipe_summary, ) from lagunavision.data.manifest import load_manifest from lagunavision.data.sources import DATASET_SOURCES from lagunavision.defaults import DEFAULT_VISION_TOWER from lagunavision.data.spatial_ocr import generate_spatial_ocr_manifest from lagunavision.eval.ablation import AblationConfig, run_ablation from lagunavision.eval.demo_set import generate_demo_eval from lagunavision.eval.run_eval import run_text_eval from lagunavision.eval.scene_probe import generate_scene_probe from lagunavision.eval.scene_probe import generate_scene_dataset from lagunavision.eval.score_eval import score_answer from lagunavision.eval.visual_overfit import write_visual_overfit_dataset from lagunavision.eval.web_probe import generate_web_probe from lagunavision.hub import resolve_checkpoint_reference from lagunavision.model import LagunaVisionTextPipeline from lagunavision.positions.normalized_2d import Normalized2DPositionEncoder from lagunavision.projectors.features import stack_visual_features from lagunavision.tiling.anyres import AnyResTiler from lagunavision.train.visual_bridge import ( VisualBridgeTrainConfig, _feature_cache_path, _save_feature_tensor, _tiles_for_item, train_visual_bridge, ) from lagunavision.visual_pipeline import LagunaVisionImagePipeline, VisualProjectorSpec def _add_backbone_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--backbone", choices=available_backbones(), default="laguna") parser.add_argument("--model-id", default="") def main() -> None: parser = argparse.ArgumentParser(prog="laguna-vision") subcommands = parser.add_subparsers(dest="command", required=True) datasets = subcommands.add_parser("datasets") datasets.set_defaults(func=_datasets) tile = subcommands.add_parser("tile") tile.add_argument("--width", type=int, required=True) tile.add_argument("--height", type=int, required=True) tile.set_defaults(func=_tile) spatial = subcommands.add_parser("spatial-ocr") spatial.add_argument("--output-dir", type=Path, required=True) spatial.add_argument("--count", type=int, default=200) spatial.set_defaults(func=_spatial_ocr) ask = subcommands.add_parser("ask") ask.add_argument("question") ask.add_argument("--context", default="") _add_backbone_args(ask) ask.add_argument("--device", default="auto") ask.set_defaults(func=_ask) eval_cmd = subcommands.add_parser("eval-text") eval_cmd.add_argument("--manifest", type=Path, required=True) eval_cmd.add_argument("--output", type=Path, required=True) _add_backbone_args(eval_cmd) eval_cmd.add_argument("--device", default="auto") eval_cmd.add_argument("--ocr-context", action="store_true") eval_cmd.set_defaults(func=_eval_text) demo_eval = subcommands.add_parser("demo-eval") demo_eval.add_argument("--output-dir", type=Path, required=True) demo_eval.set_defaults(func=_demo_eval) web_probe = subcommands.add_parser("web-probe") web_probe.add_argument("--output-dir", type=Path, required=True) web_probe.add_argument("--limit", type=int, default=5) web_probe.add_argument("--width", type=int, default=1920) web_probe.add_argument("--height", type=int, default=1080) web_probe.set_defaults(func=_web_probe) scene_probe = subcommands.add_parser("scene-probe") scene_probe.add_argument("--output-dir", type=Path, required=True) scene_probe.add_argument("--limit", type=int, default=5) scene_probe.set_defaults(func=_scene_probe) scene_dataset = subcommands.add_parser("scene-dataset") scene_dataset.add_argument("--output-dir", type=Path, required=True) scene_dataset.add_argument("--train-count", type=int, default=40) scene_dataset.add_argument("--eval-count", type=int, default=10) scene_dataset.set_defaults(func=_scene_dataset) visual_overfit_dataset = subcommands.add_parser("visual-overfit-dataset") visual_overfit_dataset.add_argument("--output-dir", type=Path, required=True) visual_overfit_dataset.add_argument("--train-count", type=int, default=32) visual_overfit_dataset.add_argument("--eval-count", type=int, default=16) visual_overfit_dataset.set_defaults(func=_visual_overfit_dataset) hf_dataset = subcommands.add_parser("hf-materialize") hf_dataset.add_argument("--output-dir", type=Path, required=True) hf_dataset.add_argument("--train-count", type=int, default=100) hf_dataset.add_argument("--eval-count", type=int, default=20) hf_dataset.add_argument("--dataset", action="append", default=[]) hf_dataset.set_defaults(func=_hf_materialize) llava_dataset = subcommands.add_parser("llava-materialize") llava_dataset.add_argument("--output-dir", type=Path, required=True) llava_dataset.add_argument("--source-json", type=Path) llava_dataset.add_argument("--image-root", type=Path, action="append", default=[]) llava_dataset.add_argument("--dataset", default="") llava_dataset.add_argument("--split", default="train") llava_dataset.add_argument("--limit", type=int, default=0) llava_dataset.add_argument("--eval-count", type=int, default=0) llava_dataset.add_argument("--image-mode", choices=("reference", "copy", "symlink"), default="reference") llava_dataset.set_defaults(func=_llava_materialize) general_dataset = subcommands.add_parser("general-materialize") general_dataset.add_argument("--output-dir", type=Path, required=True) general_dataset.add_argument("--recipe", default=PILOT_300K_RECIPE) general_dataset.add_argument("--sample-per-source", type=int, default=0) general_dataset.add_argument("--train-budget", type=int, default=0) general_dataset.add_argument("--download-assets", action="store_true") general_dataset.add_argument("--coco-train2017-root", type=Path) general_dataset.add_argument("--llava-pretrain-image-root", type=Path) general_dataset.add_argument("--seed", type=int, default=7) general_dataset.add_argument("--dry-run", action="store_true") general_dataset.set_defaults(func=_general_materialize) cache_features = subcommands.add_parser("cache-visual-features") cache_features.add_argument("--manifest", type=Path, required=True) cache_features.add_argument("--output-dir", type=Path, required=True) cache_features.add_argument("--encoder", choices=("pil", "hf"), default="hf") cache_features.add_argument("--encoder-id", default=DEFAULT_VISION_TOWER) cache_features.add_argument("--patch-px", type=int, default=32) cache_features.add_argument("--max-tiles", type=int, default=9) cache_features.add_argument("--device", default="cuda") cache_features.add_argument("--shard-index", type=int, default=0) cache_features.add_argument("--num-shards", type=int, default=1) cache_features.add_argument("--overwrite", action="store_true") cache_features.set_defaults(func=_cache_visual_features) score_cmd = subcommands.add_parser("score") score_cmd.add_argument("--manifest", type=Path, required=True) score_cmd.add_argument("--answers", type=Path, required=True) score_cmd.set_defaults(func=_score) train_bridge = subcommands.add_parser("train-visual-bridge") train_bridge.add_argument("--manifest", type=Path, required=True) train_bridge.add_argument("--output-dir", type=Path, required=True) _add_backbone_args(train_bridge) train_bridge.add_argument("--epochs", type=int, default=1) train_bridge.add_argument("--max-items", type=int, default=0) train_bridge.add_argument("--lr", type=float, default=1e-3) train_bridge.add_argument("--eval-manifest", type=Path) train_bridge.add_argument("--encoder", choices=("pil", "hf"), default="pil") train_bridge.add_argument("--encoder-id", default="") train_bridge.add_argument("--max-tiles", type=int, default=4) train_bridge.add_argument("--projector", choices=("mlp", "resampler"), default="mlp") train_bridge.add_argument("--visual-tokens", type=int, default=64) train_bridge.add_argument("--batch-size", type=int, default=4) train_bridge.add_argument("--grad-accum", type=int, default=1) train_bridge.add_argument("--warmup-ratio", type=float, default=0.03) train_bridge.add_argument("--num-workers", type=int, default=0) train_bridge.add_argument("--save-every", type=int, default=0) train_bridge.add_argument("--feature-cache-dir", type=Path) train_bridge.add_argument("--init-checkpoint") train_bridge.add_argument("--init-lora-dir", type=Path) train_bridge.add_argument("--no-grad-checkpointing", dest="grad_checkpointing", action="store_false") train_bridge.add_argument("--lora-rank", type=int, default=0) train_bridge.add_argument("--lora-alpha", type=int, default=16) train_bridge.add_argument("--lora-dropout", type=float, default=0.05) train_bridge.add_argument("--lora-targets", nargs="*", default=[]) train_bridge.add_argument("--device", default="auto") train_bridge.add_argument("--vision-device", default="auto") train_bridge.set_defaults(func=_train_visual_bridge, grad_checkpointing=True) ask_image = subcommands.add_parser("ask-image") ask_image.add_argument("--image", type=Path, required=True) ask_image.add_argument("--checkpoint", required=True) ask_image.add_argument("--question", required=True) ask_image.add_argument("--backbone", choices=available_backbones(), default="") ask_image.add_argument("--model-id", default="") ask_image.add_argument("--device", default="auto") ask_image.add_argument("--vision-device", default="auto") ask_image.set_defaults(func=_ask_image) ablation = subcommands.add_parser("eval-ablation") ablation.add_argument("--manifest", type=Path, required=True) ablation.add_argument("--checkpoint", required=True) ablation.add_argument("--output", type=Path, required=True) ablation.add_argument("--backbone", choices=available_backbones(), default="") ablation.add_argument("--model-id", default="") ablation.add_argument("--limit", type=int, default=0) ablation.add_argument("--threshold", type=float, default=0.15) ablation.add_argument("--device", default="auto") ablation.add_argument("--vision-device", default="auto") ablation.set_defaults(func=_eval_ablation) args = parser.parse_args() args.func(args) def _datasets(_: argparse.Namespace) -> None: rows = [ { "id": source.id, "stage": source.stage.value, "required": source.required, "use": source.use, } for source in DATASET_SOURCES ] print(json.dumps(rows, indent=2, sort_keys=True)) def _tile(args: argparse.Namespace) -> None: tiler = AnyResTiler() positions = Normalized2DPositionEncoder() tiles = tiler.tiles_for_size(args.width, args.height) rows = [ { "id": tile.id, "crop": [tile.crop.left, tile.crop.top, tile.crop.right, tile.crop.bottom], "position": positions.encode_tile(tile).values, } for tile in tiles ] print(json.dumps(rows, indent=2)) def _spatial_ocr(args: argparse.Namespace) -> None: examples = generate_spatial_ocr_manifest(args.output_dir, args.count) print(json.dumps({"examples": len(examples), "manifest": str(args.output_dir / "manifest.jsonl")})) def _ask(args: argparse.Namespace) -> None: async def run() -> None: backbone = build_backbone(args.backbone, args.model_id, device=args.device) pipeline = LagunaVisionTextPipeline(backbone) print(await pipeline.answer(args.question, extracted_context=args.context)) asyncio.run(run()) def _eval_text(args: argparse.Namespace) -> None: async def run() -> None: passed, total = await run_text_eval( args.manifest, build_backbone(args.backbone, args.model_id, device=args.device), args.output, use_ocr_context=args.ocr_context, ) print(json.dumps({"passed": passed, "total": total})) asyncio.run(run()) def _demo_eval(args: argparse.Namespace) -> None: manifest = generate_demo_eval(args.output_dir) print(json.dumps({"manifest": str(manifest), "items": 15})) def _web_probe(args: argparse.Namespace) -> None: async def run() -> None: manifest = await generate_web_probe(args.output_dir, args.limit, width=args.width, height=args.height) print(json.dumps({"manifest": str(manifest), "items": args.limit})) asyncio.run(run()) def _scene_probe(args: argparse.Namespace) -> None: manifest = generate_scene_probe(args.output_dir, args.limit) print(json.dumps({"manifest": str(manifest), "items": args.limit})) def _scene_dataset(args: argparse.Namespace) -> None: train_manifest, eval_manifest = generate_scene_dataset( args.output_dir, train_count=args.train_count, eval_count=args.eval_count, ) print( json.dumps( { "train_manifest": str(train_manifest), "eval_manifest": str(eval_manifest), "train_items": args.train_count, "eval_items": args.eval_count, } ) ) def _visual_overfit_dataset(args: argparse.Namespace) -> None: write_visual_overfit_dataset(args.output_dir, train_count=args.train_count, eval_count=args.eval_count) print( json.dumps( { "train_manifest": str(args.output_dir / "train.jsonl"), "eval_manifest": str(args.output_dir / "eval.jsonl"), "wrong_manifest": str(args.output_dir / "wrong.jsonl"), "blank_manifest": str(args.output_dir / "blank.jsonl"), "train_items": args.train_count, "eval_items": args.eval_count, } ) ) def _hf_materialize(args: argparse.Namespace) -> None: datasets = parse_dataset_requests(args.dataset) if args.dataset else None train_manifest, eval_manifest = materialize_hf_dataset( args.output_dir, train_count=args.train_count, eval_count=args.eval_count, datasets=datasets or DEFAULT_HF_DATASETS, ) print( json.dumps( { "train_manifest": str(train_manifest), "eval_manifest": str(eval_manifest), "train_items": args.train_count, "eval_items": args.eval_count, } ) ) def _llava_materialize(args: argparse.Namespace) -> None: if bool(args.source_json) == bool(args.dataset): raise SystemExit("provide exactly one of --source-json or --dataset") if args.source_json: result = materialize_llava_json( args.source_json, args.output_dir, image_roots=args.image_root, limit=args.limit, eval_count=args.eval_count, image_mode=args.image_mode, ) else: result = materialize_llava_hf( args.dataset, args.output_dir, split=args.split, limit=args.limit, eval_count=args.eval_count, ) print( json.dumps( { "train_manifest": str(result.train_manifest), "eval_manifest": str(result.eval_manifest) if result.eval_manifest else "", "train_items": result.train_count, "eval_items": result.eval_count, } ) ) def _general_materialize(args: argparse.Namespace) -> None: if args.dry_run: print(json.dumps(recipe_summary(args.recipe, sample_per_source=args.sample_per_source, train_budget=args.train_budget), indent=2)) return result = materialize_general_recipe( args.output_dir, recipe=args.recipe, sample_per_source=args.sample_per_source, train_budget=args.train_budget, download_assets=args.download_assets, coco_train2017_root=args.coco_train2017_root, llava_pretrain_image_root=args.llava_pretrain_image_root, seed=args.seed, ) print( json.dumps( { "recipe_file": str(result.recipe_file), "alignment_train_manifest": str(result.alignment_train_manifest), "alignment_eval_manifest": str(result.alignment_eval_manifest), "instruction_train_manifest": str(result.instruction_train_manifest), "instruction_eval_manifest": str(result.instruction_eval_manifest), "wrong_manifest": str(result.wrong_manifest), "blank_manifest": str(result.blank_manifest), "counts": dict(result.counts), } ) ) def _cache_visual_features(args: argparse.Namespace) -> None: if args.num_shards < 1: raise SystemExit("--num-shards must be >= 1") if not 0 <= args.shard_index < args.num_shards: raise SystemExit("--shard-index must be between 0 and --num-shards - 1") async def run() -> None: from lagunavision.encoders.factory import build_vision_encoder items = load_manifest(args.manifest) args.output_dir.mkdir(parents=True, exist_ok=True) encoder = build_vision_encoder(args.encoder, args.encoder_id, args.patch_px, args.device) tiler = AnyResTiler(max_tiles=args.max_tiles) positioner = Normalized2DPositionEncoder() written = 0 skipped = 0 for index, item in enumerate(items): if index % args.num_shards != args.shard_index: continue feature_path = _feature_cache_path(args.output_dir, item) if feature_path.exists() and not args.overwrite: skipped += 1 continue tiles = _tiles_for_item(tiler, item) encoded = await encoder.encode(item.image, tiles) positions = positioner.encode_tiles(tiles) _save_feature_tensor(stack_visual_features(encoded, positions, "cpu"), feature_path) written += 1 print(json.dumps({"written": written, "skipped": skipped, "shard_index": args.shard_index, "num_shards": args.num_shards})) asyncio.run(run()) def _score(args: argparse.Namespace) -> None: items = {item.id: item for item in load_manifest(args.manifest)} total = 0 passed = 0 with args.answers.open("r", encoding="utf-8") as handle: for line in handle: row = json.loads(line) item_id = str(row["id"]) score = score_answer(items[item_id], str(row["answer"])) total += 1 passed += int(score.passed) print(json.dumps({"passed": passed, "total": total})) def _train_visual_bridge(args: argparse.Namespace) -> None: async def run() -> None: checkpoint = await train_visual_bridge( VisualBridgeTrainConfig( manifest=args.manifest, output_dir=args.output_dir, backbone=args.backbone, model_id=args.model_id, epochs=args.epochs, max_items=args.max_items, lr=args.lr, eval_manifest=args.eval_manifest, encoder=args.encoder, encoder_id=args.encoder_id, max_tiles=args.max_tiles, projector=args.projector, visual_tokens=args.visual_tokens, batch_size=args.batch_size, grad_accum=args.grad_accum, warmup_ratio=args.warmup_ratio, num_workers=args.num_workers, save_every=args.save_every, feature_cache_dir=args.feature_cache_dir, init_checkpoint=resolve_checkpoint_reference(args.init_checkpoint) if args.init_checkpoint else None, init_lora_dir=args.init_lora_dir, grad_checkpointing=args.grad_checkpointing, lora_rank=args.lora_rank, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, lora_targets=tuple(args.lora_targets), device=args.device, vision_device=args.vision_device, ) ) print(json.dumps({"checkpoint": str(checkpoint)})) asyncio.run(run()) def _load_projector_spec(checkpoint: Path) -> tuple[VisualProjectorSpec, dict]: spec_row = json.loads((checkpoint.parent / "projector_spec.json").read_text(encoding="utf-8")) spec = VisualProjectorSpec( input_dim=int(spec_row["input_dim"]), embedding_dim=int(spec_row["embedding_dim"]), hidden_dim=int(spec_row["hidden_dim"]), projector=spec_row.get("projector", "mlp"), visual_tokens=int(spec_row.get("visual_tokens", 64)), encoder=spec_row.get("encoder", "pil"), encoder_id=spec_row.get("encoder_id", ""), max_tiles=int(spec_row.get("max_tiles", 4)), patch_px=int(spec_row.get("patch_px", 32)), ) return spec, spec_row def _lora_dir(checkpoint: Path, spec_row: dict) -> Path | None: directory = checkpoint.parent / "lora" return directory if int(spec_row.get("lora_rank", 0)) > 0 or directory.exists() else None def _ask_image(args: argparse.Namespace) -> None: async def run() -> None: checkpoint = resolve_checkpoint_reference(args.checkpoint) spec, spec_row = _load_projector_spec(checkpoint) pipeline = await LagunaVisionImagePipeline.from_checkpoint( checkpoint=checkpoint, spec=spec, backbone_name=args.backbone or spec_row.get("backbone", "laguna"), model_id=args.model_id or spec_row["model_id"], device=args.device, vision_device=args.vision_device, lora_dir=_lora_dir(checkpoint, spec_row), ) print(await pipeline.answer_image(args.image, args.question)) asyncio.run(run()) def _eval_ablation(args: argparse.Namespace) -> None: async def run() -> None: checkpoint = resolve_checkpoint_reference(args.checkpoint) spec, spec_row = _load_projector_spec(checkpoint) summary = await run_ablation( AblationConfig( manifest=args.manifest, checkpoint=checkpoint, output=args.output, spec=spec, backbone_name=args.backbone or spec_row.get("backbone", "laguna"), model_id=args.model_id or spec_row["model_id"], device=args.device, vision_device=args.vision_device, limit=args.limit, capability_threshold=args.threshold, lora_dir=_lora_dir(checkpoint, spec_row), ) ) print(json.dumps(summary, indent=2)) asyncio.run(run()) if __name__ == "__main__": main()