{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "j:\\AI\\ComfyUI_windows_portable\\python_embeded\\Lib\\site-packages\\transformers\\utils\\hub.py:124: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "# Function to load and preprocess the image\n",
    "def load_image_as_tensor(image_path, batch_size=1):\n",
    "    # Open the image using PIL\n",
    "    image = Image.open(image_path)\n",
    "    \n",
    "    # Define the transformation pipeline\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((256, 256)),\n",
    "        transforms.ToTensor(),  # Convert PIL Image to tensor and scale to [0, 1]\n",
    "        # transforms.Lambda(lambda x: x.unsqueeze(0).repeat(batch_size, 1, 1, 1))  # Add batch dimension and repeat\n",
    "    ])\n",
    "    \n",
    "    # Apply the transformation\n",
    "    tensor = transform(image)\n",
    "    \n",
    "    return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = load_image_as_tensor(\"examples/006.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 256, 256])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "scale_max = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"models/anysr/anysr_edsr_500.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_load = torch.load(model_path)['model']\n",
    "model_args = model_load['args']['encoder_spec']['args']\n",
    "pretrain_dict = model_load['sd']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_tensor = image.squeeze(dim=0).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 256, 256])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_tensor.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from anysr_module.edsr_anysr import make_edsr\n",
    "\n",
    "\n",
    "model = make_edsr(scale=2, **model_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from anysr_module.anysr_model import replace\n",
    "\n",
    "\n",
    "model_dict = model.state_dict()\n",
    "model_dict = replace(pretrain_dict, model_dict)\n",
    "model.load_state_dict(model_dict, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.scale = 2\n",
    "model.scale2 = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 256, 256])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_tensor.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(\n",
    "    image_tensor.unsqueeze(dim=0).cuda(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 256, 256, 256])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 256, 256, 256])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "shape '[1, 256, 256, 3, 9, 9]' is invalid for input of size 16777216",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[14], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m reshaped_tensor \u001b[38;5;241m=\u001b[39m \u001b[43mpred\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreshape\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m9\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m9\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: shape '[1, 256, 256, 3, 9, 9]' is invalid for input of size 16777216"
     ]
    }
   ],
   "source": [
    "reshaped_tensor = pred.reshape(1, 256, 256, 3, 9, 9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2023, -0.2465, -0.1883,  ...,  0.0145,  0.0291,  0.0131],\n",
       "          [-0.2058, -0.3171, -0.2691,  ...,  0.0181,  0.0362,  0.0219],\n",
       "          [-0.1518, -0.3128, -0.2633,  ...,  0.0039,  0.0277,  0.0250],\n",
       "          ...,\n",
       "          [-0.0476,  0.0612,  0.0954,  ..., -0.1948, -0.2188, -0.1279],\n",
       "          [-0.0426,  0.0376,  0.0445,  ..., -0.1631, -0.1892, -0.0890],\n",
       "          [ 0.0051,  0.0672,  0.0702,  ..., -0.0010, -0.0048,  0.0314]],\n",
       "\n",
       "         [[-0.1818, -0.2641, -0.2487,  ..., -0.1725, -0.1746, -0.1650],\n",
       "          [-0.0194, -0.1230, -0.1304,  ..., -0.2200, -0.2187, -0.1621],\n",
       "          [-0.0632, -0.1623, -0.1519,  ..., -0.2183, -0.2122, -0.1603],\n",
       "          ...,\n",
       "          [-0.1712, -0.1563, -0.1632,  ..., -0.1246, -0.1534, -0.1306],\n",
       "          [-0.2195, -0.2493, -0.2216,  ..., -0.1596, -0.1834, -0.1382],\n",
       "          [-0.1708, -0.2002, -0.1785,  ..., -0.1240, -0.1355, -0.1496]],\n",
       "\n",
       "         [[ 0.0384, -0.1395, -0.1434,  ...,  0.0056, -0.0025, -0.0254],\n",
       "          [ 0.0819, -0.0770, -0.1296,  ..., -0.0080, -0.0228, -0.0414],\n",
       "          [ 0.0429, -0.0804, -0.1248,  ..., -0.0178, -0.0256, -0.0443],\n",
       "          ...,\n",
       "          [ 0.1172,  0.0857,  0.0598,  ..., -0.0916, -0.1025, -0.0719],\n",
       "          [ 0.0778,  0.0410,  0.0348,  ..., -0.0644, -0.0767, -0.0587],\n",
       "          [ 0.0415,  0.0601,  0.0726,  ...,  0.0271,  0.0418,  0.0895]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 0.0076, -0.1249, -0.1511,  ..., -0.0041, -0.0041, -0.0131],\n",
       "          [-0.1826, -0.2713, -0.2433,  ...,  0.0897,  0.0770,  0.0489],\n",
       "          [-0.1696, -0.2405, -0.2122,  ...,  0.0730,  0.0656,  0.0382],\n",
       "          ...,\n",
       "          [ 0.0398,  0.0766,  0.1059,  ..., -0.0038, -0.0599, -0.0317],\n",
       "          [ 0.0230,  0.0949,  0.1537,  ..., -0.0311, -0.0921, -0.0622],\n",
       "          [ 0.0083,  0.0128,  0.0458,  ..., -0.1552, -0.1739, -0.1277]],\n",
       "\n",
       "         [[ 0.1756,  0.2339,  0.2451,  ...,  0.1755,  0.1732,  0.1503],\n",
       "          [ 0.0185, -0.0176,  0.0165,  ...,  0.2317,  0.2341,  0.1946],\n",
       "          [-0.0086, -0.0330,  0.0083,  ...,  0.2167,  0.2240,  0.1849],\n",
       "          ...,\n",
       "          [ 0.1407,  0.1760,  0.2159,  ...,  0.0839,  0.0529,  0.0533],\n",
       "          [ 0.1521,  0.2473,  0.3145,  ...,  0.0638,  0.0380,  0.0335],\n",
       "          [ 0.1696,  0.2355,  0.2771,  ...,  0.0139, -0.0056,  0.0147]],\n",
       "\n",
       "         [[-0.1895, -0.2206, -0.2065,  ..., -0.0522, -0.0691, -0.0413],\n",
       "          [-0.1096, -0.1073, -0.1033,  ..., -0.1391, -0.1615, -0.1086],\n",
       "          [-0.1121, -0.1327, -0.1424,  ..., -0.1380, -0.1547, -0.1007],\n",
       "          ...,\n",
       "          [ 0.0065, -0.0212, -0.0236,  ..., -0.1827, -0.1751, -0.0828],\n",
       "          [-0.0189, -0.0623, -0.0892,  ..., -0.1393, -0.1310, -0.0676],\n",
       "          [-0.0020, -0.0198, -0.0286,  ..., -0.0580, -0.0464, -0.0142]]]],\n",
       "       device='cuda:0', grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}