Spaces:
Running on L40S
Running on L40S
File size: 5,429 Bytes
9f818c5 | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | #!/usr/bin/env python
"""Run one local Action Viz generation without opening the browser UI."""
from __future__ import annotations
import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any
from cosmos_framework.data.vfm.action.action_viz.adapters import build_adapter, sample_action_to_numpy
from cosmos_framework.data.vfm.action.action_viz.local_worker import default_local_worker_from_env
from cosmos_framework.data.vfm.action.action_viz.state import (
GenerationRequest,
control_points_from_action,
make_generation_id,
)
from cosmos_framework.data.vfm.action.urdf_visualizer.viewer import _build_datasets, _create_dataset
def main() -> None:
_configure_cache_env()
args = _parse_args()
datasets = _build_datasets()
if args.dataset not in datasets:
raise ValueError(f"Unknown dataset {args.dataset!r}; expected one of {sorted(datasets)}")
entry = datasets[args.dataset]
sample_index = int(args.sample_index)
if sample_index < 0:
sample_index = int(entry.initial_index)
dataset = _create_dataset(entry, int(args.chunk_length))
sample = dataset[sample_index]
adapter = build_adapter(args.dataset, entry)
baked_action = sample_action_to_numpy(sample).astype("float32", copy=True)
generation_id = args.generation_id or make_generation_id()
generation_dir = Path(args.output_root) / generation_id
if generation_dir.exists():
shutil.rmtree(generation_dir)
request = GenerationRequest(
generation_id=generation_id,
model_mode=args.model_mode,
dataset=args.dataset,
sample_index=sample_index,
experiment_name="",
s3_checkpoint_dir=args.checkpoint,
checkpoint_cache_dir=None,
output_dir=str(generation_dir),
seed=int(args.seed),
num_steps=int(args.num_steps),
guidance=float(args.guidance),
control_points=control_points_from_action(baked_action, baked_action.shape[1]),
baked_action=baked_action.astype(float).tolist(),
prompt_description=_extract_prompt_description(sample.get("ai_caption", "")),
dataset_split="full",
dataset_selector=adapter.dataset_selector,
dataset_kwargs=entry.dataset_kwargs,
use_torch_compile=False,
)
progress: list[dict[str, Any]] = []
def _progress(percent: int, message: str) -> None:
progress.append({"percent": int(percent), "message": str(message)})
print(f"progress {percent:3d}% {message}", flush=True)
worker = default_local_worker_from_env()
try:
result = worker.run(request, progress_callback=_progress, queue_callback=lambda state: print(f"queue {state}"))
finally:
worker.close()
summary = {
"status": result.status,
"message": result.message,
"generation_id": result.generation_id,
"result_path": result.result_path,
"video_path": result.video_path,
"generated_action_path": result.generated_action_path,
"progress": progress,
}
print(json.dumps(summary, indent=2, sort_keys=True))
if result.status != "success":
raise RuntimeError(f"Generation failed: {result.message}")
if result.video_path is None or not Path(result.video_path).is_file():
raise FileNotFoundError(f"Generation did not produce a video file: {result.video_path}")
if args.model_mode == "policy" and (result.generated_action_path is None or not Path(result.generated_action_path).is_file()):
raise FileNotFoundError(f"Policy generation did not produce an action file: {result.generated_action_path}")
if not args.keep_output:
shutil.rmtree(generation_dir, ignore_errors=True)
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", default="bridge")
parser.add_argument("--sample-index", type=int, default=-1)
parser.add_argument("--chunk-length", type=int, default=16)
parser.add_argument("--model-mode", choices=("forward_dynamics", "policy"), default="forward_dynamics")
parser.add_argument("--checkpoint", default="nvidia/Cosmos3-Nano")
parser.add_argument("--output-root", default="/tmp/action_viz_generation_smoke")
parser.add_argument("--generation-id", default="")
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--num-steps", type=int, default=1)
parser.add_argument("--guidance", type=float, default=1.0)
parser.add_argument("--keep-output", action="store_true")
return parser.parse_args()
def _configure_cache_env() -> None:
app_data_root = Path(os.environ.get("ACTION_VIZ_APP_DATA_ROOT", "/app_data"))
hf_home = Path(os.environ.setdefault("HF_HOME", str(app_data_root / "huggingface")))
hf_hub_cache = Path(os.environ.setdefault("HF_HUB_CACHE", str(hf_home / "hub")))
hf_home.mkdir(parents=True, exist_ok=True)
hf_hub_cache.mkdir(parents=True, exist_ok=True)
def _extract_prompt_description(prompt: object) -> str:
if isinstance(prompt, dict):
prompt_obj = prompt
elif isinstance(prompt, str):
try:
prompt_obj = json.loads(prompt)
except json.JSONDecodeError:
return prompt
else:
return ""
value = prompt_obj.get("description", "")
return value if isinstance(value, str) else ""
if __name__ == "__main__":
main()
|