343 lines
9.9 KiB
Plaintext
343 lines
9.9 KiB
Plaintext
{
|
|
"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
|
|
}
|