{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd12eb72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "allocated MB: 0.0\n",
      "reserved  MB: 0.0\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "module 'dpm.model' has no attribute 'parameters'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[8]\u001b[39m\u001b[32m, line 12\u001b[39m\n\u001b[32m      9\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[33mallocated MB:\u001b[39m\u001b[33m\"\u001b[39m, torch.cuda.memory_allocated()/\u001b[32m1024\u001b[39m**\u001b[32m2\u001b[39m)\n\u001b[32m     10\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[33mreserved  MB:\u001b[39m\u001b[33m\"\u001b[39m, torch.cuda.memory_reserved()/\u001b[32m1024\u001b[39m**\u001b[32m2\u001b[39m)\n\u001b[32m---> \u001b[39m\u001b[32m12\u001b[39m params_bytes = \u001b[38;5;28msum\u001b[39m(p.numel()*p.element_size() \u001b[38;5;28;01mfor\u001b[39;00m p \u001b[38;5;129;01min\u001b[39;00m \u001b[43mmodel\u001b[49m\u001b[43m.\u001b[49m\u001b[43mparameters\u001b[49m())\n\u001b[32m     13\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[33mparams MB:\u001b[39m\u001b[33m\"\u001b[39m, params_bytes/\u001b[32m1024\u001b[39m**\u001b[32m2\u001b[39m)\n\u001b[32m     15\u001b[39m dtype_bytes = {}\n",
      "\u001b[31mAttributeError\u001b[39m: module 'dpm.model' has no attribute 'parameters'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "# Load model using the project's loader (may download if missing)\n",
    "from gradio_demo import load_cfg_from_cli, load_model\n",
    "cfg = load_cfg_from_cli()\n",
    "vdpm_model = load_model(cfg)\n",
    "vdpm_model.eval()\n",
    "\n",
    "# Lightweight diagnostics: VRAM and dtype breakdown\n",
    "import gc\n",
    "from collections import Counter\n",
    "gc.collect()\n",
    "torch.cuda.synchronize()\n",
    "torch.cuda.empty_cache()\n",
    "torch.cuda.synchronize()\n",
    "print('allocated MB:', torch.cuda.memory_allocated()/1024**2)\n",
    "print('reserved  MB:', torch.cuda.memory_reserved()/1024**2)\n",
    "# Ensure we have an nn.Module instance (some imports may shadow the name 'model')\n",
    "import types\n",
    "if isinstance(vdpm_model, types.ModuleType):\n",
    "    print('Warning: loader returned a module, not an instance. Trying to instantiate via load_model again.')\n",
    "    vdpm_model = load_model(cfg)\n",
    "assert hasattr(vdpm_model, 'parameters'), f'Loaded object has no parameters: {type(vdpm_model)}'\n",
    "params_bytes = sum(p.numel()*p.element_size() for p in vdpm_model.parameters())\n",
    "print('params MB:', params_bytes/1024**2)\n",
    "dtype_bytes = {}\n",
    "for p in vdpm_model.parameters():\n",
    "    k = str(p.dtype)\n",
    "    dtype_bytes[k] = dtype_bytes.get(k,0) + p.numel()*p.element_size()\n",
    "print('param bytes by dtype (MB):', {k:v/1024**2 for k,v in dtype_bytes.items()})\n",
    "print('buffer dtype counts:', Counter(getattr(b,'dtype',None) for b in vdpm_model.buffers()))\n",
    "\n",
    "# Small aggregator inspect (one small dummy batch)\n",
    "B, S, H, W = 1, min(4, 1 if not torch.cuda.is_available() else 4), 518, 518\n",
    "dummy = torch.rand(B, S, 3, H, W, device='cuda' if torch.cuda.is_available() else 'cpu')\n",
    "agg, patch_start = vdpm_model.aggregator(dummy)\n",
    "for k, t in sorted(agg.items()):\n",
    "    print(k, t.device, t.dtype, tuple(t.shape), f\"{t.numel()*t.element_size()/1024**2:.2f} MB\")\n",
    "\n",
    "# Warmup + one timed inference\n",
    "import time\n",
    "for _ in range(2):\n",
    "    _ = vdpm_model.inference([{'img': torch.rand(1,3,518,518, device='cuda' if torch.cuda.is_available() else 'cpu')}])\n",
    "torch.cuda.synchronize()\n",
    "t0 = time.time()\n",
    "_ = vdpm_model.inference([{'img': torch.rand(1,3,518,518, device='cuda' if torch.cuda.is_available() else 'cpu')}])\n",
    "torch.cuda.synchronize()\n",
    "print('inference time (s):', time.time() - t0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "4dgs-dpm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}