jaejung.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# make dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_opt():\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--path', default='./data/photoshapes/shape09096_rank00',type=str)\n",
    "    parser.add_argument('-O', action='store_true', help=\"equals --fp16 --cuda_ray --preload\")\n",
    "    parser.add_argument('--test', action='store_true', help=\"test mode\")\n",
    "    parser.add_argument('--workspace', type=str, default='workspace')\n",
    "    parser.add_argument('--seed', type=int, default=0)\n",
    "    ### training options\n",
    "    parser.add_argument('--iters', type=int, default=15000, help=\"training iters\")\n",
    "    parser.add_argument('--lr', type=float, default=1e-2, help=\"initial learning rate\")\n",
    "    parser.add_argument('--ckpt', type=str, default='best')\n",
    "    parser.add_argument('--num_rays', type=int, default=4096, help=\"num rays sampled per image for each training step\")\n",
    "    parser.add_argument('--cuda_ray', action='store_true', help=\"use CUDA raymarching instead of pytorch\")\n",
    "    parser.add_argument('--max_steps', type=int, default=1024, help=\"max num steps sampled per ray (only valid when using --cuda_ray)\")\n",
    "    parser.add_argument('--num_steps', type=int, default=512, help=\"num steps sampled per ray (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--upsample_steps', type=int, default=0, help=\"num steps up-sampled per ray (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--update_extra_interval', type=int, default=16, help=\"iter interval to update extra status (only valid when using --cuda_ray)\")\n",
    "    parser.add_argument('--max_ray_batch', type=int, default=4096, help=\"batch size of rays at inference to avoid OOM (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--patch_size', type=int, default=1, help=\"[experimental] render patches in training, so as to apply LPIPS loss. 1 means disabled, use [64, 32, 16] to enable\")\n",
    "    ### network backbone options\n",
    "    parser.add_argument('--fp16', action='store_true', help=\"use amp mixed precision training\")\n",
    "    parser.add_argument('--ff', action='store_true', help=\"use fully-fused MLP\")\n",
    "    parser.add_argument('--tcnn', action='store_true', help=\"use TCNN backend\")\n",
    "    ### dataset options\n",
    "    parser.add_argument('--color_space', type=str, default='srgb', help=\"Color space, supports (linear, srgb)\")\n",
    "    parser.add_argument('--preload', action='store_true', help=\"preload all data into GPU, accelerate training but use more GPU memory\")\n",
    "    # (the default value is for the fox dataset)\n",
    "    parser.add_argument('--bound', type=float, default=2, help=\"assume the scene is bounded in box[-bound, bound]^3, if > 1, will invoke adaptive ray marching.\")\n",
    "    parser.add_argument('--scale', type=float, default=0.33, help=\"scale camera location into box[-bound, bound]^3\")\n",
    "    parser.add_argument('--offset', type=float, nargs='*', default=[0, 0, 0], help=\"offset of camera location\")\n",
    "    parser.add_argument('--dt_gamma', type=float, default=1/128, help=\"dt_gamma (>=0) for adaptive ray marching. set to 0 to disable, >0 to accelerate rendering (but usually with worse quality)\")\n",
    "    parser.add_argument('--min_near', type=float, default=0.2, help=\"minimum near distance for camera\")\n",
    "    parser.add_argument('--density_thresh', type=float, default=10, help=\"threshold for density grid to be occupied\")\n",
    "    parser.add_argument('--bg_radius', type=float, default=-1, help=\"if positive, use a background model at sphere(bg_radius)\")\n",
    "    ### GUI options\n",
    "    parser.add_argument('--gui', action='store_true', help=\"start a GUI\")\n",
    "    parser.add_argument('--W', type=int, default=1920, help=\"GUI width\")\n",
    "    parser.add_argument('--H', type=int, default=1080, help=\"GUI height\")\n",
    "    parser.add_argument('--radius', type=float, default=5, help=\"default GUI camera radius from center\")\n",
    "    parser.add_argument('--fovy', type=float, default=50, help=\"default GUI camera fovy\")\n",
    "    parser.add_argument('--max_spp', type=int, default=64, help=\"GUI rendering max sample per pixel\")\n",
    "    ### experimental\n",
    "    parser.add_argument('--error_map', action='store_true', help=\"use error map to sample rays\")\n",
    "    parser.add_argument('--clip_text', type=str, default='', help=\"text input for CLIP guidance\")\n",
    "    parser.add_argument('--rand_pose', type=int, default=-1, help=\"<0 uses no rand pose, =0 only uses rand pose, >0 sample one rand pose every $ known poses\")\n",
    "\n",
    "\n",
    "\n",
    "    return parser.parse_args(args=[])\n",
    "opt=get_opt()\n",
    "\n",
    "opt.fp16 = True\n",
    "opt.cuda_ray = True\n",
    "opt.preload = True\n",
    "\n",
    "opt.test=True\n",
    "\n",
    "if opt.patch_size > 1:\n",
    "    opt.error_map = False # do not use error_map if use patch-based training\n",
    "    # assert opt.patch_size > 16, \"patch_size should > 16 to run LPIPS loss.\"\n",
    "    assert opt.num_rays % (opt.patch_size ** 2) == 0, \"patch_size ** 2 should be dividable by num_rays.\"\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checking the number of parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param=0\n",
    "for name,child in model.named_children():\n",
    "    print(f'module_name={name}, parameter={sum(p.numel() for p in child.parameters() if p.requires_grad)}')\n",
    "    param+= sum(p.numel() for p in child.parameters() if p.requires_grad)\n",
    "\n",
    "print(param)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Slerp, Rotation\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "from nerf.utils import *\n",
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "model=NeRFNetwork(\n",
    "        encoding=\"hashgrid\",\n",
    "        bound=opt.bound,\n",
    "        cuda_ray=opt.cuda_ray,\n",
    "        density_scale=1,\n",
    "        min_near=opt.min_near,\n",
    "        density_thresh=opt.density_thresh,\n",
    "        bg_radius=opt.bg_radius,\n",
    "    ).to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name,child in model.named_children():\n",
    "    print(f'module_name={name}, parameter={sum(p.numel() for p in child.parameters() if p.requires_grad)}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Random generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Slerp, Rotation\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "from nerf.utils import *\n",
    "\n",
    "device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')\n",
    "opt.path='./data/photoshapes/shape09096_rank01'\n",
    "opt.cuda_ray=True\n",
    "model = NeRFNetwork(\n",
    "        encoding=\"hashgrid\",\n",
    "        bound=opt.bound,\n",
    "        cuda_ray=opt.cuda_ray,\n",
    "        density_scale=1,\n",
    "        min_near=opt.min_near,\n",
    "        density_thresh=opt.density_thresh,\n",
    "        bg_radius=opt.bg_radius,\n",
    "    ).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "ckpt=torch.load('./data/photoshape_weight/shape09096_rank01/checkpoints/ngp.pth')\n",
    "\n",
    "model.load_state_dict(ckpt['model'],strict=False)\n",
    "model.to(device)\n",
    "model.eval()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- if you have \"transforms_test.json\" file in path, this code make video of your dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- This code is making random pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.cuda.amp import autocast\n",
    "\n",
    "poses=rand_poses(5,device,theta_range=[np.pi/3,np.pi/3],radius=0.5)\n",
    "# make intrinsic\n",
    "focal=245/(2*np.tan(np.radians(60)/2))\n",
    "intrinsics=np.array([focal,focal,256//2,256//2])\n",
    "data=get_rays(poses,intrinsics,256,256)\n",
    "\n",
    "with autocast():\n",
    "    with torch.no_grad():\n",
    "        result=model.render(data['rays_o'].float(),data['rays_d'].float())\n",
    "\n",
    "pred=result['image'].reshape(-1,256,256,3)\n",
    "pred=srgb_to_linear(pred).detach().cpu().numpy()\n",
    "pred=(pred*255).astype(np.uint8)\n",
    "\n",
    "cv2.imwrite(os.path.join(\"/home/poong/hyperdiff\",'sample1.png'), cv2.cvtColor(pred[3], cv2.COLOR_RGB2BGR))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Line Pose generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.cuda.amp import autocast\n",
    "\n",
    "focal=245/(2*np.tan(np.radians(60)/2))\n",
    "intrinsics=np.array([focal,focal,256//2,256//2])\n",
    "poses=line_poses(10,device,radius=0.5)\n",
    "data=get_rays(poses,intrinsics,256,256)\n",
    "with autocast():\n",
    "    with torch.no_grad():\n",
    "        result=model.render(data['rays_o'].float(),data['rays_d'].float())\n",
    "\n",
    "pred=result['image'].reshape(-1,256,256,3)\n",
    "pred=srgb_to_linear(pred).detach().cpu().numpy()\n",
    "pred=(pred*255).astype(np.uint8)\n",
    "\n",
    "#cv2.imwrite(os.path.join(\"/home/poong/hyperdiff\",'sample0.png'), cv2.cvtColor(pred[3], cv2.COLOR_RGB2BGR))\n",
    "imageio.mimwrite(os.path.join(\"/home/poong/hyperdiff\",'sample.mp4'),pred, fps=2, quality=8, macro_block_size=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = [PSNRMeter(), LPIPSMeter(device=device)]\n",
    "test_loader = NeRFDataset(opt, device=device, type='test').dataloader()\n",
    "trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, fp16=opt.fp16, metrics=metrics, use_checkpoint=opt.ckpt)\n",
    "trainer.test(test_loader, write_video=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Seed Checking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Slerp, Rotation\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "opt=get_opt()\n",
    "device = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#seed_everything(opt.seed)\n",
    "model = NeRFNetwork(\n",
    "        encoding=\"hashgrid\",\n",
    "        bound=opt.bound,\n",
    "        cuda_ray=opt.cuda_ray,\n",
    "        density_scale=1,\n",
    "        min_near=opt.min_near,\n",
    "        density_thresh=opt.density_thresh,\n",
    "        bg_radius=opt.bg_radius,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3072]), torch.Size([7168]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.sigma_net.params.shape, model.color_net.params.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tinycudann as tcnn\n",
    "encoder = tcnn.Encoding(\n",
    "    n_input_dims=3,\n",
    "    seed=1337,\n",
    "    encoding_config={\n",
    "        \"otype\": \"HashGrid\",\n",
    "        \"n_levels\": 16,\n",
    "        \"n_features_per_level\": 2,\n",
    "        \"log2_hashmap_size\": 14,\n",
    "        \"base_resolution\": 16,\n",
    "        \"per_level_scale\": np.exp2(np.log2(2048 * 1/ 16) / (16 - 1)),\n",
    "    },\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from weight_encoding.transformer.encoder import *\n",
    "from weight_encoding.transformer.block import Weight_Split,DecoderLayer\n",
    "from weight_encoding.transformer.decoder import *\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import dataset\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "from nerf.utils import *\n",
    "from tqdm import tqdm\n",
    "\n",
    "ckpt=torch.load('./data/photoshape_weight/shape09096_rank00/checkpoints/ngp.pth')\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "param=torch.concat((ckpt['model']['encoder.params'],ckpt['model']['sigma_net.params'],ckpt['model']['color_net.params']),-1)\n",
    "splited_param=Weight_Split(param.size(),param.size(), chunk_size=1024)(param.unsqueeze(0)).to(device)                               \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from weight_encoding.transformer.transformer import *\n",
    "import torch.optim as optim\n",
    "\n",
    "transformer = Transformer(device).to(device)\n",
    "optimizer = torch.optim.AdamW(transformer.parameters(), lr=5e-5)\n",
    "criterion =nn.L1Loss()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pbar = tqdm(range(5000))\n",
    "for epoch in pbar:\n",
    "    optimizer.zero_grad()\n",
    "    output=transformer(splited_param)\n",
    "    loss = criterion(output,splited_param)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    pbar.set_description('loss: %.6f'                    %loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result=transformer(splited_param)\n",
    "index1=len(ckpt['model']['encoder.params'])\n",
    "index2=index1+len(ckpt['model']['sigma_net.params'])\n",
    "index3=index2+len(ckpt['model']['color_net.params'])\n",
    "ckpt['model']['encoder.params']=result.flatten()[:index1]\n",
    "ckpt['model']['sigma_net.params']=result.flatten()[index1:index2]\n",
    "ckpt['model']['color_net.params']=result.flatten()[index2:]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reconstruct test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_opt():\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--path', default='./data/photoshapes/shape09096_rank00',type=str)\n",
    "    parser.add_argument('-O', action='store_true', help=\"equals --fp16 --cuda_ray --preload\")\n",
    "    parser.add_argument('--test', action='store_true', help=\"test mode\")\n",
    "    parser.add_argument('--workspace', type=str, default='workspace')\n",
    "    parser.add_argument('--seed', type=int, default=0)\n",
    "    ### training options\n",
    "    parser.add_argument('--iters', type=int, default=15000, help=\"training iters\")\n",
    "    parser.add_argument('--lr', type=float, default=1e-2, help=\"initial learning rate\")\n",
    "    parser.add_argument('--ckpt', type=str, default='best')\n",
    "    parser.add_argument('--num_rays', type=int, default=4096, help=\"num rays sampled per image for each training step\")\n",
    "    parser.add_argument('--cuda_ray', action='store_true', help=\"use CUDA raymarching instead of pytorch\")\n",
    "    parser.add_argument('--max_steps', type=int, default=1024, help=\"max num steps sampled per ray (only valid when using --cuda_ray)\")\n",
    "    parser.add_argument('--num_steps', type=int, default=512, help=\"num steps sampled per ray (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--upsample_steps', type=int, default=0, help=\"num steps up-sampled per ray (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--update_extra_interval', type=int, default=16, help=\"iter interval to update extra status (only valid when using --cuda_ray)\")\n",
    "    parser.add_argument('--max_ray_batch', type=int, default=4096, help=\"batch size of rays at inference to avoid OOM (only valid when NOT using --cuda_ray)\")\n",
    "    parser.add_argument('--patch_size', type=int, default=1, help=\"[experimental] render patches in training, so as to apply LPIPS loss. 1 means disabled, use [64, 32, 16] to enable\")\n",
    "    ### network backbone options\n",
    "    parser.add_argument('--fp16', action='store_true', help=\"use amp mixed precision training\")\n",
    "    parser.add_argument('--ff', action='store_true', help=\"use fully-fused MLP\")\n",
    "    parser.add_argument('--tcnn', action='store_true', help=\"use TCNN backend\")\n",
    "    ### dataset options\n",
    "    parser.add_argument('--color_space', type=str, default='srgb', help=\"Color space, supports (linear, srgb)\")\n",
    "    parser.add_argument('--preload', action='store_true', help=\"preload all data into GPU, accelerate training but use more GPU memory\")\n",
    "    # (the default value is for the fox dataset)\n",
    "    parser.add_argument('--bound', type=float, default=2, help=\"assume the scene is bounded in box[-bound, bound]^3, if > 1, will invoke adaptive ray marching.\")\n",
    "    parser.add_argument('--scale', type=float, default=0.33, help=\"scale camera location into box[-bound, bound]^3\")\n",
    "    parser.add_argument('--offset', type=float, nargs='*', default=[0, 0, 0], help=\"offset of camera location\")\n",
    "    parser.add_argument('--dt_gamma', type=float, default=1/128, help=\"dt_gamma (>=0) for adaptive ray marching. set to 0 to disable, >0 to accelerate rendering (but usually with worse quality)\")\n",
    "    parser.add_argument('--min_near', type=float, default=0.2, help=\"minimum near distance for camera\")\n",
    "    parser.add_argument('--density_thresh', type=float, default=10, help=\"threshold for density grid to be occupied\")\n",
    "    parser.add_argument('--bg_radius', type=float, default=-1, help=\"if positive, use a background model at sphere(bg_radius)\")\n",
    "    ### GUI options\n",
    "    parser.add_argument('--gui', action='store_true', help=\"start a GUI\")\n",
    "    parser.add_argument('--W', type=int, default=1920, help=\"GUI width\")\n",
    "    parser.add_argument('--H', type=int, default=1080, help=\"GUI height\")\n",
    "    parser.add_argument('--radius', type=float, default=5, help=\"default GUI camera radius from center\")\n",
    "    parser.add_argument('--fovy', type=float, default=50, help=\"default GUI camera fovy\")\n",
    "    parser.add_argument('--max_spp', type=int, default=64, help=\"GUI rendering max sample per pixel\")\n",
    "    ### experimental\n",
    "    parser.add_argument('--error_map', action='store_true', help=\"use error map to sample rays\")\n",
    "    parser.add_argument('--clip_text', type=str, default='', help=\"text input for CLIP guidance\")\n",
    "    parser.add_argument('--rand_pose', type=int, default=-1, help=\"<0 uses no rand pose, =0 only uses rand pose, >0 sample one rand pose every $ known poses\")\n",
    "\n",
    "\n",
    "\n",
    "    return parser.parse_args(args=[])\n",
    "opt=get_opt()\n",
    "\n",
    "opt.fp16 = True\n",
    "opt.cuda_ray = True\n",
    "opt.preload = True\n",
    "\n",
    "opt.test=True\n",
    "\n",
    "if opt.patch_size > 1:\n",
    "    opt.error_map = False # do not use error_map if use patch-based training\n",
    "    # assert opt.patch_size > 16, \"patch_size should > 16 to run LPIPS loss.\"\n",
    "    assert opt.num_rays % (opt.patch_size ** 2) == 0, \"patch_size ** 2 should be dividable by num_rays.\"\n",
    "    \n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NeRFNetwork(\n",
       "  (encoder): Encoding(n_input_dims=3, n_output_dims=32, seed=1337, dtype=torch.float16, hyperparams={'base_resolution': 16, 'hash': 'CoherentPrime', 'interpolation': 'Linear', 'log2_hashmap_size': 14, 'n_features_per_level': 2, 'n_levels': 16, 'otype': 'Grid', 'per_level_scale': 1.4472692012786865, 'type': 'Hash'})\n",
       "  (sigma_net): Network(n_input_dims=32, n_output_dims=16, seed=1337, dtype=torch.float16, hyperparams={'encoding': {'offset': 0.0, 'otype': 'Identity', 'scale': 1.0}, 'network': {'activation': 'ReLU', 'n_hidden_layers': 1, 'n_neurons': 64, 'otype': 'FullyFusedMLP', 'output_activation': 'None'}, 'otype': 'NetworkWithInputEncoding'})\n",
       "  (encoder_dir): Encoding(n_input_dims=3, n_output_dims=16, seed=1337, dtype=torch.float16, hyperparams={'degree': 4, 'otype': 'SphericalHarmonics'})\n",
       "  (color_net): Network(n_input_dims=31, n_output_dims=3, seed=1337, dtype=torch.float16, hyperparams={'encoding': {'offset': 0.0, 'otype': 'Identity', 'scale': 1.0}, 'network': {'activation': 'ReLU', 'n_hidden_layers': 2, 'n_neurons': 64, 'otype': 'FullyFusedMLP', 'output_activation': 'None'}, 'otype': 'NetworkWithInputEncoding'})\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "\n",
    " \n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "from nerf.utils import *\n",
    "from torch.cuda.amp import autocast\n",
    "opt.cuda_ray=False\n",
    "\n",
    "#ckpt_recon= torch.load(\"/nas2/lait/1000_Members/proinit/pred.bin\")\n",
    "ckpt=torch.load(\"./data/photoshape_weight//shape09096_rank01/checkpoints/ngp.pth\",map_location='cuda')\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "model = NeRFNetwork(\n",
    "        encoding=\"hashgrid\",\n",
    "        bound=opt.bound,\n",
    "        cuda_ray=opt.cuda_ray,\n",
    "        density_scale=1,\n",
    "        min_near=opt.min_near,\n",
    "        density_thresh=opt.density_thresh,\n",
    "        bg_radius=opt.bg_radius,\n",
    "    ).to(device)\n",
    "#model.load_state_dict(ckpt['model'],strict=False)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#ckpt_recon= torch.load(\"/data1/poong/hyperdiff/data/photoshape_weight/weight/shape09096_rank00.bin\")\n",
    "ckpt_recon= torch.load(\"/nas2/lait/tjfwownd/hyperdiff/weight_encoding/weight_transformer/result/test_data_all_MSE.bin\",map_location=device)[2]\n",
    "index1=len(ckpt['model']['encoder.params'])\n",
    "index2=index1+len(ckpt['model']['sigma_net.params'])\n",
    "index3=index2+len(ckpt['model']['color_net.params'])\n",
    "ckpt['model']['encoder.params']=ckpt_recon.flatten()[:index1]\n",
    "ckpt['model']['sigma_net.params']=ckpt_recon.flatten()[index1:index2]\n",
    "ckpt['model']['color_net.params']=ckpt_recon.flatten()[index2:]\n",
    "model.load_state_dict(ckpt['model'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "focal=245/(2*np.tan(np.radians(60)/2))\n",
    "intrinsics=np.array([focal,focal,256//2,256//2])\n",
    "poses=line_poses(40,device,radius=0.5)\n",
    "data=get_rays(poses,intrinsics,256,256)\n",
    "with autocast():\n",
    "    with torch.no_grad():\n",
    "        result=model.render(data['rays_o'],data['rays_d'],staged=True,perturb=False,bg_color=1)\n",
    "        #result=model.render(data['rays_o'].float(),data['rays_d'].float(),staged=True) # when cuda_ray=True, staged must be False\n",
    "\n",
    "pred=result['image'].reshape(-1,256,256,3)\n",
    "pred=pred.detach().cpu().numpy()\n",
    "#pred=pred.detach().cpu().numpy()\n",
    "pred=(pred*255).astype(np.uint8)\n",
    "\n",
    "imageio.mimwrite(os.path.join(\"/nas2/lait/tjfwownd/hyperdiff/workspace\",'sample.mp4'),pred, fps=8, quality=8, macro_block_size=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Slerp, Rotation\n",
    "\n",
    " \n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "from nerf.utils import *\n",
    "from torch.cuda.amp import autocast\n",
    "\n",
    "\n",
    "from nerf.provider import *\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import tqdm\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Slerp, Rotation\n",
    "\n",
    " \n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from nerf.network_tcnn import NeRFNetwork\n",
    "import matplotlib.pyplot as plt\n",
    "from nerf.utils import *\n",
    "from torch.cuda.amp import autocast\n",
    "opt.cuda_ray=True\n",
    "\n",
    "#ckpt_recon= torch.load(\"/nas2/lait/1000_Members/proinit/pred.bin\")\n",
    "ckpt=torch.load('./result/shape09096_rank00/checkpoints/ngp.pth')\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "model = NeRFNetwork(\n",
    "        encoding=\"hashgrid\",\n",
    "        bound=opt.bound,\n",
    "        cuda_ray=opt.cuda_ray,\n",
    "        density_scale=1,\n",
    "        min_near=opt.min_near,\n",
    "        density_thresh=opt.density_thresh,\n",
    "        bg_radius=opt.bg_radius,\n",
    "    ).to(device)\n",
    "model.eval()\n",
    "model.load_state_dict(ckpt['model'],strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt['model'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "ckpt=torch.load(\"./result/shape09182_rank02/checkpoints/ngp.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt['model'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Given normalized_shape=[1024], expected input with shape [*, 1024], but got input of size[504832]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 11\u001b[0m\n\u001b[1;32m      9\u001b[0m     encoder \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m[BasicTransformerBlock(\u001b[39m1024\u001b[39m,\u001b[39m16\u001b[39m,\u001b[39m64\u001b[39m,dropout\u001b[39m=\u001b[39m\u001b[39m0.2\u001b[39m)]\n\u001b[1;32m     10\u001b[0m encoder \u001b[39m=\u001b[39m nn\u001b[39m.\u001b[39mSequential(\u001b[39m*\u001b[39mencoder)\u001b[39m.\u001b[39mto(device)\n\u001b[0;32m---> 11\u001b[0m encoder(ckpt_recon)\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1195\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/torch/nn/modules/container.py:204\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    202\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m):\n\u001b[1;32m    203\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m:\n\u001b[0;32m--> 204\u001b[0m         \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m module(\u001b[39minput\u001b[39;49m)\n\u001b[1;32m    205\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39minput\u001b[39m\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1195\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/diffusers/models/attention.py:141\u001b[0m, in \u001b[0;36mBasicTransformerBlock.forward\u001b[0;34m(self, hidden_states, attention_mask, encoder_hidden_states, encoder_attention_mask, timestep, cross_attention_kwargs, class_labels)\u001b[0m\n\u001b[1;32m    137\u001b[0m     norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnorm1(\n\u001b[1;32m    138\u001b[0m         hidden_states, timestep, class_labels, hidden_dtype\u001b[39m=\u001b[39mhidden_states\u001b[39m.\u001b[39mdtype\n\u001b[1;32m    139\u001b[0m     )\n\u001b[1;32m    140\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 141\u001b[0m     norm_hidden_states \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnorm1(hidden_states)\n\u001b[1;32m    143\u001b[0m cross_attention_kwargs \u001b[39m=\u001b[39m cross_attention_kwargs \u001b[39mif\u001b[39;00m cross_attention_kwargs \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39melse\u001b[39;00m {}\n\u001b[1;32m    144\u001b[0m attn_output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mattn1(\n\u001b[1;32m    145\u001b[0m     norm_hidden_states,\n\u001b[1;32m    146\u001b[0m     encoder_hidden_states\u001b[39m=\u001b[39mencoder_hidden_states \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39monly_cross_attention \u001b[39melse\u001b[39;00m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m    147\u001b[0m     attention_mask\u001b[39m=\u001b[39mattention_mask,\n\u001b[1;32m    148\u001b[0m     \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mcross_attention_kwargs,\n\u001b[1;32m    149\u001b[0m )\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1195\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/torch/nn/modules/normalization.py:190\u001b[0m, in \u001b[0;36mLayerNorm.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    189\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m: Tensor) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Tensor:\n\u001b[0;32m--> 190\u001b[0m     \u001b[39mreturn\u001b[39;00m F\u001b[39m.\u001b[39;49mlayer_norm(\n\u001b[1;32m    191\u001b[0m         \u001b[39minput\u001b[39;49m, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mnormalized_shape, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbias, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49meps)\n",
      "File \u001b[0;32m~/anaconda3/envs/ngp/lib/python3.9/site-packages/torch/nn/functional.py:2515\u001b[0m, in \u001b[0;36mlayer_norm\u001b[0;34m(input, normalized_shape, weight, bias, eps)\u001b[0m\n\u001b[1;32m   2511\u001b[0m \u001b[39mif\u001b[39;00m has_torch_function_variadic(\u001b[39minput\u001b[39m, weight, bias):\n\u001b[1;32m   2512\u001b[0m     \u001b[39mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m   2513\u001b[0m         layer_norm, (\u001b[39minput\u001b[39m, weight, bias), \u001b[39minput\u001b[39m, normalized_shape, weight\u001b[39m=\u001b[39mweight, bias\u001b[39m=\u001b[39mbias, eps\u001b[39m=\u001b[39meps\n\u001b[1;32m   2514\u001b[0m     )\n\u001b[0;32m-> 2515\u001b[0m \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39;49mlayer_norm(\u001b[39minput\u001b[39;49m, normalized_shape, weight, bias, eps, torch\u001b[39m.\u001b[39;49mbackends\u001b[39m.\u001b[39;49mcudnn\u001b[39m.\u001b[39;49menabled)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Given normalized_shape=[1024], expected input with shape [*, 1024], but got input of size[504832]"
     ]
    }
   ],
   "source": [
    "from diffusers.models.attention import BasicTransformerBlock\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "ckpt_recon= torch.load(\"/data1/poong/hyperdiff/weight_encoding/weight_transformer/best.bin\",map_location=device)[0]\n",
    "encoder = []\n",
    "for i in range(14):\n",
    "    encoder +=[BasicTransformerBlock(1024,16,64,dropout=0.2)]\n",
    "encoder = nn.Sequential(*encoder).to(device)\n",
    "encoder(ckpt_recon)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ModuleList(\n",
       "  (0): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (1): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (2): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (3): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (4): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (5): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (6): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (7): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (8): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (9): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (10): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (11): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (12): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (13): BasicTransformerBlock(\n",
       "    (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (attn1): Attention(\n",
       "      (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (to_out): ModuleList(\n",
       "        (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (norm3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (ff): FeedForward(\n",
       "      (net): ModuleList(\n",
       "        (0): GEGLU(\n",
       "          (proj): Linear(in_features=1024, out_features=8192, bias=True)\n",
       "        )\n",
       "        (1): Dropout(p=0.2, inplace=False)\n",
       "        (2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hyperdiff1",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}