demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "de794129",
   "metadata": {},
   "source": [
    "# Load the dependcy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "45564f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import io\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "from PIL import Image, ImageChops\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e3166d",
   "metadata": {},
   "source": [
    "Map Encoder dependcy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6c86c8d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CLIPSurgery(\n",
       "  (visual): VisionTransformer(\n",
       "    (conv1): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16), bias=False)\n",
       "    (ln_pre): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (transformer): Transformer(\n",
       "      (resblocks): Sequential(\n",
       "        (0): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (1): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (2): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (3): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (4): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (5): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (6): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (7): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (8): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (9): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (10): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (11): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_post): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (transformer): Transformer(\n",
       "    (resblocks): Sequential(\n",
       "      (0): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (1): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (2): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (3): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (4): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (5): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (6): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (7): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (8): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (9): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (10): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (11): ResidualAttentionBlock(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (token_embedding): Embedding(49408, 512)\n",
       "  (ln_final): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "import clip\n",
    "import torch\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from  matplotlib import pyplot as plt\n",
    "from torchvision.transforms import Compose, Resize, ToTensor, Normalize\n",
    "from torchvision.transforms import InterpolationMode\n",
    "BICUBIC = InterpolationMode.BICUBIC\n",
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "\n",
    "preprocess =  Compose([Resize((224, 224), interpolation=BICUBIC), ToTensor(),\n",
    "    Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))])\n",
    "seg_model, preprocess = clip.load(\"CS-ViT-B/16\", device=device)\n",
    "seg_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16ec52c7",
   "metadata": {},
   "source": [
    "Decoder dependcy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "93f25d6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "A matching Triton is not available, some optimizations will not be enabled.\n",
      "Error caught was: No module named 'triton'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ControlLDM: Running in eps-prediction mode\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "DiffusionWrapper has 865.91 M params.\n",
      "making attention of type 'vanilla-xformers' with 512 in_channels\n",
      "building MemoryEfficientAttnBlock with 512 in_channels...\n",
      "Working with z of shape (1, 4, 32, 32) = 4096 dimensions.\n",
      "making attention of type 'vanilla-xformers' with 512 in_channels\n",
      "building MemoryEfficientAttnBlock with 512 in_channels...\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is None and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 320, context_dim is 1024 and using 5 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is None and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 640, context_dim is 1024 and using 10 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is None and using 20 heads.\n",
      "Setting up MemoryEfficientCrossAttention. Query dim is 1280, context_dim is 1024 and using 20 heads.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Anaconda3\\lib\\site-packages\\torch\\functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at C:\\actions-runner\\_work\\pytorch\\pytorch\\builder\\windows\\pytorch\\aten\\src\\ATen\\native\\TensorShape.cpp:3484.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting up [LPIPS] perceptual loss: trunk [alex], v[0.1], spatial [off]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Anaconda3\\lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "D:\\Anaconda3\\lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=AlexNet_Weights.IMAGENET1K_V1`. You can also use `weights=AlexNet_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from: D:\\Anaconda3\\lib\\site-packages\\lpips\\weights\\v0.1\\alex.pth\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ControlLDM(\n",
       "  (model): DiffusionWrapper(\n",
       "    (diffusion_model): ControlledUnetModel(\n",
       "      (time_embed): Sequential(\n",
       "        (0): Linear(in_features=320, out_features=1280, bias=True)\n",
       "        (1): SiLU()\n",
       "        (2): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "      )\n",
       "      (input_blocks): ModuleList(\n",
       "        (0): TimestepEmbedSequential(\n",
       "          (0): Conv2d(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (1-2): 2 x TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=320, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Identity()\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 320, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=320, out_features=320, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_k): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_v): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=320, out_features=2560, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=1280, out_features=320, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=320, out_features=320, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (3): TimestepEmbedSequential(\n",
       "          (0): Downsample(\n",
       "            (op): Conv2d(320, 320, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (4): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(320, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(320, 640, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (5): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Identity()\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (6): TimestepEmbedSequential(\n",
       "          (0): Downsample(\n",
       "            (op): Conv2d(640, 640, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (7): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(640, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(640, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (8): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Identity()\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (9): TimestepEmbedSequential(\n",
       "          (0): Downsample(\n",
       "            (op): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (10-11): 2 x TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Identity()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (middle_block): TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Identity()\n",
       "        )\n",
       "        (1): SpatialTransformer(\n",
       "          (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "          (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (transformer_blocks): ModuleList(\n",
       "            (0): BasicTransformerBlock(\n",
       "              (attn1): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (ff): FeedForward(\n",
       "                (net): Sequential(\n",
       "                  (0): GEGLU(\n",
       "                    (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                  )\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                  (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                )\n",
       "              )\n",
       "              (attn2): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "        (2): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Identity()\n",
       "        )\n",
       "      )\n",
       "      (output_blocks): ModuleList(\n",
       "        (0-1): 2 x TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 2560, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(2560, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(2560, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (2): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 2560, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(2560, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(2560, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): Upsample(\n",
       "            (conv): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (3-4): 2 x TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 2560, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(2560, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(2560, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (5): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1920, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(1920, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(1920, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "          (2): Upsample(\n",
       "            (conv): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (6): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1920, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(1920, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(1920, 640, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (7): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(1280, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(1280, 640, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (8): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 960, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(960, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(960, 640, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "          )\n",
       "          (2): Upsample(\n",
       "            (conv): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (9): TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 960, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(960, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=320, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(960, 320, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 320, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=320, out_features=320, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_k): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_v): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=320, out_features=2560, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=1280, out_features=320, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=320, out_features=320, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (10-11): 2 x TimestepEmbedSequential(\n",
       "          (0): ResBlock(\n",
       "            (in_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Conv2d(640, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (h_upd): Identity()\n",
       "            (x_upd): Identity()\n",
       "            (emb_layers): Sequential(\n",
       "              (0): SiLU()\n",
       "              (1): Linear(in_features=1280, out_features=320, bias=True)\n",
       "            )\n",
       "            (out_layers): Sequential(\n",
       "              (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "              (1): SiLU()\n",
       "              (2): Dropout(p=0, inplace=False)\n",
       "              (3): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "            (skip_connection): Conv2d(640, 320, kernel_size=(1, 1), stride=(1, 1))\n",
       "          )\n",
       "          (1): SpatialTransformer(\n",
       "            (norm): GroupNorm(32, 320, eps=1e-06, affine=True)\n",
       "            (proj_in): Linear(in_features=320, out_features=320, bias=True)\n",
       "            (transformer_blocks): ModuleList(\n",
       "              (0): BasicTransformerBlock(\n",
       "                (attn1): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_k): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_v): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (ff): FeedForward(\n",
       "                  (net): Sequential(\n",
       "                    (0): GEGLU(\n",
       "                      (proj): Linear(in_features=320, out_features=2560, bias=True)\n",
       "                    )\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                    (2): Linear(in_features=1280, out_features=320, bias=True)\n",
       "                  )\n",
       "                )\n",
       "                (attn2): MemoryEfficientCrossAttention(\n",
       "                  (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                  (to_k): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                  (to_v): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                  (to_out): Sequential(\n",
       "                    (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                    (1): Dropout(p=0.0, inplace=False)\n",
       "                  )\n",
       "                )\n",
       "                (norm1): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm2): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "                (norm3): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "            (proj_out): Linear(in_features=320, out_features=320, bias=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (out): Sequential(\n",
       "        (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "        (1): SiLU()\n",
       "        (2): Conv2d(320, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (first_stage_model): AutoencoderKL(\n",
       "    (encoder): Encoder(\n",
       "      (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (down): ModuleList(\n",
       "        (0): Module(\n",
       "          (block): ModuleList(\n",
       "            (0-1): 2 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (downsample): Downsample(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
       "          )\n",
       "        )\n",
       "        (1): Module(\n",
       "          (block): ModuleList(\n",
       "            (0): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (nin_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            )\n",
       "            (1): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (downsample): Downsample(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
       "          )\n",
       "        )\n",
       "        (2): Module(\n",
       "          (block): ModuleList(\n",
       "            (0): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (nin_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            )\n",
       "            (1): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (downsample): Downsample(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2))\n",
       "          )\n",
       "        )\n",
       "        (3): Module(\n",
       "          (block): ModuleList(\n",
       "            (0-1): 2 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "        )\n",
       "      )\n",
       "      (mid): Module(\n",
       "        (block_1): ResnetBlock(\n",
       "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (attn_1): MemoryEfficientAttnBlock(\n",
       "          (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (block_2): ResnetBlock(\n",
       "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "      (conv_out): Conv2d(512, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    )\n",
       "    (decoder): Decoder(\n",
       "      (conv_in): Conv2d(4, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (mid): Module(\n",
       "        (block_1): ResnetBlock(\n",
       "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (attn_1): MemoryEfficientAttnBlock(\n",
       "          (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (block_2): ResnetBlock(\n",
       "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (up): ModuleList(\n",
       "        (0): Module(\n",
       "          (block): ModuleList(\n",
       "            (0): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (nin_shortcut): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            )\n",
       "            (1-2): 2 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "        )\n",
       "        (1): Module(\n",
       "          (block): ModuleList(\n",
       "            (0): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (nin_shortcut): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            )\n",
       "            (1-2): 2 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (upsample): Upsample(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "        (2-3): 2 x Module(\n",
       "          (block): ModuleList(\n",
       "            (0-2): 3 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (upsample): Upsample(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (norm_out): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "      (conv_out): Conv2d(128, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    )\n",
       "    (loss): Identity()\n",
       "    (quant_conv): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (post_quant_conv): Conv2d(4, 4, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (cond_stage_model): FrozenOpenCLIPEmbedder(\n",
       "    (model): CLIP(\n",
       "      (transformer): Transformer(\n",
       "        (resblocks): ModuleList(\n",
       "          (0-23): 24 x ResidualAttentionBlock(\n",
       "            (ln_1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (attn): MultiheadAttention(\n",
       "              (out_proj): NonDynamicallyQuantizableLinear(in_features=1024, out_features=1024, bias=True)\n",
       "            )\n",
       "            (ls_1): Identity()\n",
       "            (ln_2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (mlp): Sequential(\n",
       "              (c_fc): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "              (gelu): GELU(approximate='none')\n",
       "              (c_proj): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "            )\n",
       "            (ls_2): Identity()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (token_embedding): Embedding(49408, 1024)\n",
       "      (ln_final): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (control_model): ControlNet(\n",
       "    (time_embed): Sequential(\n",
       "      (0): Linear(in_features=320, out_features=1280, bias=True)\n",
       "      (1): SiLU()\n",
       "      (2): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "    )\n",
       "    (input_blocks): ModuleList(\n",
       "      (0): TimestepEmbedSequential(\n",
       "        (0): Conv2d(8, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "      (1-2): 2 x TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=320, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Identity()\n",
       "        )\n",
       "        (1): SpatialTransformer(\n",
       "          (norm): GroupNorm(32, 320, eps=1e-06, affine=True)\n",
       "          (proj_in): Linear(in_features=320, out_features=320, bias=True)\n",
       "          (transformer_blocks): ModuleList(\n",
       "            (0): BasicTransformerBlock(\n",
       "              (attn1): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                (to_k): Linear(in_features=320, out_features=320, bias=False)\n",
       "                (to_v): Linear(in_features=320, out_features=320, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (ff): FeedForward(\n",
       "                (net): Sequential(\n",
       "                  (0): GEGLU(\n",
       "                    (proj): Linear(in_features=320, out_features=2560, bias=True)\n",
       "                  )\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                  (2): Linear(in_features=1280, out_features=320, bias=True)\n",
       "                )\n",
       "              )\n",
       "              (attn2): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=320, out_features=320, bias=False)\n",
       "                (to_k): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                (to_v): Linear(in_features=1024, out_features=320, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=320, out_features=320, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (norm1): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm2): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm3): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (proj_out): Linear(in_features=320, out_features=320, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (3): TimestepEmbedSequential(\n",
       "        (0): Downsample(\n",
       "          (op): Conv2d(320, 320, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (4): TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 320, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(320, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Conv2d(320, 640, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (1): SpatialTransformer(\n",
       "          (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "          (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "          (transformer_blocks): ModuleList(\n",
       "            (0): BasicTransformerBlock(\n",
       "              (attn1): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (ff): FeedForward(\n",
       "                (net): Sequential(\n",
       "                  (0): GEGLU(\n",
       "                    (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                  )\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                  (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                )\n",
       "              )\n",
       "              (attn2): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (5): TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=640, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Identity()\n",
       "        )\n",
       "        (1): SpatialTransformer(\n",
       "          (norm): GroupNorm(32, 640, eps=1e-06, affine=True)\n",
       "          (proj_in): Linear(in_features=640, out_features=640, bias=True)\n",
       "          (transformer_blocks): ModuleList(\n",
       "            (0): BasicTransformerBlock(\n",
       "              (attn1): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_k): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_v): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (ff): FeedForward(\n",
       "                (net): Sequential(\n",
       "                  (0): GEGLU(\n",
       "                    (proj): Linear(in_features=640, out_features=5120, bias=True)\n",
       "                  )\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                  (2): Linear(in_features=2560, out_features=640, bias=True)\n",
       "                )\n",
       "              )\n",
       "              (attn2): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=640, out_features=640, bias=False)\n",
       "                (to_k): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                (to_v): Linear(in_features=1024, out_features=640, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=640, out_features=640, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (proj_out): Linear(in_features=640, out_features=640, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (6): TimestepEmbedSequential(\n",
       "        (0): Downsample(\n",
       "          (op): Conv2d(640, 640, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (7): TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 640, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(640, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Conv2d(640, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (1): SpatialTransformer(\n",
       "          (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "          (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (transformer_blocks): ModuleList(\n",
       "            (0): BasicTransformerBlock(\n",
       "              (attn1): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (ff): FeedForward(\n",
       "                (net): Sequential(\n",
       "                  (0): GEGLU(\n",
       "                    (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                  )\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                  (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                )\n",
       "              )\n",
       "              (attn2): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (8): TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Identity()\n",
       "        )\n",
       "        (1): SpatialTransformer(\n",
       "          (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "          (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (transformer_blocks): ModuleList(\n",
       "            (0): BasicTransformerBlock(\n",
       "              (attn1): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (ff): FeedForward(\n",
       "                (net): Sequential(\n",
       "                  (0): GEGLU(\n",
       "                    (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                  )\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                  (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "                )\n",
       "              )\n",
       "              (attn2): MemoryEfficientCrossAttention(\n",
       "                (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "                (to_out): Sequential(\n",
       "                  (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (1): Dropout(p=0.0, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (9): TimestepEmbedSequential(\n",
       "        (0): Downsample(\n",
       "          (op): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (10-11): 2 x TimestepEmbedSequential(\n",
       "        (0): ResBlock(\n",
       "          (in_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (h_upd): Identity()\n",
       "          (x_upd): Identity()\n",
       "          (emb_layers): Sequential(\n",
       "            (0): SiLU()\n",
       "            (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          )\n",
       "          (out_layers): Sequential(\n",
       "            (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "            (1): SiLU()\n",
       "            (2): Dropout(p=0, inplace=False)\n",
       "            (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          )\n",
       "          (skip_connection): Identity()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (zero_convs): ModuleList(\n",
       "      (0-3): 4 x TimestepEmbedSequential(\n",
       "        (0): Conv2d(320, 320, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (4-6): 3 x TimestepEmbedSequential(\n",
       "        (0): Conv2d(640, 640, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "      (7-11): 5 x TimestepEmbedSequential(\n",
       "        (0): Conv2d(1280, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "    )\n",
       "    (middle_block): TimestepEmbedSequential(\n",
       "      (0): ResBlock(\n",
       "        (in_layers): Sequential(\n",
       "          (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "          (1): SiLU()\n",
       "          (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (h_upd): Identity()\n",
       "        (x_upd): Identity()\n",
       "        (emb_layers): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "        (out_layers): Sequential(\n",
       "          (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "          (1): SiLU()\n",
       "          (2): Dropout(p=0, inplace=False)\n",
       "          (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (skip_connection): Identity()\n",
       "      )\n",
       "      (1): SpatialTransformer(\n",
       "        (norm): GroupNorm(32, 1280, eps=1e-06, affine=True)\n",
       "        (proj_in): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        (transformer_blocks): ModuleList(\n",
       "          (0): BasicTransformerBlock(\n",
       "            (attn1): MemoryEfficientCrossAttention(\n",
       "              (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "              (to_k): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "              (to_v): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "              (to_out): Sequential(\n",
       "                (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                (1): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (ff): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): GEGLU(\n",
       "                  (proj): Linear(in_features=1280, out_features=10240, bias=True)\n",
       "                )\n",
       "                (1): Dropout(p=0.0, inplace=False)\n",
       "                (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (attn2): MemoryEfficientCrossAttention(\n",
       "              (to_q): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1280, bias=False)\n",
       "              (to_out): Sequential(\n",
       "                (0): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                (1): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (norm1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            (norm2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            (norm3): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "          )\n",
       "        )\n",
       "        (proj_out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "      )\n",
       "      (2): ResBlock(\n",
       "        (in_layers): Sequential(\n",
       "          (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "          (1): SiLU()\n",
       "          (2): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (h_upd): Identity()\n",
       "        (x_upd): Identity()\n",
       "        (emb_layers): Sequential(\n",
       "          (0): SiLU()\n",
       "          (1): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "        (out_layers): Sequential(\n",
       "          (0): GroupNorm32(32, 1280, eps=1e-05, affine=True)\n",
       "          (1): SiLU()\n",
       "          (2): Dropout(p=0, inplace=False)\n",
       "          (3): Conv2d(1280, 1280, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (skip_connection): Identity()\n",
       "      )\n",
       "    )\n",
       "    (middle_block_out): TimestepEmbedSequential(\n",
       "      (0): Conv2d(1280, 1280, kernel_size=(1, 1), stride=(1, 1))\n",
       "    )\n",
       "  )\n",
       "  (preprocess_model): SwinIR(\n",
       "    (conv_first): Sequential(\n",
       "      (0): PixelUnshuffle(downscale_factor=8)\n",
       "      (1): Conv2d(192, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    )\n",
       "    (patch_embed): PatchEmbed(\n",
       "      (norm): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (patch_unembed): PatchUnEmbed()\n",
       "    (pos_drop): Dropout(p=0.0, inplace=False)\n",
       "    (layers): ModuleList(\n",
       "      (0): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): Identity()\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.002)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.004)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.006)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.009)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.011)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (1): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.013)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.015)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.017)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.019)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.021)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.023)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (2): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.026)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.028)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.030)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.032)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.034)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.036)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (3): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.038)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.040)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.043)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.045)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.047)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.049)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (4): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.051)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.053)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.055)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.057)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.060)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.062)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (5): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.064)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.066)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.068)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.070)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.072)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.074)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (6): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.077)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.079)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.081)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.083)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.085)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.087)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "      (7): RSTB(\n",
       "        (residual_group): BasicLayer(\n",
       "          dim=180, input_resolution=(64, 64), depth=6\n",
       "          (blocks): ModuleList(\n",
       "            (0): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.089)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (1): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.091)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (2): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.094)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (3): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.096)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (4): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=0, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.098)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (5): SwinTransformerBlock(\n",
       "              dim=180, input_resolution=(64, 64), num_heads=6, window_size=8, shift_size=4, mlp_ratio=2\n",
       "              (norm1): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): WindowAttention(\n",
       "                dim=180, window_size=(8, 8), num_heads=6\n",
       "                (qkv): Linear(in_features=180, out_features=540, bias=True)\n",
       "                (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "                (proj): Linear(in_features=180, out_features=180, bias=True)\n",
       "                (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "                (softmax): Softmax(dim=-1)\n",
       "              )\n",
       "              (drop_path): DropPath(drop_prob=0.100)\n",
       "              (norm2): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): Mlp(\n",
       "                (fc1): Linear(in_features=180, out_features=360, bias=True)\n",
       "                (act): GELU(approximate='none')\n",
       "                (fc2): Linear(in_features=360, out_features=180, bias=True)\n",
       "                (drop): Dropout(p=0.0, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        (patch_embed): PatchEmbed()\n",
       "        (patch_unembed): PatchUnEmbed()\n",
       "      )\n",
       "    )\n",
       "    (norm): LayerNorm((180,), eps=1e-05, elementwise_affine=True)\n",
       "    (conv_after_body): Conv2d(180, 180, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (conv_before_upsample): Sequential(\n",
       "      (0): Conv2d(180, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_up1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (conv_up2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (conv_up3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (conv_hr): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (conv_last): Conv2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (lrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "  )\n",
       "  (cond_encoder): Sequential(\n",
       "    (encoder): Encoder(\n",
       "      (conv_in): Conv2d(3, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (down): ModuleList(\n",
       "        (0): Module(\n",
       "          (block): ModuleList(\n",
       "            (0-1): 2 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (downsample): Downsample(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2))\n",
       "          )\n",
       "        )\n",
       "        (1): Module(\n",
       "          (block): ModuleList(\n",
       "            (0): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 128, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (nin_shortcut): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            )\n",
       "            (1): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (downsample): Downsample(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2))\n",
       "          )\n",
       "        )\n",
       "        (2): Module(\n",
       "          (block): ModuleList(\n",
       "            (0): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 256, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (nin_shortcut): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            )\n",
       "            (1): ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "          (downsample): Downsample(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2))\n",
       "          )\n",
       "        )\n",
       "        (3): Module(\n",
       "          (block): ModuleList(\n",
       "            (0-1): 2 x ResnetBlock(\n",
       "              (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "              (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            )\n",
       "          )\n",
       "          (attn): ModuleList()\n",
       "        )\n",
       "      )\n",
       "      (mid): Module(\n",
       "        (block_1): ResnetBlock(\n",
       "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "        (attn_1): MemoryEfficientAttnBlock(\n",
       "          (norm): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (q): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (k): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (v): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (proj_out): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (block_2): ResnetBlock(\n",
       "          (norm1): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (norm2): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (norm_out): GroupNorm(32, 512, eps=1e-06, affine=True)\n",
       "      (conv_out): Conv2d(512, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    )\n",
       "    (quant_conv): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from utils.image import auto_resize, pad\n",
    "import einops\n",
    "from utils.common import instantiate_from_config, load_state_dict\n",
    "from omegaconf import OmegaConf\n",
    "from model.cldm import ControlLDM\n",
    "from model.spaced_sampler import SpacedSampler\n",
    "model: ControlLDM = instantiate_from_config(OmegaConf.load('./config/cldm.yaml'))\n",
    "ckpt_swinir='./weights/general_full_v1.ckpt'\n",
    "ckpt_net='./weights/cheng_small.pth.tar'\n",
    "load_state_dict(model, torch.load(ckpt_swinir, map_location=\"cpu\"), strict=True)\n",
    "model.freeze()\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fe59143",
   "metadata": {},
   "source": [
    "function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a2fe5c35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.transform import resize\n",
    "\n",
    "import shutil\n",
    "def clear_directory(directory_path):\n",
    "    for file in os.listdir(directory_path):\n",
    "        file_path = os.path.join(directory_path, file)\n",
    "        if os.path.isdir(file_path):\n",
    "            # 如果是子文件夹，递归调用删除文件\n",
    "            \n",
    "            for root, dirs, files in os.walk(file_path):\n",
    "                for file in files:\n",
    "                    os.remove(os.path.join(root, file))\n",
    "            \n",
    "def to_block(img,grid=32,level=8):\n",
    "    g_w=int(img.size[0]/grid)\n",
    "    g_h=int(img.size[1]/grid)\n",
    "    img_resize=img.resize((g_w, g_h))\n",
    "    img_np=np.floor(np.array(img_resize)/level)*level\n",
    "    img_np=img_np.astype (np.uint8)\n",
    "    img_reference = Image.fromarray(img_np).resize(img.size)\n",
    "    return img_reference, g_w, g_h\n",
    "\n",
    "def divide_integer(num, n):\n",
    "    quotient = num // n  # 整数除法，计算商\n",
    "    remainder = num % n  # 取余数\n",
    "    result = [quotient] * n  # 创建一个包含n个quotient的列表\n",
    "\n",
    "    # 将余数平均分配给前几个数\n",
    "    for i in range(remainder):\n",
    "        result[i] += 1\n",
    "\n",
    "    return result\n",
    "def mask_block(mask,num=8,level=0.35):\n",
    "    tmp=resize(mask, (num, num), mode='reflect')\n",
    "    tmp[tmp>level]=255\n",
    "    tmp[tmp<=level]=0\n",
    "    rp_mat_0=np.array(divide_integer(mask.shape[0], num),dtype='int')\n",
    "    rp_mat_1=np.array(divide_integer(mask.shape[1], num),dtype='int')\n",
    "    return tmp.repeat(rp_mat_1,axis=1).repeat(rp_mat_0,axis=0)\n",
    "\n",
    "def image_paddle_in(image, num=32):\n",
    "    # 计算扩充后的宽度和高度\n",
    "    new_width = ((image.width-1) // num + 1) * num\n",
    "    new_height = ((image.height-1) // num + 1) * num\n",
    "\n",
    "    # 创建一个新的扩充后的图像，用空值填充\n",
    "    new_image = Image.new(\"RGB\", (new_width, new_height), (0, 0, 0))\n",
    "\n",
    "    # 将原始图像粘贴到扩充后的图像左上角\n",
    "    new_image.paste(image, (0, 0))\n",
    "    return new_image,image.width,image.height\n",
    "\n",
    "def image_paddle_out(image, old_width, old_height):\n",
    "\n",
    "    return image.crop((0,0,old_width,old_height))\n",
    "\n",
    "def compute_bpp(out_net):\n",
    "    size = out_net['x_hat'].size()\n",
    "    num_pixels = size[0] * size[2] * size[3]\n",
    "    return sum(torch.log(likelihoods).sum() / (-math.log(2) * num_pixels)\n",
    "              for likelihoods in out_net['likelihoods'].values()).item()\n",
    "\n",
    "def clip_map(img,texts,mask_num=8):\n",
    "    image = preprocess(img).unsqueeze(0).to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        # CLIP architecture surgery acts on the image encoder\n",
    "        cv2_img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)\n",
    "        image_features = seg_model.encode_image(image)\n",
    "        image_features = image_features / image_features.norm(dim=1, keepdim=True)\n",
    "\n",
    "        # Prompt ensemble for text features with normalization\n",
    "        text_features = clip.encode_text_with_prompt_ensemble(seg_model, texts, device)\n",
    "\n",
    "        # Extract redundant features from an empty string\n",
    "        redundant_features = clip.encode_text_with_prompt_ensemble(seg_model, [\"\"], device)\n",
    "\n",
    "        # Apply feature surgery for single text\n",
    "        similarity = clip.clip_feature_surgery(image_features, text_features, redundant_features)\n",
    "        similarity_map = clip.get_similarity_map(similarity[:, 1:, :], cv2_img.shape[:2])\n",
    "        \n",
    "        mask_0=(similarity_map[0,:,:,0].cpu().numpy() * 255).astype('uint8')\n",
    "        mask_1=(similarity_map[0,:,:,1].cpu().numpy() * 255).astype('uint8')\n",
    "        mask_2=(similarity_map[0,:,:,2].cpu().numpy() * 255).astype('uint8')\n",
    "        mask_0=Image.fromarray(mask_block(mask_0,num=mask_num))\n",
    "        mask_1=Image.fromarray(mask_block(mask_1,num=mask_num))\n",
    "        mask_2=Image.fromarray(mask_block(mask_2,num=mask_num))\n",
    "        return mask_0,mask_1,mask_2\n",
    "    \n",
    "def sr_pipe(img_reference,positive_prompt=\"\",cfg=1.0,steps=40,res=512):\n",
    "    lq_resized = auto_resize(img_reference, res)\n",
    "    x = pad(np.array(lq_resized), scale=64)\n",
    "    sampler = SpacedSampler(model, var_type=\"fixed_small\")\n",
    "    control = torch.tensor(np.stack([x]) / 255.0, dtype=torch.float32, device=device).clamp_(0, 1)\n",
    "    control = einops.rearrange(control, \"n h w c -> n c h w\").contiguous()\n",
    "    control = model.preprocess_model(control)\n",
    "    model.control_scales = [1] * 13\n",
    "    height, width = control.size(-2), control.size(-1)\n",
    "    shape = (1, 4, height // 8, width // 8)\n",
    "    x_T = torch.randn(shape, device=device, dtype=torch.float32)\n",
    "    samples = sampler.sample(\n",
    "            steps=steps, shape=shape, cond_img=control,\n",
    "            positive_prompt=positive_prompt, negative_prompt=\"Blurry, Low Quality\", x_T=x_T,\n",
    "            cfg_scale=cfg, cond_fn=None,\n",
    "            color_fix_type=\"wavelet\"\n",
    "        )\n",
    "    x_samples = samples.clamp(0, 1)\n",
    "    x_samples = (einops.rearrange(x_samples, \"b c h w -> b h w c\") * 255).cpu().numpy().clip(0, 255).astype(np.uint8)\n",
    "\n",
    "    img_sr=Image.fromarray(x_samples[0][:,0:lq_resized.size[0],:])\n",
    "    return img_sr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "276737eb",
   "metadata": {},
   "source": [
    "# Encode and Decode Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "91131e57",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bpp=0.0469675537109375\n",
      "timesteps used in spaced sampler: \n",
      "\t[0, 26, 51, 77, 102, 128, 154, 179, 205, 231, 256, 282, 307, 333, 359, 384, 410, 435, 461, 487, 512, 538, 564, 589, 615, 640, 666, 692, 717, 743, 768, 794, 820, 845, 871, 897, 922, 948, 973, 999]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Spaced Sampler: 100%|██████████████████████████████████████████████████████████████████| 40/40 [00:17<00:00,  2.28it/s]\n"
     ]
    }
   ],
   "source": [
    "mode='ref'\n",
    "using_map=False\n",
    "\n",
    "if mode=='net':\n",
    "    checkpoint_path=ckpt_net\n",
    "    checkpoint = torch.load(checkpoint_path, map_location=\"cpu\")\n",
    "    state_dict = checkpoint\n",
    "    for key in [\"network\", \"state_dict\", \"model_state_dict\"]:\n",
    "        if key in checkpoint:\n",
    "            state_dict = checkpoint[key]\n",
    "    arch='cheng2020-attn'\n",
    "    from compressai.zoo.image import model_architectures as architectures\n",
    "    model_cls = architectures[arch]\n",
    "    comp_net = model_cls.from_state_dict(state_dict).eval().to(device)\n",
    "elif mode=='ref':\n",
    "    ref_path='./ref/example-reference.png'\n",
    "    ref_bpp=0.0421\n",
    "elif mode=='pixel':\n",
    "    block_level=3\n",
    "    block_num_min=32\n",
    "\n",
    "    \n",
    "mask_num=8\n",
    "res=1024\n",
    "\n",
    "import pandas as pd\n",
    "df = pd.read_csv(\"GPT-example.csv\")\n",
    "\n",
    "image_path='./example-input.png'\n",
    "\n",
    "\n",
    "img = Image.open(image_path).convert('RGB')\n",
    "\n",
    "########################## Encoder ##########################\n",
    "\n",
    "#GPT prompt processing\n",
    "prompt=df['prompt'][0]\n",
    "prompt_list=prompt.split('\\n')\n",
    "prompt_list = [element for element in prompt_list if element != '']\n",
    "name_0,name_1,name_2=prompt_list[0].split('.')[0].split(',')\n",
    "detail_0,detail_1,detail_2=prompt_list[1],prompt_list[2],prompt_list[3]\n",
    "detail_all=prompt_list[4]\n",
    "mask_0,mask_1,mask_2=clip_map(img,[name_0,name_1,name_2],mask_num)\n",
    "\n",
    "#reference\n",
    "if mode=='pixel':\n",
    "    old_width, old_height=img.size\n",
    "    block_num=max(int(max(old_width, old_height)/16),block_num_min)\n",
    "    img_reference=to_block(img,block_num,2**block_level)\n",
    "\n",
    "    b_image=block_level*block_num**2\n",
    "elif mode=='net':\n",
    "    img, old_width, old_height = image_paddle_in(img, 64)\n",
    "    x = transforms.ToTensor()(img).unsqueeze(0).to(device)\n",
    "    with torch.no_grad():\n",
    "        out_net = comp_net.forward(x)\n",
    "    out_net['x_hat'].clamp_(0, 1)\n",
    "    img_reference = transforms.ToPILImage()(out_net['x_hat'].squeeze().cpu())\n",
    "    img_reference = image_paddle_out(img_reference, old_width, old_height)\n",
    "    b_image=compute_bpp(out_net)*img.size[0]*img.size[1]\n",
    "elif mode=='ref':\n",
    "    old_width, old_height=img.size\n",
    "    img_reference = Image.open(ref_path).convert('RGB')\n",
    "    b_image=ref_bpp*img.size[0]*img.size[1]\n",
    "\n",
    "\n",
    "#############################################################\n",
    "\n",
    "########################## Decoder ##########################\n",
    "num_inference_steps=40\n",
    "exag=1024/max(img_reference.size)\n",
    "height=int(img_reference.size[1]*exag/8)*8\n",
    "width=int(img_reference.size[0]*exag/8)*8\n",
    "\n",
    "#    img_reference=img_reference.resize([width,height])\n",
    "mask_0=mask_0.resize([width,height])\n",
    "mask_1=mask_1.resize([width,height])\n",
    "mask_2=mask_2.resize([width,height])\n",
    "mask_all=Image.new(\"RGB\", img_reference.size, (255, 255, 255))\n",
    "\n",
    "image = img_reference \n",
    "\n",
    "if using_map:\n",
    "    b_mask=mask_num*mask_num*3\n",
    "    b_word=(len(detail_0)+len(detail_1)+len(detail_2)+len(detail_all))*8\n",
    "    bpp=(b_image+b_mask+b_word)/(img.size[0]*img.size[1])\n",
    "    print('bpp='+str(bpp))\n",
    "    \n",
    "    image_tmp = sr_pipe(image,positive_prompt=detail_0,cfg=3.5,steps=3,res=res)\n",
    "    image = ImageChops.add(ImageChops.multiply(image_tmp,mask_0.convert(\"RGB\")), \n",
    "                          ImageChops.multiply(image,Image.fromarray(255-np.array(mask_0)).convert(\"RGB\"))\n",
    "                          ).resize((old_width, old_height))\n",
    "#    image.resize((old_width, old_height)).save(output_folder+'Mask0/'+image_name)\n",
    "\n",
    "    image_tmp = sr_pipe(image,positive_prompt=detail_1,cfg=3.5,steps=3,res=res)\n",
    "    image = ImageChops.add(ImageChops.multiply(image_tmp,mask_1.convert(\"RGB\")), \n",
    "                          ImageChops.multiply(image,Image.fromarray(255-np.array(mask_1)).convert(\"RGB\"))\n",
    "                          ).resize((old_width, old_height))\n",
    "#    image.resize((old_width, old_height)).save(output_folder+'Mask1/'+image_name)\n",
    "                          \n",
    "    image_tmp = sr_pipe(image,positive_prompt=detail_2,cfg=3.5,steps=3,res=res)\n",
    "    image = ImageChops.add(ImageChops.multiply(image_tmp,mask_2.convert(\"RGB\")), \n",
    "                          ImageChops.multiply(image,Image.fromarray(255-np.array(mask_2)).convert(\"RGB\"))\n",
    "                          ).resize((old_width, old_height))\n",
    "#    image.resize((old_width, old_height)).save(output_folder+'Mask2/'+image_name)\n",
    "    \n",
    "    image = sr_pipe(image,positive_prompt=detail_all,cfg=7,steps=40,res=res).resize((old_width, old_height))\n",
    "#    image.resize((old_width, old_height)).save(output_folder+'SR/'+image_name)\n",
    "    \n",
    "else:  \n",
    "    b_word=(len(detail_all))*8\n",
    "    bpp=(b_image+b_word)/(img.size[0]*img.size[1])\n",
    "    print('bpp='+str(bpp))\n",
    "\n",
    "    image = sr_pipe(image,positive_prompt=detail_all,cfg=7,steps=40,res=res)\n",
    "    image = image.resize((old_width, old_height))\n",
    "    \n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}