Upload colab_checkpoint_eval.ipynb with huggingface_hub

colab_checkpoint_eval.ipynb

@@ -0,0 +1,268 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Cosmos-Predict2.5 Checkpoint Evaluation\n",
+    "\n",
+    "Evaluate fine-tuned checkpoints for iterative video generation on RoboCasa tasks.\n",
+    "\n",
+    "**Prerequisites:** Upload data to HuggingFace Hub first (see upload commands in the repo)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Setup Environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install cosmos-predict2.5\n",
+    "!git clone https://github.com/nvidia-cosmos/cosmos-predict2.5.git /content/cosmos-predict2.5\n",
+    "%cd /content/cosmos-predict2.5\n",
+    "!pip install -e '.[all]' 2>&1 | tail -5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Login to HuggingFace to download data\n",
+    "from huggingface_hub import login, hf_hub_download, snapshot_download\n",
+    "login()  # Enter your HF token when prompted"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "import os\n\n# === CONFIGURE THESE ===\nHF_REPO = \"bosungkim/cosmos-predict2.5-eval\"\nDATA_DIR = \"/content/data\"\nos.makedirs(DATA_DIR, exist_ok=True)\n\n# Download checkpoints\nprint(\"Downloading checkpoints...\")\nsnapshot_download(\n    repo_id=HF_REPO,\n    repo_type=\"model\",\n    local_dir=DATA_DIR,\n    allow_patterns=[\"checkpoints/*\", \"robocasa_data/*\", \"initial_images/*\"],\n)\nprint(\"Download complete!\")\n\n# Verify files\n!find {DATA_DIR} -name \"*.pt\" -exec ls -lh {} \\;\n!find {DATA_DIR} -name \"*.json\" -exec ls -lh {} \\;\n!find {DATA_DIR} -name \"*.jpg\" -exec ls -lh {} \\;"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# === CONFIGURE THESE ===\n",
+    "HF_REPO = \"bkim/cosmos-predict2.5-eval\"  # Change to your HF repo\n",
+    "DATA_DIR = \"/content/data\"\n",
+    "os.makedirs(DATA_DIR, exist_ok=True)\n",
+    "\n",
+    "# Download checkpoints\n",
+    "print(\"Downloading checkpoints...\")\n",
+    "snapshot_download(\n",
+    "    repo_id=HF_REPO,\n",
+    "    repo_type=\"model\",\n",
+    "    local_dir=DATA_DIR,\n",
+    "    allow_patterns=[\"checkpoints/*\", \"robocasa_data/*\", \"initial_images/*\"],\n",
+    ")\n",
+    "print(\"Download complete!\")\n",
+    "\n",
+    "# Verify files\n",
+    "!find {DATA_DIR} -name \"*.pt\" -exec ls -lh {} \\;\n",
+    "!find {DATA_DIR} -name \"*.json\" -exec ls -lh {} \\;\n",
+    "!find {DATA_DIR} -name \"*.jpg\" -exec ls -lh {} \\;"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ========== CONFIGURATION ==========\n",
+    "# Task selection: uncomment the task you want to evaluate\n",
+    "\n",
+    "# --- ArrangeVegetables ---\n",
+    "INITIAL_IMAGE = f\"{DATA_DIR}/initial_images/ArrangeVegetables__2024-05-11__demo_1_grid_first.jpg\"\n",
+    "TASK_NAME = \"ArrangeVegetables\"\n",
+    "\n",
+    "# --- PrepareCoffee ---\n",
+    "# INITIAL_IMAGE = f\"{DATA_DIR}/initial_images/PrepareCoffee__2024-05-07__demo_1_grid_first.jpg\"\n",
+    "# TASK_NAME = \"PrepareCoffee\"\n",
+    "\n",
+    "DATASET_PATH = f\"{DATA_DIR}/robocasa_data\"\n",
+    "JSON_FILE = f\"{DATA_DIR}/robocasa_data/robocasa_composite_to_atomic_decomposition_finegrained.json\"\n",
+    "RESOLUTION = \"432,768\"\n",
+    "GUIDANCE = 7\n",
+    "NUM_STEPS = 50\n",
+    "NUM_VIDEO_FRAMES = 77\n",
+    "\n",
+    "# Checkpoint paths\n",
+    "CHECKPOINT_ITERS = [2000, 2200, 2400, 2600]\n",
+    "CHECKPOINTS = [\n",
+    "    f\"{DATA_DIR}/checkpoints/iter_{it:09d}/model_ema_bf16.pt\"\n",
+    "    for it in CHECKPOINT_ITERS\n",
+    "]\n",
+    "\n",
+    "# Verify all files exist\n",
+    "for f in [INITIAL_IMAGE, JSON_FILE] + CHECKPOINTS:\n",
+    "    assert os.path.exists(f), f\"Missing: {f}\"\n",
+    "print(f\"Task: {TASK_NAME}\")\n",
+    "print(f\"Checkpoints: {len(CHECKPOINTS)}\")\n",
+    "print(\"All files verified!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Run Evaluation (Sequential)\n",
+    "\n",
+    "Colab has a single GPU, so we run checkpoints sequentially."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import subprocess\n",
+    "import time\n",
+    "\n",
+    "OUTPUT_BASE = \"/content/outputs/checkpoint_comparison\"\n",
+    "os.makedirs(OUTPUT_BASE, exist_ok=True)\n",
+    "\n",
+    "results = []\n",
+    "\n",
+    "for i, (ckpt_path, ckpt_iter) in enumerate(zip(CHECKPOINTS, CHECKPOINT_ITERS)):\n",
+    "    iter_str = f\"iter_{ckpt_iter:09d}\"\n",
+    "    output_dir = f\"{OUTPUT_BASE}/{iter_str}\"\n",
+    "    log_file = f\"{output_dir}/run.log\"\n",
+    "    os.makedirs(output_dir, exist_ok=True)\n",
+    "\n",
+    "    print(f\"\\n{'='*60}\")\n",
+    "    print(f\"[{i+1}/{len(CHECKPOINTS)}] Running {iter_str}\")\n",
+    "    print(f\"{'='*60}\")\n",
+    "\n",
+    "    cmd = [\n",
+    "        \"python\", \"examples/iterative_video_gen.py\",\n",
+    "        \"--output-dir\", output_dir,\n",
+    "        \"--dataset-path\", DATASET_PATH,\n",
+    "        \"--json-file\", JSON_FILE,\n",
+    "        \"--initial-image-path\", INITIAL_IMAGE,\n",
+    "        \"--checkpoint-path\", ckpt_path,\n",
+    "        \"--resolution\", RESOLUTION,\n",
+    "        \"--num-video-frames\", str(NUM_VIDEO_FRAMES),\n",
+    "        \"--task-name\", TASK_NAME,\n",
+    "        \"--guidance\", str(GUIDANCE),\n",
+    "        \"--num-steps\", str(NUM_STEPS),\n",
+    "    ]\n",
+    "\n",
+    "    start_time = time.time()\n",
+    "    with open(log_file, \"w\") as lf:\n",
+    "        proc = subprocess.run(\n",
+    "            cmd,\n",
+    "            stdout=lf,\n",
+    "            stderr=subprocess.STDOUT,\n",
+    "            env={**os.environ, \"PYTHONPATH\": \"/content/cosmos-predict2.5\"},\n",
+    "            cwd=\"/content/cosmos-predict2.5\",\n",
+    "        )\n",
+    "    elapsed = time.time() - start_time\n",
+    "\n",
+    "    status = \"SUCCESS\" if proc.returncode == 0 else \"FAILED\"\n",
+    "    results.append((iter_str, status, elapsed))\n",
+    "    print(f\"  {status} ({elapsed:.1f}s) - log: {log_file}\")\n",
+    "\n",
+    "print(f\"\\n{'='*60}\")\n",
+    "print(\"Summary:\")\n",
+    "for iter_str, status, elapsed in results:\n",
+    "    print(f\"  {iter_str}: {status} ({elapsed:.1f}s)\")\n",
+    "print(f\"{'='*60}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. View Results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "from IPython.display import display, HTML, Video\n",
+    "\n",
+    "# Find all generated videos\n",
+    "video_files = sorted(glob.glob(f\"{OUTPUT_BASE}/*/task_{TASK_NAME}/*.mp4\"))\n",
+    "print(f\"Found {len(video_files)} videos:\\n\")\n",
+    "for vf in video_files:\n",
+    "    print(f\"  {vf}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display videos inline (final_concatenated for each checkpoint)\n",
+    "for vf in video_files:\n",
+    "    if \"final_concatenated\" in vf:\n",
+    "        iter_name = vf.split(\"/checkpoint_comparison/\")[1].split(\"/\")[0]\n",
+    "        print(f\"\\n--- {iter_name} ---\")\n",
+    "        display(Video(vf, embed=True, width=768))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Download Results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Zip results for download\n",
+    "!cd /content && zip -r checkpoint_eval_results.zip outputs/checkpoint_comparison/\n",
+    "\n",
+    "from google.colab import files\n",
+    "files.download(\"/content/checkpoint_eval_results.zip\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "A100",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}