add inference script

app.py

@@ -1,95 +1,44 @@
-# Part of the code is from: fashn-ai/sapiens-body-part-segmentation
 import os
-
 import gradio as gr
 import numpy as np
-import spaces
-import torch
-from gradio.themes.utils import sizes
 from PIL import Image
-from torchvision import transforms
-from utils.vis_utils import get_palette, visualize_mask_with_overlay
-from config import SAPIENS_LITE_MODELS_PATH
-
-if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-
-CHECKPOINTS_DIR = "checkpoints"
-
-def load_model(checkpoint_name: str):
-    checkpoint_path = os.path.join(CHECKPOINTS_DIR, CHECKPOINTS[checkpoint_name])
-    model = torch.jit.load(checkpoint_path)
-    model.eval()
-    model.to("cuda")
-    return model
-
-
-#MODELS = {name: load_model(name) for name in CHECKPOINTS.keys()}
-
-@torch.inference_mode()
-def run_model(model, input_tensor, height, width):
-    output = model(input_tensor)
-    output = torch.nn.functional.interpolate(output, size=(height, width), mode="bilinear", align_corners=False)
-    _, preds = torch.max(output, 1)
-    return preds
-
-
-transform_fn = transforms.Compose(
-    [
-        transforms.Resize((1024, 768)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ]
-)
-
-@spaces.GPU
-def segment(image: Image.Image, model_name: str) -> Image.Image:
-    input_tensor = transform_fn(image).unsqueeze(0).to("cuda")
-    model = MODELS[model_name]
-    preds = run_model(model, input_tensor, height=image.height, width=image.width)
-    mask = preds.squeeze(0).cpu().numpy()
-    mask_image = Image.fromarray(mask.astype("uint8"))
-    blended_image = visualize_mask_with_overlay(image, mask_image, LABELS_TO_IDS, alpha=0.5)
-    return blended_image
 
+from inference.seg import process_image_or_video
+from config import SAPIENS_LITE_MODELS_PATH
 
 def update_model_choices(task):
     model_choices = list(SAPIENS_LITE_MODELS_PATH[task.lower()].keys())
     return gr.Dropdown(choices=model_choices, value=model_choices[0] if model_choices else None)
 
+def gradio_wrapper(input_image, task, version):
+    if isinstance(input_image, np.ndarray):
+        input_image = Image.fromarray(input_image)
+    
+    result = process_image_or_video(input_image, task=task.lower(), version=version)
+    
+    return result
+
 with gr.Blocks() as demo:
     gr.Markdown("# Sapiens Arena 🤸🏽‍♂️ - WIP devmode- Not yet available")
     with gr.Tabs():
         with gr.TabItem('Image'):
             with gr.Row():
                 with gr.Column():
-                    input_image = gr.Image(label="Input Image", type="pil", format="png")
+                    input_image = gr.Image(label="Input Image", type="pil")
                     select_task = gr.Radio(
-                        ["Seg", "Pose", "Depth", "Normal"], 
+                        ["seg", "pose", "depth", "normal"], 
                         label="Task", 
-                        info="Choose the task to perfom",
-                        choices=list(SAPIENS_LITE_MODELS_PATH.keys())
+                        info="Choose the task to perform",
+                        value="seg"
                     )
                     model_name = gr.Dropdown(
                         label="Model Version",
                         choices=list(SAPIENS_LITE_MODELS_PATH["seg"].keys()),
-                        value="0.3B",
+                        value="sapiens_0.3b",
                     )
-
-                    # example_model = gr.Examples(
-                    #     inputs=input_image,
-                    #     examples_per_page=10,
-                    #     examples=[
-                    #         os.path.join(ASSETS_DIR, "examples", img)
-                    #         for img in os.listdir(os.path.join(ASSETS_DIR, "examples"))
-                    #     ],
-                    # )
                 with gr.Column():
-                    result_image = gr.Image(label="Segmentation Result", format="png")
+                    result_image = gr.Image(label="Result")
                     run_button = gr.Button("Run")
-
-                    #gr.Image(os.path.join(ASSETS_DIR, "legend.png"), label="Legend", type="filepath")
         
         with gr.TabItem('Video'):
             gr.Markdown("In construction")
@@ -97,11 +46,10 @@ with gr.Blocks() as demo:
     select_task.change(fn=update_model_choices, inputs=select_task, outputs=model_name)
 
     run_button.click(
-        fn=segment,
-        inputs=[input_image, model_name],
+        fn=gradio_wrapper,
+        inputs=[input_image, select_task, model_name],
         outputs=[result_image],
     )
 
-
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch(share=True)

inference/seg.py

@@ -0,0 +1,67 @@
+import torch
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+from config import LABELS_TO_IDS
+from utils.vis_utils import visualize_mask_with_overlay
+
+def load_model(task, version):
+    from config import SAPIENS_LITE_MODELS_PATH
+    import os
+
+    try:
+        model_path = SAPIENS_LITE_MODELS_PATH[task][version]
+        if not os.path.exists(model_path):
+            print(f"Advertencia: El archivo del modelo no existe en {model_path}")
+            return None, None
+        
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = torch.jit.load(model_path)
+        model.eval()
+        model.to(device)
+        return model, device
+    except KeyError as e:
+        print(f"Error: Tarea o versión inválida. {e}")
+        return None, None
+
+def process_image_or_video(input_data, task='seg', version='sapiens_0.3b'):
+    # Configurar el modelo
+    model, device = load_model(task, version)
+    if model is None or device is None:
+        return None
+
+    # Configurar la transformación de entrada
+    transform_fn = transforms.Compose([
+        transforms.Resize((1024, 768)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    # Función para procesar un solo frame
+    def process_frame(frame):
+        if isinstance(frame, np.ndarray):
+            frame = Image.fromarray(frame)
+        
+        if frame.mode == 'RGBA':
+            frame = frame.convert('RGB')
+        
+        input_tensor = transform_fn(frame).unsqueeze(0).to(device)
+        
+        with torch.inference_mode():
+            output = model(input_tensor)
+            output = torch.nn.functional.interpolate(output, size=(frame.height, frame.width), mode="bilinear", align_corners=False)
+            _, preds = torch.max(output, 1)
+        
+        mask = preds.squeeze(0).cpu().numpy()
+        mask_image = Image.fromarray(mask.astype("uint8"))
+        blended_image = visualize_mask_with_overlay(frame, mask_image, LABELS_TO_IDS, alpha=0.5)
+        return blended_image
+
+    # Procesar imagen o video
+    if isinstance(input_data, np.ndarray):  # Video frame
+        return process_frame(input_data)
+    elif isinstance(input_data, Image.Image):  # Imagen
+        return process_frame(input_data)
+    else:
+        print("Tipo de entrada no soportado. Por favor, proporcione una imagen PIL o un frame de video numpy.")
+        return None

load_and_test.ipynb

@@ -3146,7 +3146,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 83,
+   "execution_count": 90,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -3155,6 +3155,104 @@
     "import numpy as np\n",
     "import cv2\n",
     "\n",
+    "def get_depth(image, depth_model, input_shape=(3, 1024, 768), device=\"cuda\"):\n",
+    "    # Preprocess the image\n",
+    "    img = preprocess_image(image, input_shape)\n",
+    "    \n",
+    "    # Run the model\n",
+    "    with torch.no_grad():\n",
+    "        result = depth_model(img.to(device))\n",
+    "    \n",
+    "    # Post-process the output\n",
+    "    depth_map = post_process_depth(result, (image.shape[0], image.shape[1]))\n",
+    "    \n",
+    "    # Visualize the depth map\n",
+    "    depth_image = visualize_depth(depth_map)\n",
+    "    \n",
+    "    return depth_image, depth_map\n",
+    "\n",
+    "def preprocess_image(image, input_shape):\n",
+    "    img = cv2.resize(image, (input_shape[2], input_shape[1]), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1)\n",
+    "    img = torch.from_numpy(img)\n",
+    "    img = img[[2, 1, 0], ...].float()\n",
+    "    mean = torch.tensor([123.5, 116.5, 103.5]).view(-1, 1, 1)\n",
+    "    std = torch.tensor([58.5, 57.0, 57.5]).view(-1, 1, 1)\n",
+    "    img = (img - mean) / std\n",
+    "    return img.unsqueeze(0)\n",
+    "\n",
+    "def post_process_depth(result, original_shape):\n",
+    "    # Check the dimensionality of the result\n",
+    "    if result.dim() == 3:\n",
+    "        result = result.unsqueeze(0)\n",
+    "    elif result.dim() == 4:\n",
+    "        pass\n",
+    "    else:\n",
+    "        raise ValueError(f\"Unexpected result dimension: {result.dim()}\")\n",
+    "    \n",
+    "    # Ensure we're interpolating to the correct dimensions\n",
+    "    seg_logits = F.interpolate(result, size=original_shape, mode=\"bilinear\", align_corners=False).squeeze(0)\n",
+    "    depth_map = seg_logits.data.float().cpu().numpy()\n",
+    "    \n",
+    "    # If depth_map has an extra dimension, squeeze it\n",
+    "    if depth_map.ndim == 3 and depth_map.shape[0] == 1:\n",
+    "        depth_map = depth_map.squeeze(0)\n",
+    "    \n",
+    "    return depth_map\n",
+    "\n",
+    "# def visualize_depth(depth_map):\n",
+    "#     # Normalize the depth map\n",
+    "#     min_val, max_val = np.nanmin(depth_map), np.nanmax(depth_map)\n",
+    "#     depth_normalized = (depth_map - min_val) / (max_val - min_val)\n",
+    "#     depth_normalized = (depth_normalized * 255.0).astype(np.uint8)\n",
+    "    \n",
+    "#     # Apply color map\n",
+    "#     depth_colored = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)\n",
+    "    \n",
+    "#     return depth_colored\n",
+    "\n",
+    "# def post_process_depth(result, original_shape):\n",
+    "#     seg_logits = F.interpolate(result.unsqueeze(0), size=original_shape, mode=\"bilinear\").squeeze(0)\n",
+    "#     depth_map = seg_logits.data.float().cpu().numpy()[0]  # H x W\n",
+    "#     return depth_map\n",
+    "\n",
+    "def visualize_depth(depth_map):\n",
+    "    # Normalize the depth map\n",
+    "    min_val, max_val = np.nanmin(depth_map), np.nanmax(depth_map)\n",
+    "    depth_normalized = 1 - ((depth_map - min_val) / (max_val - min_val))\n",
+    "    \n",
+    "    # Convert to uint8\n",
+    "    depth_normalized = (depth_normalized * 255).astype(np.uint8)\n",
+    "    \n",
+    "    # Apply colormap\n",
+    "    depth_colored = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)\n",
+    "    \n",
+    "    return depth_colored\n",
+    "\n",
+    "# You can add the surface normal calculation if needed\n",
+    "def calculate_surface_normal(depth_map):\n",
+    "    kernel_size = 7\n",
+    "    grad_x = cv2.Sobel(depth_map.astype(np.float32), cv2.CV_32F, 1, 0, ksize=kernel_size)\n",
+    "    grad_y = cv2.Sobel(depth_map.astype(np.float32), cv2.CV_32F, 0, 1, ksize=kernel_size)\n",
+    "    z = np.full(grad_x.shape, -1)\n",
+    "    normals = np.dstack((-grad_x, -grad_y, z))\n",
+    "\n",
+    "    normals_mag = np.linalg.norm(normals, axis=2, keepdims=True)\n",
+    "    with np.errstate(divide=\"ignore\", invalid=\"ignore\"):\n",
+    "        normals_normalized = normals / (normals_mag + 1e-5)\n",
+    "\n",
+    "    normals_normalized = np.nan_to_num(normals_normalized, nan=-1, posinf=-1, neginf=-1)\n",
+    "    normal_from_depth = ((normals_normalized + 1) / 2 * 255).astype(np.uint8)\n",
+    "    normal_from_depth = normal_from_depth[:, :, ::-1]  # RGB to BGR for cv2\n",
+    "\n",
+    "    return normal_from_depth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 94,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "def get_depth(image, depth_model, input_shape=(3, 1024, 768), device=\"cuda\"):\n",
     "    # Preprocess the image\n",
     "    img = preprocess_image(image, input_shape)\n",
@@ -3202,18 +3300,38 @@
     "def visualize_depth(depth_map):\n",
     "    # Normalize the depth map\n",
     "    min_val, max_val = np.nanmin(depth_map), np.nanmax(depth_map)\n",
-    "    depth_normalized = (depth_map - min_val) / (max_val - min_val)\n",
-    "    depth_normalized = (depth_normalized * 255.0).astype(np.uint8)\n",
+    "    depth_normalized = 1 - ((depth_map - min_val) / (max_val - min_val))\n",
     "    \n",
-    "    # Apply color map\n",
+    "    # Convert to uint8\n",
+    "    depth_normalized = (depth_normalized * 255).astype(np.uint8)\n",
+    "    \n",
+    "    # Apply colormap\n",
     "    depth_colored = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)\n",
     "    \n",
-    "    return depth_colored"
+    "    return depth_colored\n",
+    "\n",
+    "# You can add the surface normal calculation if needed\n",
+    "def calculate_surface_normal(depth_map):\n",
+    "    kernel_size = 7\n",
+    "    grad_x = cv2.Sobel(depth_map.astype(np.float32), cv2.CV_32F, 1, 0, ksize=kernel_size)\n",
+    "    grad_y = cv2.Sobel(depth_map.astype(np.float32), cv2.CV_32F, 0, 1, ksize=kernel_size)\n",
+    "    z = np.full(grad_x.shape, -1)\n",
+    "    normals = np.dstack((-grad_x, -grad_y, z))\n",
+    "\n",
+    "    normals_mag = np.linalg.norm(normals, axis=2, keepdims=True)\n",
+    "    with np.errstate(divide=\"ignore\", invalid=\"ignore\"):\n",
+    "        normals_normalized = normals / (normals_mag + 1e-5)\n",
+    "\n",
+    "    normals_normalized = np.nan_to_num(normals_normalized, nan=-1, posinf=-1, neginf=-1)\n",
+    "    normal_from_depth = ((normals_normalized + 1) / 2 * 255).astype(np.uint8)\n",
+    "    normal_from_depth = normal_from_depth[:, :, ::-1]  # RGB to BGR for cv2\n",
+    "\n",
+    "    return normal_from_depth"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 84,
+   "execution_count": 99,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -3222,31 +3340,913 @@
     "pil_image = Image.open('/home/user/app/assets/image.webp')\n",
     "\n",
     "# Load and process an image\n",
-    "image = cv2.imread('/home/user/app/assets/image.webp')\n",
+    "image = cv2.imread('/home/user/app/assets/frame.png')\n",
     "depth_image, depth_map = get_depth(image, model)\n",
     "\n",
+    "surface_normal = calculate_surface_normal(depth_map)\n",
+    "cv2.imwrite(\"output_surface_normal.jpg\", surface_normal)\n",
     "# Save the results\n",
-    "output_im = cv2.imwrite(\"output_depth_image.jpg\", depth_image)"
+    "output_im = cv2.imwrite(\"output_depth_image2.jpg\", depth_image)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Normal"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 85,
+   "execution_count": 100,
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 85,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "checkpoints/normal/sapiens_0.3b_torchscript.pt2\n"
+     ]
     }
    ],
    "source": [
-    "output_im"
+    "# Example usage\n",
+    "TASK = 'normal'\n",
+    "VERSION = 'sapiens_0.3b'\n",
+    "\n",
+    "model_path = get_model_path(TASK, VERSION)\n",
+    "print(model_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 101,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "RecursiveScriptModule(\n",
+       "  original_name=DepthEstimator\n",
+       "  (data_preprocessor): RecursiveScriptModule(original_name=SegDataPreProcessor)\n",
+       "  (backbone): RecursiveScriptModule(\n",
+       "    original_name=VisionTransformer\n",
+       "    (patch_embed): RecursiveScriptModule(\n",
+       "      original_name=PatchEmbed\n",
+       "      (projection): RecursiveScriptModule(original_name=Conv2d)\n",
+       "    )\n",
+       "    (drop_after_pos): RecursiveScriptModule(original_name=Dropout)\n",
+       "    (layers): RecursiveScriptModule(\n",
+       "      original_name=ModuleList\n",
+       "      (0): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (1): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (2): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (3): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (4): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (5): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (6): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (7): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (8): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (9): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (10): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (11): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (12): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (13): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (14): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (15): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (16): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (17): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (18): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (19): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (20): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (21): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (22): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "      (23): RecursiveScriptModule(\n",
+       "        original_name=TransformerEncoderLayer\n",
+       "        (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (attn): RecursiveScriptModule(\n",
+       "          original_name=MultiheadAttention\n",
+       "          (qkv): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj): RecursiveScriptModule(original_name=Linear)\n",
+       "          (proj_drop): RecursiveScriptModule(original_name=Dropout)\n",
+       "          (out_drop): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma1): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "        (ln2): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "        (ffn): RecursiveScriptModule(\n",
+       "          original_name=FFN\n",
+       "          (layers): RecursiveScriptModule(\n",
+       "            original_name=Sequential\n",
+       "            (0): RecursiveScriptModule(\n",
+       "              original_name=Sequential\n",
+       "              (0): RecursiveScriptModule(original_name=Linear)\n",
+       "              (1): RecursiveScriptModule(original_name=GELU)\n",
+       "              (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "            )\n",
+       "            (1): RecursiveScriptModule(original_name=Linear)\n",
+       "            (2): RecursiveScriptModule(original_name=Dropout)\n",
+       "          )\n",
+       "          (dropout_layer): RecursiveScriptModule(original_name=DropPath)\n",
+       "          (gamma2): RecursiveScriptModule(original_name=Identity)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (pre_norm): RecursiveScriptModule(original_name=Identity)\n",
+       "    (ln1): RecursiveScriptModule(original_name=LayerNorm)\n",
+       "  )\n",
+       "  (decode_head): RecursiveScriptModule(\n",
+       "    original_name=VitNormalHead\n",
+       "    (loss_decode): RecursiveScriptModule(\n",
+       "      original_name=ModuleList\n",
+       "      (0): RecursiveScriptModule(original_name=CosineSimilarityLoss)\n",
+       "      (1): RecursiveScriptModule(original_name=L1Loss)\n",
+       "    )\n",
+       "    (conv_seg): RecursiveScriptModule(original_name=Conv2d)\n",
+       "    (dropout): RecursiveScriptModule(original_name=Dropout2d)\n",
+       "    (deconv_layers): RecursiveScriptModule(\n",
+       "      original_name=Sequential\n",
+       "      (0): RecursiveScriptModule(original_name=ConvTranspose2d)\n",
+       "      (1): RecursiveScriptModule(original_name=InstanceNorm2d)\n",
+       "      (2): RecursiveScriptModule(original_name=SiLU)\n",
+       "      (3): RecursiveScriptModule(original_name=ConvTranspose2d)\n",
+       "      (4): RecursiveScriptModule(original_name=InstanceNorm2d)\n",
+       "      (5): RecursiveScriptModule(original_name=SiLU)\n",
+       "      (6): RecursiveScriptModule(original_name=ConvTranspose2d)\n",
+       "      (7): RecursiveScriptModule(original_name=InstanceNorm2d)\n",
+       "      (8): RecursiveScriptModule(original_name=SiLU)\n",
+       "    )\n",
+       "    (conv_layers): RecursiveScriptModule(\n",
+       "      original_name=Sequential\n",
+       "      (0): RecursiveScriptModule(original_name=Conv2d)\n",
+       "      (1): RecursiveScriptModule(original_name=InstanceNorm2d)\n",
+       "      (2): RecursiveScriptModule(original_name=SiLU)\n",
+       "      (3): RecursiveScriptModule(original_name=Conv2d)\n",
+       "      (4): RecursiveScriptModule(original_name=InstanceNorm2d)\n",
+       "      (5): RecursiveScriptModule(original_name=SiLU)\n",
+       "      (6): RecursiveScriptModule(original_name=Conv2d)\n",
+       "      (7): RecursiveScriptModule(original_name=InstanceNorm2d)\n",
+       "      (8): RecursiveScriptModule(original_name=SiLU)\n",
+       "    )\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 101,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model = torch.jit.load(model_path)\n",
+    "model.eval()\n",
+    "model.to(\"cuda\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 105,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "import numpy as np\n",
+    "import cv2\n",
+    "\n",
+    "def get_normal(image, normal_model, input_shape=(3, 1024, 768), device=\"cuda\"):\n",
+    "    # Preprocess the image\n",
+    "    img = preprocess_image(image, input_shape)\n",
+    "    \n",
+    "    # Run the model\n",
+    "    with torch.no_grad():\n",
+    "        result = normal_model(img.to(device))\n",
+    "    \n",
+    "    # Post-process the output\n",
+    "    normal_map = post_process_normal(result, (image.shape[0], image.shape[1]))\n",
+    "    \n",
+    "    # Visualize the normal map\n",
+    "    normal_image = visualize_normal(normal_map)\n",
+    "    \n",
+    "    return normal_image, normal_map\n",
+    "\n",
+    "def preprocess_image(image, input_shape):\n",
+    "    img = cv2.resize(image, (input_shape[2], input_shape[1]), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1)\n",
+    "    img = torch.from_numpy(img)\n",
+    "    img = img[[2, 1, 0], ...].float()\n",
+    "    mean = torch.tensor([123.5, 116.5, 103.5]).view(-1, 1, 1)\n",
+    "    std = torch.tensor([58.5, 57.0, 57.5]).view(-1, 1, 1)\n",
+    "    img = (img - mean) / std\n",
+    "    return img.unsqueeze(0)\n",
+    "\n",
+    "def post_process_normal(result, original_shape):\n",
+    "    # Check the dimensionality of the result\n",
+    "    if result.dim() == 3:\n",
+    "        result = result.unsqueeze(0)\n",
+    "    elif result.dim() == 4:\n",
+    "        pass\n",
+    "    else:\n",
+    "        raise ValueError(f\"Unexpected result dimension: {result.dim()}\")\n",
+    "    \n",
+    "    # Ensure we're interpolating to the correct dimensions\n",
+    "    seg_logits = F.interpolate(result, size=original_shape, mode=\"bilinear\", align_corners=False).squeeze(0)\n",
+    "    normal_map = seg_logits.float().cpu().numpy().transpose(1, 2, 0)  # H x W x 3\n",
+    "    return normal_map\n",
+    "\n",
+    "def visualize_normal(normal_map):\n",
+    "    normal_map_norm = np.linalg.norm(normal_map, axis=-1, keepdims=True)\n",
+    "    normal_map_normalized = normal_map / (normal_map_norm + 1e-5)  # Add a small epsilon to avoid division by zero\n",
+    "    \n",
+    "    # Convert to 0-255 range and BGR format for visualization\n",
+    "    normal_map_vis = ((normal_map_normalized + 1) / 2 * 255).astype(np.uint8)\n",
+    "    normal_map_vis = normal_map_vis[:, :, ::-1]  # RGB to BGR\n",
+    "    \n",
+    "    return normal_map_vis\n",
+    "\n",
+    "def load_normal_model(checkpoint, use_torchscript=False):\n",
+    "    if use_torchscript:\n",
+    "        return torch.jit.load(checkpoint)\n",
+    "    else:\n",
+    "        model = torch.export.load(checkpoint).module()\n",
+    "        model = model.to(\"cuda\")\n",
+    "        model = torch.compile(model, mode=\"max-autotune\", fullgraph=True)\n",
+    "        return model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 107,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 107,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import cv2\n",
+    "import numpy as np\n",
+    "\n",
+    "# Load the model\n",
+    "normal_model = load_normal_model(model_path, use_torchscript='_torchscript')\n",
+    "\n",
+    "# Load the image\n",
+    "image = cv2.imread(\"/home/user/app/assets/image.webp\")\n",
+    "\n",
+    "# Get the normal map and visualization\n",
+    "normal_image, normal_map = get_normal(image, normal_model)\n",
+    "\n",
+    "# Save the results\n",
+    "cv2.imwrite(\"output_normal_image.png\", normal_image)"
    ]
   },
   {