Upload weights, notebooks, sample images

notebooks/UnReflectAnything.ipynb

@@ -0,0 +1,485 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# UnReflectAnything API & CLI Examples\n",
+        "---\n",
+        "\n",
+        "### 1. Installation and assets download\n",
+        "Ensure you have installed UnReflectAnything with \n",
+        "```bash\n",
+        "pip install unreflectanything\n",
+        "```\n",
+        "this will also install the CLI, which is also callable with aliases `unreflect` and `ura`. Verify installation and check the version with:\n",
+        "```bash\n",
+        "unreflectanything --help\n",
+        "```\n",
+        "```bash\n",
+        "unreflect --version\n",
+        "```\n",
+        "```bash\n",
+        "ura --version\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 31,
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Using device: cuda\n"
+          ]
+        }
+      ],
+      "source": [
+        "import torch\n",
+        "from pathlib import Path\n",
+        "\n",
+        "# Import UnreflectAnything!\n",
+        "import unreflectanything\n",
+        "\n",
+        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+        "print(f\"Using device: {device}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "`pip install`ing UnReflectAnything does not download the pretrained model weights. Download them with the cli command\n",
+        "```bash\n",
+        "unrefleactanything download --weights\n",
+        "```\n",
+        "or"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 32,
+      "metadata": {},
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "e2614e279f5f40609b18ac82d696b01d",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Fetching 4 files:   0%|          | 0/4 [00:00<?, ?it/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "7293fbadfcc74c438a3cb5a0e4df1ac2",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "weights/diffuse_decoder.pt:   0%|          | 0.00/418M [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "f1253cb4b21d48c2b0c3780f4c4060ef",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "weights/token_inpainter.pt:   0%|          | 0.00/307M [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "74ed792f61264beab4a890ea5ede8a58",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "weights/highlight_decoder.pt:   0%|          | 0.00/54.8M [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "1ea92cd82ad34af68f06be823dd12060",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "weights/full_model_weights.pt:   0%|          | 0.00/3.55G [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Weights saved to /home/arota/.cache/unreflectanything/weights\n"
+          ]
+        }
+      ],
+      "source": [
+        "weights_dir = unreflectanything.download(\"weights\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "511c670a",
+      "metadata": {},
+      "source": [
+        "Download some sample images which will be used in this notebook with \n",
+        "```bash\n",
+        "unrefleactanything download --images\n",
+        "```\n",
+        "or"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "e0c60b28",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "images_dir = unreflectanything.download(\"images\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "36417577",
+      "metadata": {},
+      "source": [
+        "### 2. Running UnReflectAnything with pretrained weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "c318961c",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Instantating the pretrained default UnreflectAnything model. \n",
+        "unreflect = unreflectanything.model(device=device)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "448d5456",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from PIL import Image\n",
+        "import numpy as np\n",
+        "\n",
+        "# Building a simple dataloader on a simple dataset that loads from a dir of images\n",
+        "sample_dataset = unreflectanything.ImageDirDataset(images_dir)\n",
+        "sample_dataloader = torch.utils.data.DataLoader(\n",
+        "    sample_dataset, batch_size=1, shuffle=False\n",
+        ")\n",
+        "\n",
+        "# Threshold and Dilation in inpaint mask can be overridden; defaults 0.2 and 40\n",
+        "THRESHOLD = 0.2\n",
+        "DILATION = 40\n",
+        "\n",
+        "# Process and display only N images out of the full sample dataset\n",
+        "DISPLAY_N_IMAGES = 2\n",
+        "\n",
+        "outputs = []\n",
+        "for batch in sample_dataloader:\n",
+        "    # Forward pass\n",
+        "    batch_output = unreflect(\n",
+        "        batch.to(device), return_dict=True, threshold=THRESHOLD, dilation=DILATION\n",
+        "    )\n",
+        "    outputs.append(batch_output)\n",
+        "    if len(outputs) >= DISPLAY_N_IMAGES:\n",
+        "        break\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "98c60bf4",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Helper: Convert tensor [H, W, C] in [0,1] float32 to uint8 to display them \n",
+        "def tensor_to_uint8_img(t):\n",
+        "    arr = t.permute(1, 2, 0).detach().numpy()\n",
+        "    arr = np.clip(arr, 0, 1)\n",
+        "    arr = (arr * 255).round().astype(np.uint8)\n",
+        "    return arr\n",
+        "\n",
+        "# Plotting a collage of the input, the diffuse output, and the highlight mask\n",
+        "for input_batch, output_batch in zip(sample_dataloader, outputs):\n",
+        "    concat_images = torch.cat(\n",
+        "        [\n",
+        "            input_batch.cpu(),\n",
+        "            output_batch[\"diffuse\"].cpu(),\n",
+        "            output_batch[\"highlight\"].repeat(1, 3, 1, 1).cpu(), # \n",
+        "        ],\n",
+        "        dim=3,\n",
+        "    )\n",
+        "    for sample in concat_images:\n",
+        "        img_uint8 = tensor_to_uint8_img(sample)\n",
+        "        display(Image.fromarray(img_uint8))\n",
+        "    # break\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "cf0e6ac6",
+      "metadata": {},
+      "source": [
+        "### 3. Inference API and CLI endpoint\n",
+        "The `inference` wrapper instantiates the UnReflectAnything model and calls its forward function is a single API call. It either:\n",
+        "- Inputs a batched image tensor and outputs a batched image tensor\n",
+        "- Inputs the path to an image (or directory of images) and saves the output results at a given path (of file or directory)\n",
+        "- Inputs the path to an image and outputs a batched image tensor \n",
+        "\n",
+        "Some example CLI calls:\n",
+        "```bash\n",
+        "unreflect inference path/to/image/dir/ -o output/dir/ --threshold 0.3 --dilation 40\n",
+        "```\n",
+        "```bash\n",
+        "unreflect inference path/to/image.png -o path/to/output.png --threshold 0.3 --dilation 40\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Pick a sample image from the downloaded assets. `input` can also be the path to a dir\n",
+        "input_path = list(images_dir.glob(\"*.png\"))[0]\n",
+        "print(\"Input file: \", input_path)\n",
+        "# Specify the outptut name. If `input` is a path to a dir, `output` should be too.\n",
+        "output_path = Path(\"output_example.png\").resolve()\n",
+        "print(\"Output file: \", output_path)\n",
+        "\n",
+        "unreflectanything.inference(\n",
+        "    input=input_path,\n",
+        "    output=output_path,\n",
+        "    device=device,\n",
+        "    threshold=THRESHOLD, \n",
+        "    dilation=DILATION,  \n",
+        ")\n",
+        "\n",
+        "# Loading the saved output and original input from files, then displaying them\n",
+        "input_img = Image.open(input_path).convert(\"RGB\")\n",
+        "output_img = Image.open(output_path).convert(\"RGB\")\n",
+        "\n",
+        "def to_tensor(img):\n",
+        "    return torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.\n",
+        "\n",
+        "input_tensor = to_tensor(input_img)\n",
+        "output_tensor = to_tensor(output_img)\n",
+        "concat = torch.cat([input_tensor, output_tensor], dim=2)\n",
+        "concat_uint8 = (concat.permute(1,2,0).numpy() * 255).clip(0,255).astype(np.uint8)\n",
+        "display(Image.fromarray(concat_uint8))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "5118ea92",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "print(\"Equivalent CLI command:\\n\")\n",
+        "print(f\"unreflect inference {input_path} -o {output_path} --threshold {THRESHOLD} --dilation {DILATION}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "562c17f1",
+      "metadata": {},
+      "source": [
+        "`inference` initializes the model every time by default. To run it without this step, pass them model to the API call"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "89fbbf62",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Pick a sample image from the downloaded assets. `input` can also be the path to a dir\n",
+        "input_path = list(images_dir.glob(\"*.png\"))[6]\n",
+        "# Specify the outptu name\n",
+        "output_path = Path(\"output_example.png\")\n",
+        "                   \n",
+        "unreflectanything.inference(\n",
+        "    model=unreflect, # <<<<<<<<< Pass the model instance and it won't be loaded at every `inference` call\n",
+        "    input=input_path,\n",
+        "    output=output_path,\n",
+        "    device=device,\n",
+        "    threshold=THRESHOLD, \n",
+        "    dilation=DILATION,  \n",
+        ")\n",
+        "\n",
+        "# Loading the saved output and original input from files, then displaying them\n",
+        "input_img = Image.open(input_path).convert(\"RGB\")\n",
+        "output_img = Image.open(output_path).convert(\"RGB\")\n",
+        "\n",
+        "def to_tensor(img):\n",
+        "    return torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.\n",
+        "\n",
+        "input_tensor = to_tensor(input_img)\n",
+        "output_tensor = to_tensor(output_img)\n",
+        "concat = torch.cat([input_tensor, output_tensor], dim=2)\n",
+        "concat_uint8 = (concat.permute(1,2,0).numpy() * 255).clip(0,255).astype(np.uint8)\n",
+        "display(Image.fromarray(concat_uint8))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "57441af0",
+      "metadata": {},
+      "source": [
+        "### 4. The Cache Directory\n",
+        "\n",
+        "`unreflectanything download` saves the downloaded asset in your system cache. Print this path with\n",
+        "```bash\n",
+        "unreflectanything cache --dir\n",
+        "```\n",
+        "or clear the cache with \n",
+        "```bash\n",
+        "unreflectanything cache --clear\n",
+        "```\n",
+        "The same endopoints are also on the API"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "b331050d",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "unreflectanything.cache(\"dir\")  # Also unreflectanything.cache()\n",
+        "unreflectanything.cache(\"clear\")\n",
+        "# unreflectanything.cache.clear()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 4. Verify Assets\n",
+        "\n",
+        "You can verify that the weights are correctly downloaded and loadable."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "is_valid = unreflectanything.verify(\"weights\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### CLI Equivalent\n",
+        "\n",
+        "```bash\n",
+        "unreflect verify --weights\n",
+        "```\n",
+        "```bash\n",
+        "unreflect verify --weights\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "85947e38",
+      "metadata": {},
+      "outputs": [],
+      "source": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 5. Cite\n",
+        "\n",
+        "If you use UnReflectAnything in your research, please cite it:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "print(ura.cite(format=\"bibtex\"))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### CLI Equivalent\n",
+        "\n",
+        "```bash\n",
+        "unreflect cite --bibtex\n",
+        "```"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3 (ipykernel)",
+      "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.3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}

notebooks/api_examples.ipynb

@@ -1,253 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "id": "d5e78019",
-      "metadata": {},
-      "source": [
-        "# UnReflectAnything API Examples\n",
-        "---"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "d423248d",
-      "metadata": {},
-      "source": [
-        "### Package Import"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 1,
-      "id": "db2eda79",
-      "metadata": {},
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Using device: cuda\n"
-          ]
-        }
-      ],
-      "source": [
-        "import unreflectanything\n",
-        "import torch\n",
-        "\n",
-        "%load_ext autoreload\n",
-        "%autoreload 2\n",
-        "\n",
-        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-        "print(f\"Using device: {device}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "c3828c5e",
-      "metadata": {},
-      "source": [
-        "### Model Loading"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "cabb1b8a",
-      "metadata": {},
-      "source": [
-        "If you haven't downloaded the pre-trained weights yet, do so with \n",
-        "\n",
-        "`unreflectanything download --weights` from the terminal\n",
-        "\n",
-        "\n",
-        "or with `unreflectanything.download(\"weights\")` from Python."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 6,
-      "id": "d58ad7f1",
-      "metadata": {},
-      "outputs": [
-        {
-          "data": {
-            "text/html": [
-              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">MODEL    <span style=\"font-weight: bold\">[</span><span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\">18:45:03</span><span style=\"font-weight: bold\">]</span> ✓ Decoder <span style=\"color: #008000; text-decoration-color: #008000\">'diffuse'</span>: Successfully loaded all <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">54</span> state dict keys from weights/rgb_decoder.pth\n",
-              "</pre>\n"
-            ],
-            "text/plain": [
-              "MODEL    \u001b[1m[\u001b[0m\u001b[1;92m18:45:03\u001b[0m\u001b[1m]\u001b[0m ✓ Decoder \u001b[32m'diffuse'\u001b[0m: Successfully loaded all \u001b[1;36m54\u001b[0m state dict keys from weights/rgb_decoder.pth\n"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "data": {
-            "text/html": [
-              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">MODEL    <span style=\"font-weight: bold\">[</span><span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\">18:45:03</span><span style=\"font-weight: bold\">]</span> Loaded pre-trained decoder weights from weights/rgb_decoder.pth\n",
-              "</pre>\n"
-            ],
-            "text/plain": [
-              "MODEL    \u001b[1m[\u001b[0m\u001b[1;92m18:45:03\u001b[0m\u001b[1m]\u001b[0m Loaded pre-trained decoder weights from weights/rgb_decoder.pth\n"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "data": {
-            "text/html": [
-              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">MODEL    <span style=\"font-weight: bold\">[</span><span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\">18:45:03</span><span style=\"font-weight: bold\">]</span> ✓ Token Inpainter: Successfully loaded all <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">78</span> state dict keys from weights/token_inpainter.pth\n",
-              "</pre>\n"
-            ],
-            "text/plain": [
-              "MODEL    \u001b[1m[\u001b[0m\u001b[1;92m18:45:03\u001b[0m\u001b[1m]\u001b[0m ✓ Token Inpainter: Successfully loaded all \u001b[1;36m78\u001b[0m state dict keys from weights/token_inpainter.pth\n"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "data": {
-            "text/html": [
-              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">MODEL    <span style=\"font-weight: bold\">[</span><span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\">18:45:03</span><span style=\"font-weight: bold\">]</span> Loaded pretrained token inpainter weights from weights/token_inpainter.pth\n",
-              "</pre>\n"
-            ],
-            "text/plain": [
-              "MODEL    \u001b[1m[\u001b[0m\u001b[1;92m18:45:03\u001b[0m\u001b[1m]\u001b[0m Loaded pretrained token inpainter weights from weights/token_inpainter.pth\n"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Warning: missing keys when loading checkpoint: ['decoders.highlight.reassemble_layers.0.proj.weight', 'decoders.highlight.reassemble_layers.0.proj.bias', 'decoders.highlight.reassemble_layers.0.resample.weight', 'decoders.highlight.reassemble_layers.0.resample.bias', 'decoders.highlight.reassemble_layers.1.proj.weight', 'decoders.highlight.reassemble_layers.1.proj.bias', 'decoders.highlight.reassemble_layers.1.resample.weight', 'decoders.highlight.reassemble_layers.1.resample.bias', 'decoders.highlight.reassemble_layers.2.proj.weight', 'decoders.highlight.reassemble_layers.2.proj.bias', 'decoders.highlight.reassemble_layers.3.proj.weight', 'decoders.highlight.reassemble_layers.3.proj.bias', 'decoders.highlight.reassemble_layers.3.resample.weight', 'decoders.highlight.reassemble_layers.3.resample.bias', 'decoders.highlight.fusion_blocks.0.residual_conv1.weight', 'decoders.highlight.fusion_blocks.0.residual_conv1.bias', 'decoders.highlight.fusion_blocks.0.residual_conv2.0.weight', 'decoders.highlight.fusion_blocks.0.residual_conv2.0.bias', 'decoders.highlight.fusion_blocks.0.residual_conv2.3.weight', 'decoders.highlight.fusion_blocks.0.residual_conv2.3.bias', 'decoders.highlight.fusion_blocks.0.out_conv.weight', 'decoders.highlight.fusion_blocks.0.out_conv.bias', 'decoders.highlight.fusion_blocks.1.residual_conv1.weight', 'decoders.highlight.fusion_blocks.1.residual_conv1.bias', 'decoders.highlight.fusion_blocks.1.residual_conv2.0.weight', 'decoders.highlight.fusion_blocks.1.residual_conv2.0.bias', 'decoders.highlight.fusion_blocks.1.residual_conv2.3.weight', 'decoders.highlight.fusion_blocks.1.residual_conv2.3.bias', 'decoders.highlight.fusion_blocks.1.out_conv.weight', 'decoders.highlight.fusion_blocks.1.out_conv.bias', 'decoders.highlight.fusion_blocks.2.residual_conv1.weight', 'decoders.highlight.fusion_blocks.2.residual_conv1.bias', 'decoders.highlight.fusion_blocks.2.residual_conv2.0.weight', 'decoders.highlight.fusion_blocks.2.residual_conv2.0.bias', 'decoders.highlight.fusion_blocks.2.residual_conv2.3.weight', 'decoders.highlight.fusion_blocks.2.residual_conv2.3.bias', 'decoders.highlight.fusion_blocks.2.out_conv.weight', 'decoders.highlight.fusion_blocks.2.out_conv.bias', 'decoders.highlight.fusion_blocks.3.residual_conv1.weight', 'decoders.highlight.fusion_blocks.3.residual_conv1.bias', 'decoders.highlight.fusion_blocks.3.residual_conv2.0.weight', 'decoders.highlight.fusion_blocks.3.residual_conv2.0.bias', 'decoders.highlight.fusion_blocks.3.residual_conv2.3.weight', 'decoders.highlight.fusion_blocks.3.residual_conv2.3.bias', 'decoders.highlight.fusion_blocks.3.out_conv.weight', 'decoders.highlight.fusion_blocks.3.out_conv.bias', 'decoders.highlight.rgb_head.0.weight', 'decoders.highlight.rgb_head.0.bias', 'decoders.highlight.rgb_head.5.weight', 'decoders.highlight.rgb_head.5.bias', 'decoders.highlight.rgb_head.9.weight', 'decoders.highlight.rgb_head.9.bias', 'decoders.highlight.rgb_head.13.weight', 'decoders.highlight.rgb_head.13.bias', 'token_inpaint.mask_token', 'token_inpaint.mask_indicator', 'token_inpaint.blocks.0.attn.norm.weight', 'token_inpaint.blocks.0.attn.norm.bias', 'token_inpaint.blocks.0.attn.fn.attn.in_proj_weight', 'token_inpaint.blocks.0.attn.fn.attn.in_proj_bias', 'token_inpaint.blocks.0.attn.fn.attn.out_proj.weight', 'token_inpaint.blocks.0.attn.fn.attn.out_proj.bias', 'token_inpaint.blocks.0.mlp.norm.weight', 'token_inpaint.blocks.0.mlp.norm.bias', 'token_inpaint.blocks.0.mlp.fn.fc1.weight', 'token_inpaint.blocks.0.mlp.fn.fc1.bias', 'token_inpaint.blocks.0.mlp.fn.fc2.weight', 'token_inpaint.blocks.0.mlp.fn.fc2.bias', 'token_inpaint.blocks.1.attn.norm.weight', 'token_inpaint.blocks.1.attn.norm.bias', 'token_inpaint.blocks.1.attn.fn.attn.in_proj_weight', 'token_inpaint.blocks.1.attn.fn.attn.in_proj_bias', 'token_inpaint.blocks.1.attn.fn.attn.out_proj.weight', 'token_inpaint.blocks.1.attn.fn.attn.out_proj.bias', 'token_inpaint.blocks.1.mlp.norm.weight', 'token_inpaint.blocks.1.mlp.norm.bias', 'token_inpaint.blocks.1.mlp.fn.fc1.weight', 'token_inpaint.blocks.1.mlp.fn.fc1.bias', 'token_inpaint.blocks.1.mlp.fn.fc2.weight', 'token_inpaint.blocks.1.mlp.fn.fc2.bias', 'token_inpaint.blocks.2.attn.norm.weight', 'token_inpaint.blocks.2.attn.norm.bias', 'token_inpaint.blocks.2.attn.fn.attn.in_proj_weight', 'token_inpaint.blocks.2.attn.fn.attn.in_proj_bias', 'token_inpaint.blocks.2.attn.fn.attn.out_proj.weight', 'token_inpaint.blocks.2.attn.fn.attn.out_proj.bias', 'token_inpaint.blocks.2.mlp.norm.weight', 'token_inpaint.blocks.2.mlp.norm.bias', 'token_inpaint.blocks.2.mlp.fn.fc1.weight', 'token_inpaint.blocks.2.mlp.fn.fc1.bias', 'token_inpaint.blocks.2.mlp.fn.fc2.weight', 'token_inpaint.blocks.2.mlp.fn.fc2.bias', 'token_inpaint.blocks.3.attn.norm.weight', 'token_inpaint.blocks.3.attn.norm.bias', 'token_inpaint.blocks.3.attn.fn.attn.in_proj_weight', 'token_inpaint.blocks.3.attn.fn.attn.in_proj_bias', 'token_inpaint.blocks.3.attn.fn.attn.out_proj.weight', 'token_inpaint.blocks.3.attn.fn.attn.out_proj.bias', 'token_inpaint.blocks.3.mlp.norm.weight', 'token_inpaint.blocks.3.mlp.norm.bias', 'token_inpaint.blocks.3.mlp.fn.fc1.weight', 'token_inpaint.blocks.3.mlp.fn.fc1.bias', 'token_inpaint.blocks.3.mlp.fn.fc2.weight', 'token_inpaint.blocks.3.mlp.fn.fc2.bias', 'token_inpaint.blocks.4.attn.norm.weight', 'token_inpaint.blocks.4.attn.norm.bias', 'token_inpaint.blocks.4.attn.fn.attn.in_proj_weight', 'token_inpaint.blocks.4.attn.fn.attn.in_proj_bias', 'token_inpaint.blocks.4.attn.fn.attn.out_proj.weight', 'token_inpaint.blocks.4.attn.fn.attn.out_proj.bias', 'token_inpaint.blocks.4.mlp.norm.weight', 'token_inpaint.blocks.4.mlp.norm.bias', 'token_inpaint.blocks.4.mlp.fn.fc1.weight', 'token_inpaint.blocks.4.mlp.fn.fc1.bias', 'token_inpaint.blocks.4.mlp.fn.fc2.weight', 'token_inpaint.blocks.4.mlp.fn.fc2.bias', 'token_inpaint.blocks.5.attn.norm.weight', 'token_inpaint.blocks.5.attn.norm.bias', 'token_inpaint.blocks.5.attn.fn.attn.in_proj_weight', 'token_inpaint.blocks.5.attn.fn.attn.in_proj_bias', 'token_inpaint.blocks.5.attn.fn.attn.out_proj.weight', 'token_inpaint.blocks.5.attn.fn.attn.out_proj.bias', 'token_inpaint.blocks.5.mlp.norm.weight', 'token_inpaint.blocks.5.mlp.norm.bias', 'token_inpaint.blocks.5.mlp.fn.fc1.weight', 'token_inpaint.blocks.5.mlp.fn.fc1.bias', 'token_inpaint.blocks.5.mlp.fn.fc2.weight', 'token_inpaint.blocks.5.mlp.fn.fc2.bias', 'token_inpaint.out_proj.weight', 'token_inpaint.out_proj.bias', 'token_inpaint._final_norm.weight', 'token_inpaint._final_norm.bias']\n"
-          ]
-        }
-      ],
-      "source": [
-        "# unreflectanything.download(\"weights\")\n",
-        "# unreflectanything.download(\"images\") # --> Loads 20 sample images\n",
-        "unreflectanythingmodel = unreflectanything.model(pretrained=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "f3dfa889",
-      "metadata": {},
-      "source": [
-        "Load a dataset of images. Change `PATH_TO_IMAGE_DIR` to point to your own image directory"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "da39fa39",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from unreflectanything import ImageDirDataset, get_cache_dir\n",
-        "from torch.utils.data import DataLoader\n",
-        "\n",
-        "PATH_TO_IMAGE_DIR = get_cache_dir(\n",
-        "    \"images\"\n",
-        ")  # Modify this path to point to your image directory\n",
-        "\n",
-        "ds = ImageDirDataset(PATH_TO_IMAGE_DIR, target_size=(448, 448), return_path=False)\n",
-        "loader = DataLoader(ds, batch_size=1, shuffle=False)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "4c8312f0",
-      "metadata": {},
-      "source": [
-        "### Forward Pass / Inference"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 8,
-      "id": "34e01754",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "output_images = [unreflectanythingmodel(batch_images) for batch_images in loader]"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "94690751",
-      "metadata": {},
-      "source": [
-        "### Displaying results"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 9,
-      "id": "a130c042",
-      "metadata": {},
-      "outputs": [
-        {
-          "ename": "RuntimeError",
-          "evalue": "Sizes of tensors must match except in dimension 3. Expected size 896 but got size 448 for tensor number 1 in the list.",
-          "output_type": "error",
-          "traceback": [
-            "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
-            "\u001b[31mRuntimeError\u001b[39m                              Traceback (most recent call last)",
-            "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[9]\u001b[39m\u001b[32m, line 14\u001b[39m\n\u001b[32m     10\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m arr\n\u001b[32m     13\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m input_batch, output_batch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(loader, output_images):\n\u001b[32m---> \u001b[39m\u001b[32m14\u001b[39m     concat_images = \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcat\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m     15\u001b[39m \u001b[43m        \u001b[49m\u001b[43m[\u001b[49m\u001b[43minput_batch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcpu\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput_batch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcpu\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m3\u001b[39;49m\n\u001b[32m     16\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# (B, 3, H, 2W)\u001b[39;00m\n\u001b[32m     17\u001b[39m     \u001b[38;5;28;01mfor\u001b[39;00m sample \u001b[38;5;129;01min\u001b[39;00m concat_images:\n\u001b[32m     18\u001b[39m         img_uint8 = tensor_to_uint8_img(sample)\n",
-            "\u001b[31mRuntimeError\u001b[39m: Sizes of tensors must match except in dimension 3. Expected size 896 but got size 448 for tensor number 1 in the list."
-          ]
-        }
-      ],
-      "source": [
-        "from PIL import Image\n",
-        "import numpy as np\n",
-        "\n",
-        "\n",
-        "# Helper: Convert tensor [H, W, C] in [0,1] float32 to uint8\n",
-        "def tensor_to_uint8_img(t):\n",
-        "    arr = t.permute(1, 2, 0).cpu().detach().numpy()\n",
-        "    arr = np.clip(arr, 0, 1)\n",
-        "    arr = (arr * 255).round().astype(np.uint8)\n",
-        "    return arr\n",
-        "\n",
-        "\n",
-        "for input_batch, output_batch in zip(loader, output_images):\n",
-        "    concat_images = torch.cat(\n",
-        "        [input_batch.cpu(), output_batch.cpu()], dim=3\n",
-        "    )  # (B, 3, H, 2W)\n",
-        "    for sample in concat_images:\n",
-        "        img_uint8 = tensor_to_uint8_img(sample)\n",
-        "        display(Image.fromarray(img_uint8))\n",
-        "    break\n"
-      ]
-    }
-  ],
-  "metadata": {
-    "kernelspec": {
-      "display_name": "Python 3 (ipykernel)",
-      "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.11"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 5
-}