{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cbf9cf63",
   "metadata": {},
   "source": [
    "# Presto Free Inference (Colab)\n",
    "\n",
    "Self-contained notebook: installs dependencies, downloads Presto + weights, runs encoder inference.\n",
    "\n",
    "**No Google Cloud / Vertex AI / authentication required.**\n",
    "\n",
    "Runtime: `Runtime → Run all` (GPU optional; CPU works).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d77a6974",
   "metadata": {},
   "source": [
    "## 1. Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "326a2b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install only what Presto needs (no gcloud / Vertex)\n",
    "%pip install -q --upgrade pip\n",
    "%pip install -q torch einops numpy\n",
    "\n",
    "import torch\n",
    "print(\"torch:\", torch.__version__)\n",
    "print(\"cuda available:\", torch.cuda.is_available())\n",
    "print(\"device will be:\", \"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bd6668b",
   "metadata": {},
   "source": [
    "## 2. Download Presto source + pretrained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f077f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import urllib.request\n",
    "\n",
    "ROOT = Path(\"/content/presto_run\") if Path(\"/content\").exists() else Path.cwd() / \"presto_run\"\n",
    "ROOT.mkdir(parents=True, exist_ok=True)\n",
    "print(\"ROOT:\", ROOT)\n",
    "\n",
    "files = {\n",
    "    \"single_file_presto.py\": \"https://raw.githubusercontent.com/nasaharvest/presto/main/single_file_presto.py\",\n",
    "    \"default_model.pt\": \"https://github.com/nasaharvest/presto/raw/main/data/default_model.pt\",\n",
    "}\n",
    "\n",
    "for name, url in files.items():\n",
    "    dest = ROOT / name\n",
    "    if dest.exists() and dest.stat().st_size > 0:\n",
    "        print(f\"already present: {dest} ({dest.stat().st_size} bytes)\")\n",
    "        continue\n",
    "    print(f\"downloading {name} ...\")\n",
    "    urllib.request.urlretrieve(url, dest)\n",
    "    print(f\"saved {dest} ({dest.stat().st_size} bytes)\")\n",
    "\n",
    "assert (ROOT / \"single_file_presto.py\").exists()\n",
    "assert (ROOT / \"default_model.pt\").stat().st_size > 1_000_000\n",
    "print(\"OK: files ready\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1787c1ba",
   "metadata": {},
   "source": [
    "## 3. Load Presto encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b90d6eb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "\n",
    "ROOT = Path(\"/content/presto_run\") if Path(\"/content\").exists() else Path.cwd() / \"presto_run\"\n",
    "sys.path.insert(0, str(ROOT))\n",
    "\n",
    "from single_file_presto import NUM_DYNAMIC_WORLD_CLASSES, Presto\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = Presto.construct()\n",
    "try:\n",
    "    state = torch.load(ROOT / \"default_model.pt\", map_location=device, weights_only=True)\n",
    "except TypeError:\n",
    "    state = torch.load(ROOT / \"default_model.pt\", map_location=device)\n",
    "model.load_state_dict(state)\n",
    "model.eval()\n",
    "model.to(device)\n",
    "\n",
    "print(\"OK: model loaded on\", device)\n",
    "print(\"dynamic_world ignore class:\", NUM_DYNAMIC_WORLD_CLASSES)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b47e71",
   "metadata": {},
   "source": [
    "## 4. Run inference\n",
    "\n",
    "`x` shape must be `[batch, timesteps, 17]` bands:\n",
    "\n",
    "`VV, VH, B2, B3, B4, B5, B6, B7, B8, B8A, B11, B12, temp, precip, elev, slope, NDVI`\n",
    "\n",
    "Use Dynamic World class `9` when DW is unavailable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f332b9bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import torch\n",
    "\n",
    "NUM_BANDS = 17\n",
    "batch, timesteps = 1, 1\n",
    "\n",
    "x = torch.zeros(batch, timesteps, NUM_BANDS, device=device)\n",
    "# sample optical-ish values\n",
    "x[0, 0, 2:12] = 0.1\n",
    "x[0, 0, 8:10] = 0.2\n",
    "x[0, 0, 16] = 0.3\n",
    "\n",
    "dynamic_world = torch.full((batch, timesteps), NUM_DYNAMIC_WORLD_CLASSES, dtype=torch.long, device=device)\n",
    "latlons = torch.tensor([[42.96, -85.67]], dtype=torch.float32, device=device)\n",
    "mask = torch.zeros_like(x)\n",
    "month = torch.tensor([4], dtype=torch.long, device=device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    embeddings = model.encoder(\n",
    "        x=x,\n",
    "        dynamic_world=dynamic_world,\n",
    "        latlons=latlons,\n",
    "        mask=mask,\n",
    "        month=month,\n",
    "        eval_task=True,\n",
    "    )\n",
    "\n",
    "print(\"OK: embeddings shape:\", tuple(embeddings.shape))\n",
    "print(\"first 8 dims:\", embeddings[0, :8].detach().cpu().tolist())\n",
    "\n",
    "result = {\n",
    "    \"embeddings\": embeddings.detach().cpu().tolist(),\n",
    "    \"shape\": list(embeddings.shape),\n",
    "    \"device\": str(device),\n",
    "    \"dynamic_world_ignore_class\": NUM_DYNAMIC_WORLD_CLASSES,\n",
    "}\n",
    "\n",
    "out_path = ROOT / \"presto_embeddings.json\"\n",
    "out_path.write_text(json.dumps(result, indent=2))\n",
    "print(\"saved:\", out_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8ff9186",
   "metadata": {},
   "source": [
    "## 5. (Optional) Predict from your own JSON payload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "918166f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = {\n",
    "    \"x\": [[[0.0, 0.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.1, 0.1, 0.0, 0.0, 0.0, 0.0, 0.3]]],\n",
    "    \"dynamic_world\": [[9]],\n",
    "    \"latlons\": [[42.96, -85.67]],\n",
    "    \"mask\": [[[0] * 17]],\n",
    "    \"month\": [4],\n",
    "}\n",
    "\n",
    "x = torch.tensor(payload[\"x\"], dtype=torch.float32, device=device)\n",
    "dw = torch.tensor(payload[\"dynamic_world\"], dtype=torch.long, device=device)\n",
    "ll = torch.tensor(payload[\"latlons\"], dtype=torch.float32, device=device)\n",
    "mk = torch.tensor(payload[\"mask\"], dtype=torch.float32, device=device)\n",
    "mo = torch.tensor(payload[\"month\"], dtype=torch.long, device=device)\n",
    "\n",
    "assert x.ndim == 3 and x.shape[-1] == 17, f\"x must be [B,T,17], got {tuple(x.shape)}\"\n",
    "\n",
    "with torch.no_grad():\n",
    "    emb = model.encoder(x=x, dynamic_world=dw, latlons=ll, mask=mk, month=mo, eval_task=True)\n",
    "\n",
    "print(\"OK:\", tuple(emb.shape))\n",
    "print(emb[0, :5].detach().cpu().tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd574fa7",
   "metadata": {},
   "source": [
    "## Done\n",
    "\n",
    "If every cell printed `OK`, Presto is running correctly.\n",
    "\n",
    "Working free hosted UI (backup): https://huggingface.co/spaces/morta00/presto-free\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}