Enable torch compile for action viewer

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/app.py

@@ -773,7 +773,7 @@ def launch_action_viz(config: AppConfig) -> None:
                 dataset_split="full",
                 dataset_selector=loaded.adapter.dataset_selector,
                 dataset_kwargs=loaded.adapter.entry.dataset_kwargs,
-                use_torch_compile=False,
+                use_torch_compile=_env_enabled("ACTION_VIZ_USE_TORCH_COMPILE", default=True),
             )
             status_text.content = f"*Generating {generation_id}...*"
 

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/local_worker.py

@@ -38,6 +38,13 @@ ACTION_NORMALIZATION_STATS_PATH = Path(__file__).with_name("normalization_stats.
 PERSISTENT_RUNNER_ROOT = Path(os.environ.get("ACTION_VIZ_PERSISTENT_RUNNER_ROOT", "/tmp/action_viz_persistent_runner"))
 
 
+def _env_enabled(name: str, default: bool) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
 class LocalGenerationWorker:
     """Generation worker that runs cosmos-framework in this Viser process."""
 
@@ -814,7 +821,7 @@ class _PersistentFrameworkRunner:
                     checkpoint_path=checkpoint_path,
                     output_dir=runner_output_dir,
                     parallelism_preset="latency",
-                    use_torch_compile=False,
+                    use_torch_compile=_env_enabled("ACTION_VIZ_USE_TORCH_COMPILE", default=True),
                     guardrails=False,
                     max_num_seqs=1,
                 )

start.sh

@@ -23,6 +23,7 @@ export ACTION_VIZ_GENERATE_IMAGE_SIZE="${ACTION_VIZ_GENERATE_IMAGE_SIZE:-480}"
 export ACTION_VIZ_GENERATE_NUM_FRAMES="${ACTION_VIZ_GENERATE_NUM_FRAMES:-24}"
 export ACTION_VIZ_GENERATE_NUM_STEPS="${ACTION_VIZ_GENERATE_NUM_STEPS:-30}"
 export ACTION_VIZ_PERSISTENT_RUNNER="${ACTION_VIZ_PERSISTENT_RUNNER:-1}"
+export ACTION_VIZ_USE_TORCH_COMPILE="${ACTION_VIZ_USE_TORCH_COMPILE:-1}"
 export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
 
 if [[ -z "${ACTION_VIZ_CHECKPOINT_LOCAL_PATH:-}" && -z "${ACTION_VIZ_MODEL_MOUNT_PATH:-}" ]]; then
@@ -55,6 +56,7 @@ echo "ACTION_VIZ_GENERATE_IMAGE_SIZE=${ACTION_VIZ_GENERATE_IMAGE_SIZE}"
 echo "ACTION_VIZ_GENERATE_NUM_FRAMES=${ACTION_VIZ_GENERATE_NUM_FRAMES}"
 echo "ACTION_VIZ_GENERATE_NUM_STEPS=${ACTION_VIZ_GENERATE_NUM_STEPS}"
 echo "ACTION_VIZ_PERSISTENT_RUNNER=${ACTION_VIZ_PERSISTENT_RUNNER}"
+echo "ACTION_VIZ_USE_TORCH_COMPILE=${ACTION_VIZ_USE_TORCH_COMPILE}"
 echo "PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF}"
 echo "BRIDGE_LEROBOT_ROOT=${BRIDGE_LEROBOT_ROOT}"
 echo "COSMOS_VIEWER_ON_DEMAND_VIDEO=${COSMOS_VIEWER_ON_DEMAND_VIDEO}"