Upload weights, notebooks, sample images

notebooks/api_examples.ipynb

@@ -5,7 +5,16 @@
       "id": "d5e78019",
       "metadata": {},
       "source": [
-        "# UnReflectAnything API Examples\n"
+        "# UnReflectAnything API Examples\n",
+        "---"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "d423248d",
+      "metadata": {},
+      "source": [
+        "### Package Import"
       ]
     },
     {
@@ -35,162 +44,262 @@
     },
     {
       "cell_type": "markdown",
-      "id": "94f8c2fb",
-      "metadata": {},
-      "source": [
-        "### 1. Get the model class (for custom setup or training)\n",
-        "\n",
-        "`unreflectanything.model()` with no arguments returns the underlying model class `UnReflect_Model_TokenInpainter`. Use it when you need to build the architecture yourself (e.g. from config or for training)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 13,
-      "id": "f49c99b7",
+      "id": "c3828c5e",
       "metadata": {},
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "cuda:0\n"
-          ]
-        }
-      ],
       "source": [
-        "UnReflectModel = unreflectanything.model()\n",
-        "UnReflectModel_Pretrained = unreflectanything.model(pretrained=True)\n",
-        "print((next(UnReflectModel.parameters()).device))"
+        "### Model Loading"
       ]
     },
     {
       "cell_type": "markdown",
-      "id": "575fb9a1",
+      "id": "cabb1b8a",
       "metadata": {},
       "source": [
-        "### 2. Get a pretrained model and run on batched RGB\n",
+        "If you haven't downloaded the pre-trained weights yet, do so with \n",
         "\n",
-        "`unreflectanything.model(pretrained=True)` returns an `UnReflectModel` instance (a `torch.nn.Module`) with weights loaded. Call it with a batch of RGB tensors `[B, 3, H, W]` (values in [0, 1]); it returns the diffuse (reflection-removed) tensor."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "d1cdc14f",
-      "metadata": {},
-      "source": [
-        "#### Load pretrained model (uses cached weights; run 'unreflectanything download --weights' first)"
+        "`unreflectanything download --weights` from the terminal\n",
+        "\n",
+        "\n",
+        "or with `unreflectanything.download(\"weights\")` from Python."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 4,
       "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\">23:19:01</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;92m23:19:01\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\">23:19:01</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;92m23:19:01\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\">23:19:02</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;92m23:19:02\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\">23:19:02</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;92m23:19:02\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": [
-            "Model is nn.Module: True\n",
-            "Expected image size (side): 896\n",
-            "Device: cuda\n"
+            "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": [
-        "import torch\n",
-        "\n",
-        "# Load pretrained model (uses cached weights; run 'unreflectanything download --weights' first)\n",
-        "unreflectanythingmodel = unreflectanything.model(pretrained=True)\n",
-        "unreflectanythingmodel_scratch = unreflectanything.model(pretrained=False)\n",
-        "print(f\"Model is nn.Module: {isinstance(unreflectanythingmodel, torch.nn.Module)}\")\n",
-        "print(f\"Expected image size (side): {unreflectanythingmodel.image_size}\")\n",
-        "print(f\"Device: {unreflectanythingmodel.device}\")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "34e01754",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Batched RGB tensor [B, 3, H, W], values in [0, 1]\n",
-        "batch_size = 2\n",
-        "images = torch.rand(batch_size, 3, 448, 448, device=unreflectanythingmodel.device)\n",
-        "model_out = unreflectanythingmodel(images)  # [B, 3, H, W] diffuse tensor\n",
-        "print(f\"Input shape: {images.shape} -> Output shape: {model_out.shape}\")"
+        "# unreflectanything.download(\"weights\")\n",
+        "# unreflectanything.download(\"images\") # --> Loads 20 sample images\n",
+        "unreflectanythingmodel = unreflectanything.model(pretrained=True)"
       ]
     },
     {
       "cell_type": "markdown",
-      "id": "696bce42",
+      "id": "f3dfa889",
       "metadata": {},
       "source": [
-        "### 3. Full output dict and custom mask (optional)\n",
-        "\n",
-        "You can get the full model outputs (e.g. highlight mask, patch mask) with `return_dict=True`, or pass a custom inpainting mask with `inpaint_mask_override`."
+        "Load a dataset of images. Change `PATH_TO_IMAGE_DIR` to point to your own image directory"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "dc2ecc8a",
+      "id": "da39fa39",
       "metadata": {},
       "outputs": [],
       "source": [
-        "# Get full outputs: diffuse, highlight, patch_mask, etc.\n",
-        "outputs = unreflectanythingmodel(images, return_dict=True)\n",
-        "print(\"Keys:\", list(outputs.keys()))  # e.g. diffuse, highlight, patch_mask, tokens_completed\n",
-        "diffuse_only = outputs[\"diffuse\"]\n",
-        "highlight_mask = outputs[\"highlight\"]  # [B, 1, H, W]"
+        "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=(896, 896), return_path=False)\n",
+        "loader = DataLoader(ds, batch_size=1, shuffle=False)"
       ]
     },
     {
       "cell_type": "markdown",
-      "id": "87fe354c",
+      "id": "4c8312f0",
       "metadata": {},
       "source": [
-        "### 4. One-shot inference (no model handle)\n",
-        "\n",
-        "For a single call without keeping a model in memory, use `unreflectanything.inference()`. It accepts a file path, directory, or tensor and returns a tensor (or writes to disk if `output=` is set)."
+        "### Forward Pass / Inference"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "ff5740b8",
+      "id": "34e01754",
       "metadata": {},
       "outputs": [],
       "source": [
-        "# Tensor in -> tensor out (loads model internally, then discards)\n",
-        "result = unreflectanything.inference(images)\n",
-        "print(f\"unreflectanything.inference(images) shape: {result.shape}\")\n",
-        "\n",
-        "# File-based: save to disk\n",
-        "# unreflectanything.inference(\"input.png\", output=\"output.png\")\n",
-        "# unreflectanything.inference(\"input_dir/\", output=\"output_dir/\", batch_size=8)"
+        "output_images = [unreflectanythingmodel(batch_images) for batch_images in loader]"
       ]
     },
     {
       "cell_type": "markdown",
-      "id": "e2d1673d",
+      "id": "94690751",
       "metadata": {},
       "source": [
-        "### 5. Loading sample images (optional)\n",
-        "\n",
-        "If you have downloaded sample images with `unreflectanything download --images`, you can run inference on that directory."
+        "### Displaying results"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "1834686c",
+      "id": "b588087b",
       "metadata": {},
-      "outputs": [],
+      "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\">22:43:07</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;92m22:43:07\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\">22:43:07</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;92m22:43:07\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\">22:43:08</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;92m22:43:08\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\">22:43:08</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;92m22:43:08\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": [
+        "unreflectanythingmodel = unreflectanything.model(\n",
+        "    pretrained=True,\n",
+        "    config_path=\"huggingface/configs/pretrained_config.yaml\",\n",
+        "    weights_path=get_cache_dir(\"weights\") / \"full_model_weights.pt\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 57,
+      "id": "a130c042",
+      "metadata": {},
+      "outputs": [
+        {
+          "ename": "AttributeError",
+          "evalue": "'dict' object has no attribute 'cpu'",
+          "output_type": "error",
+          "traceback": [
+            "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+            "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
+            "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[57]\u001b[39m\u001b[32m, line 15\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     14\u001b[39m     concat_images = torch.cat(\n\u001b[32m---> \u001b[39m\u001b[32m15\u001b[39m         [input_batch.cpu(), \u001b[43moutput_batch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcpu\u001b[49m()], dim=\u001b[32m3\u001b[39m\n\u001b[32m     16\u001b[39m     )  \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[31mAttributeError\u001b[39m: 'dict' object has no attribute 'cpu'"
+          ]
+        }
+      ],
       "source": [
-        "SAMPLE_IMAGE_PATH_DIR = \"sample_images\"  # default from 'unreflectanything download --images'\n",
-        "# unreflectanything.inference(SAMPLE_IMAGE_PATH_DIR, output=\"output_sample/\", verbose=True)"
+        "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"
       ]
     }
   ],