{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d45abdaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[SPARSE] Backend: spconv, Attention: flash_attn\n",
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "# load necessary libraries\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ.setdefault(\"SPCONV_ALGO\", \"native\")\n",
    "os.environ.setdefault(\"ATTN_BACKEND\", \"flash_attn\")\n",
    "os.environ.setdefault(\"TORCH_HOME\", os.path.expanduser(\"~/.cache/torch\"))\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from pytorch3d.ops import cubify\n",
    "import trimesh\n",
    "\n",
    "from dvd import DVDImageToVoxelPipeline, as_voxel_output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe08603f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /homes/zx1321/.cache/torch/hub/facebookresearch_dinov2_main\n"
     ]
    }
   ],
   "source": [
    "# set device and load the BSP fine-tuned DVD image editing pipeline\n",
    "\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "# from local files\n",
    "\n",
    "# dvd_pipeline = DVDImageToVoxelPipeline.from_files(\n",
    "#     \"./ckpts/dvd_img_BSP_ft.json\",\n",
    "#     \"./ckpts/dvd_img_BSP_ft.safetensors\",\n",
    "#     device=device,\n",
    "# )\n",
    "\n",
    "# or from pretrained\n",
    "dvd_pipeline = DVDImageToVoxelPipeline.from_pretrained(\"Zhengrui/dvd\",variant=\"bsp\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0d798a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load a voxel grid in DVD coordinate convention and visualize it\n",
    "coord = np.load(\"./assets/example_voxel_edit/voxel64_typical_building_mushroom_dis.npy\")\n",
    "voxels = as_voxel_output(torch.from_numpy(coord), resolution=64)\n",
    "samples = voxels.samples.to(device)\n",
    "\n",
    "cubified_meshes = cubify(samples.float(), 0.5, align=\"center\")\n",
    "mesh = trimesh.Trimesh(\n",
    "    vertices=cubified_meshes.verts_packed().cpu().numpy(),\n",
    "    faces=cubified_meshes.faces_packed().cpu().numpy(),\n",
    ")\n",
    "mesh.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2bcf6a2",
   "metadata": {},
   "source": [
    "## Edit with alternative condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0d49cb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# perturb the roof part of the generated shape\n",
    "# The DVD edit sampler preserves voxels where keep_mask=True and regenerates where keep_mask=False.\n",
    "edit_samples = samples.clone().long()\n",
    "noise = torch.randint(0, 2, edit_samples.shape, device=device)\n",
    "edit_samples[:, :, 28:, :] = noise[:, :, 28:, :]\n",
    "keep_mask = torch.ones_like(edit_samples, dtype=torch.bool)\n",
    "keep_mask[:, :, 28:, :] = False\n",
    "\n",
    "# visualize perturbed mesh\n",
    "cubified_meshes = cubify(edit_samples.float(), 0.5, align=\"center\")\n",
    "mesh = trimesh.Trimesh(\n",
    "    vertices=cubified_meshes.verts_packed().cpu().numpy(),\n",
    "    faces=cubified_meshes.faces_packed().cpu().numpy(),\n",
    ")\n",
    "mesh.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e955ac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the alternative image condition\n",
    "image_path = \"./assets/example_image_edit/flower_rm.png\"\n",
    "image = Image.open(image_path)\n",
    "image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa8bd25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The DVD pipeline now obtains the image condition internally.\n",
    "voxels_to_edit = as_voxel_output(edit_samples, resolution=64)\n",
    "print(voxels_to_edit.samples.shape, keep_mask.shape)\n",
    "\n",
    "res = dvd_pipeline.edit_voxels(\n",
    "    image,\n",
    "    voxels_to_edit,\n",
    "    keep_mask=keep_mask,\n",
    "    seed=0,\n",
    "    steps=128,\n",
    "    cfg_strength=0.45,\n",
    "    preprocess_image=True,\n",
    "    verbose=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d09af34",
   "metadata": {},
   "outputs": [],
   "source": [
    "edited_samples = res.samples.to(device)\n",
    "cubified_meshes = cubify(edited_samples.float(), 0.5, align=\"center\")\n",
    "mesh = trimesh.Trimesh(\n",
    "    vertices=cubified_meshes.verts_packed().cpu().numpy(),\n",
    "    faces=cubified_meshes.faces_packed().cpu().numpy(),\n",
    ")\n",
    "mesh.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}