StyleGAN_XL_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G9fueCPLFKin"
      },
      "source": [
        "GitHub<br>\n",
        "https://github.com/autonomousvision/stylegan_xl<br>\n",
        "論文<br>\n",
        "https://arxiv.org/abs/2202.00273<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/StyleGAN_XL_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "894c0QeEFKit"
      },
      "source": [
        "# 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TmL3LoMMGcNF"
      },
      "source": [
        "## Githubからソースコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tthUF7D6FKiu",
        "outputId": "9465d383-bf6e-4fc6-babc-81c57a589a97"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/autonomousvision/stylegan_xl\n",
        "!git clone https://github.com/openai/CLIP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2-2Fh1sOHdnn"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AP8SsHYvHfkR",
        "outputId": "58a78103-e28e-4ba7-cb16-87ce1dfa1bc5"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!pip install -e ./CLIP\n",
        "!pip install einops ninja\n",
        "!pip install timm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gv7uc__wHuWd"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yayN7u44HwTZ",
        "outputId": "1f4ef424-094e-4742-d580-f1e2371f9f8c"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "import sys\n",
        "sys.path.append('./CLIP')\n",
        "sys.path.append('./stylegan_xl')\n",
        "\n",
        "import io\n",
        "import os, time, glob\n",
        "import pickle\n",
        "import shutil\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import requests\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.transforms.functional as TF\n",
        "import clip\n",
        "import unicodedata\n",
        "import re\n",
        "from tqdm.notebook import tqdm\n",
        "from torchvision.transforms import Compose, Resize, ToTensor, Normalize\n",
        "from IPython.display import display\n",
        "from einops import rearrange\n",
        "from google.colab import files\n",
        "import dnnlib\n",
        "import legacy\n",
        "from base64 import b64encode\n",
        "from IPython.display import HTML\n",
        "\n",
        "device = torch.device('cuda:0')\n",
        "print('Using device:', device, file=sys.stderr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ZKpJZIkMfGx"
      },
      "outputs": [],
      "source": [
        "def fetch(url_or_path):\n",
        "    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):\n",
        "        r = requests.get(url_or_path)\n",
        "        r.raise_for_status()\n",
        "        fd = io.BytesIO()\n",
        "        fd.write(r.content)\n",
        "        fd.seek(0)\n",
        "        return fd\n",
        "    return open(url_or_path, 'rb')\n",
        "\n",
        "def fetch_model(url_or_path):\n",
        "    !wget -c '{url_or_path}'\n",
        "\n",
        "def slugify(value, allow_unicode=False):\n",
        "    \"\"\"\n",
        "    Taken from https://github.com/django/django/blob/master/django/utils/text.py\n",
        "    Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated\n",
        "    dashes to single dashes. Remove characters that aren't alphanumerics,\n",
        "    underscores, or hyphens. Convert to lowercase. Also strip leading and\n",
        "    trailing whitespace, dashes, and underscores.\n",
        "    \"\"\"\n",
        "    value = str(value)\n",
        "    if allow_unicode:\n",
        "        value = unicodedata.normalize('NFKC', value)\n",
        "    else:\n",
        "        value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')\n",
        "    value = re.sub(r'[^\\w\\s-]', '', value.lower())\n",
        "    return re.sub(r'[-\\s]+', '-', value).strip('-_')\n",
        "\n",
        "def norm1(prompt):\n",
        "    \"Normalize to the unit sphere.\"\n",
        "    return prompt / prompt.square().sum(dim=-1,keepdim=True).sqrt()\n",
        "\n",
        "def spherical_dist_loss(x, y):\n",
        "    x = F.normalize(x, dim=-1)\n",
        "    y = F.normalize(y, dim=-1)\n",
        "    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)\n",
        "\n",
        "def prompts_dist_loss(x, targets, loss):\n",
        "    if len(targets) == 1: # Keeps consitent results vs previous method for single objective guidance \n",
        "      return loss(x, targets[0])\n",
        "    distances = [loss(x, target) for target in targets]\n",
        "    return torch.stack(distances, dim=-1).sum(dim=-1)\n",
        "\n",
        "def embed_image(image):\n",
        "  n = image.shape[0]\n",
        "  cutouts = make_cutouts(image)\n",
        "  embeds = clip_model.embed_cutout(cutouts)\n",
        "  embeds = rearrange(embeds, '(cc n) c -> cc n c', n=n)\n",
        "  return embeds\n",
        "\n",
        "def embed_url(url):\n",
        "  image = Image.open(fetch(url)).convert('RGB')\n",
        "  return embed_image(TF.to_tensor(image).to(device).unsqueeze(0)).mean(0).squeeze(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vkhs5-BuIRFw"
      },
      "source": [
        "# クラス定義"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c6L_svf7IVxH",
        "outputId": "3282198e-5265-41d3-f6da-ea2bcd4edfc5"
      },
      "outputs": [],
      "source": [
        "class MakeCutouts(torch.nn.Module):\n",
        "    def __init__(self, cut_size, cutn, cut_pow=1.):\n",
        "        super().__init__()\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "        self.cut_pow = cut_pow\n",
        "\n",
        "    def forward(self, input):\n",
        "        sideY, sideX = input.shape[2:4]\n",
        "        max_size = min(sideX, sideY)\n",
        "        min_size = min(sideX, sideY, self.cut_size)\n",
        "        cutouts = []\n",
        "        for _ in range(self.cutn):\n",
        "            size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)\n",
        "            offsetx = torch.randint(0, sideX - size + 1, ())\n",
        "            offsety = torch.randint(0, sideY - size + 1, ())\n",
        "            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
        "            cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))\n",
        "        return torch.cat(cutouts)\n",
        "\n",
        "class CLIP(object):\n",
        "  def __init__(self):\n",
        "    clip_model = \"ViT-B/16\"\n",
        "    self.model, _ = clip.load(clip_model)\n",
        "    self.model = self.model.requires_grad_(False)\n",
        "    self.normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                                          std=[0.26862954, 0.26130258, 0.27577711])\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def embed_text(self, prompt):\n",
        "      \"Normalized clip text embedding.\"\n",
        "      return norm1(self.model.encode_text(clip.tokenize(prompt).to(device)).float())\n",
        "\n",
        "  def embed_cutout(self, image):\n",
        "      \"Normalized clip image embedding.\"\n",
        "      return norm1(self.model.encode_image(self.normalize(image)))\n",
        "  \n",
        "make_cutouts = MakeCutouts(224, 32, 0.5)\n",
        "clip_model = CLIP()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yLugChQtLSxe"
      },
      "outputs": [],
      "source": [
        "def fetch_model(url_or_path):\n",
        "  !wget -c '{url_or_path}'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKbUq5APIqJW"
      },
      "source": [
        "# モデル選択"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_6A5ALMrIsAC"
      },
      "outputs": [],
      "source": [
        "Model = 'Pokemon' #@param [\"Imagenet\", \"Pokemon\", \"FFHQ\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sd3w1cEaIxPN",
        "outputId": "ae46acd1-18da-421f-954b-cea2737a1919"
      },
      "outputs": [],
      "source": [
        "network_url = {\n",
        "    \"Imagenet\": \"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl\",\n",
        "    \"Pokemon\": \"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/pokemon256.pkl\",\n",
        "    \"FFHQ\": \"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/ffhq256.pkl\"\n",
        "}\n",
        "\n",
        "network_name = network_url[Model].split(\"/\")[-1]\n",
        "fetch_model(network_url[Model])\n",
        "\n",
        "with dnnlib.util.open_url(network_name) as f:\n",
        "  G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore\n",
        "\n",
        "zs = torch.randn([10000, G.mapping.z_dim], device=device)\n",
        "cs = torch.zeros([10000, G.mapping.c_dim], device=device)\n",
        "\n",
        "if G.mapping.c_dim != 0:\n",
        "  for i in range(cs.shape[0]):\n",
        "    cs[i,i//10]=1\n",
        "\n",
        "w_stds = G.mapping(zs, cs)\n",
        "w_stds = w_stds.reshape(10, 1000, G.num_ws, -1)\n",
        "w_stds=w_stds.std(0).mean(0)[0]\n",
        "w_all_classes_avg = G.mapping.w_avg.mean(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uyqUocZVI7vh"
      },
      "source": [
        "# Text prompt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AWxDndi1JBwg"
      },
      "outputs": [],
      "source": [
        "texts = \"psychedelic pikachu.\"#@param {type:\"string\"}\n",
        "steps = 500#@param {type:\"number\"}\n",
        "seed = 12#@param {type:\"number\"}\n",
        "\n",
        "!rm -rf ./outputs\n",
        "!rm -rf result.mp4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3FcAW1itI-zh"
      },
      "outputs": [],
      "source": [
        "if seed == -1:\n",
        "    seed = np.random.randint(0,9e9)\n",
        "    print(f\"Your random seed is: {seed}\")\n",
        "\n",
        "texts = [frase.strip() for frase in texts.split(\"|\") if frase]\n",
        "\n",
        "targets = [clip_model.embed_text(text) for text in texts]\n",
        "\n",
        "tf = Compose([\n",
        "  # Resize(224),\n",
        "  lambda x: torch.clamp((x+1)/2,min=0,max=1),\n",
        "])\n",
        "\n",
        "initial_batch=4 #actually that will be multiplied by initial_image_steps\n",
        "initial_image_steps=32"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zQUBTXAGJY2c"
      },
      "source": [
        "# Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oFyuKYcbJX68"
      },
      "outputs": [],
      "source": [
        "def run(out_dir):\n",
        "  torch.manual_seed(seed)\n",
        "  with torch.no_grad():\n",
        "    qs = []\n",
        "    losses = []\n",
        "    for _ in range(initial_image_steps):\n",
        "      a = torch.randn([initial_batch, 512], device=device)*0.4 + w_stds*0.4\n",
        "      q = ((a-w_all_classes_avg)/w_stds)\n",
        "      images = G.synthesis((q * w_stds + w_all_classes_avg).unsqueeze(1).repeat([1, G.num_ws, 1]))\n",
        "      embeds = embed_image(images.add(1).div(2))\n",
        "      loss = prompts_dist_loss(embeds, targets, spherical_dist_loss).mean(0)\n",
        "      i = torch.argmin(loss)\n",
        "      qs.append(q[i])\n",
        "      losses.append(loss[i])\n",
        "    qs = torch.stack(qs)\n",
        "    losses = torch.stack(losses)\n",
        "    # print(losses)\n",
        "    # print(losses.shape, qs.shape)\n",
        "    i = torch.argmin(losses)\n",
        "    q = qs[i].unsqueeze(0).repeat([G.num_ws, 1]).requires_grad_()\n",
        "\n",
        "\n",
        "  # Sampling loop\n",
        "  q_ema = q\n",
        "  print(q.shape)\n",
        "  opt = torch.optim.AdamW([q], lr=0.05, betas=(0., 0.999), weight_decay=0.025)\n",
        "  loop = tqdm(range(steps))\n",
        "  for i in loop:\n",
        "    opt.zero_grad()\n",
        "    w = q * w_stds\n",
        "    image = G.synthesis((q * w_stds + w_all_classes_avg)[None], noise_mode='const')\n",
        "    embed = embed_image(image.add(1).div(2))\n",
        "    loss = prompts_dist_loss(embed, targets, spherical_dist_loss).mean()\n",
        "    loss.backward()\n",
        "    opt.step()\n",
        "    loop.set_postfix(loss=loss.item(), q_magnitude=q.std().item())\n",
        "\n",
        "    q_ema = q_ema * 0.98 + q * 0.02\n",
        "    image = G.synthesis((q_ema * w_stds + w_all_classes_avg)[None], noise_mode='const')\n",
        "\n",
        "    if i % 50 == 0:\n",
        "      display(TF.to_pil_image(tf(image)[0]))\n",
        "      print(f\"Image {i}/{steps} | Current loss: {loss}\")\n",
        "    pil_image = TF.to_pil_image(image[0].add(1).div(2).clamp(0,1))\n",
        "    os.makedirs(f'{out_dir}', exist_ok=True)\n",
        "    pil_image.save(f'{out_dir}/{i:04}.jpg')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "6f2005f559e445b7aa3fd305eac4a52f",
            "677e4b4d7aeb49f998a746a0724db4db",
            "6f07c24009ad4a5f997d71e50fb29598",
            "90c2c3532b6649369f97c46129b1380b",
            "76c92f6ecdd44b14940f660f527fc2ff",
            "d4be30bfc76940afba9b81c54e322986",
            "0b1b12ac8383459685ab236d00d3facb",
            "be922992ae564492b8b4e8acadbc50e2",
            "809ae75b578e428f925904c8ecfeeaf2",
            "a85a7668d4364ae18542a4918bd3862e",
            "664b182a42344b0d98558acac8786ed2"
          ]
        },
        "id": "29b-tQYUJdPi",
        "outputId": "851c8ecf-d164-4411-c778-7544d9a497c2"
      },
      "outputs": [],
      "source": [
        "run('./outputs')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sn2U0yahJyzQ"
      },
      "source": [
        "# img to video"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vj5i9sPOJ1Bk",
        "outputId": "b5146b7d-c285-4efc-f84e-ebe14853c6c0"
      },
      "outputs": [],
      "source": [
        "!ffmpeg -r 60 -i 'outputs/%04d.jpg' -vcodec libx264 -crf 18 -pix_fmt yuv420p result.mp4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EyRWaafrKH7G"
      },
      "outputs": [],
      "source": [
        "def show_mp4(filename, width):\n",
        "  mp4 = open(filename + '.mp4', 'rb').read()\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "  display(HTML(\"\"\"\n",
        "  <video width=\"%d\" controls autoplay loop>\n",
        "    <source src=\"%s\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  \"\"\" % (width, data_url)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 536
        },
        "id": "CHVgvO_PKKAA",
        "outputId": "de8af38e-9b2a-4870-ac5b-b9c3e33b5382"
      },
      "outputs": [],
      "source": [
        "show_mp4(\"result\", width=514)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "StyleGAN_XL_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0b1b12ac8383459685ab236d00d3facb": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "664b182a42344b0d98558acac8786ed2": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "677e4b4d7aeb49f998a746a0724db4db": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d4be30bfc76940afba9b81c54e322986",
            "placeholder": "​",
            "style": "IPY_MODEL_0b1b12ac8383459685ab236d00d3facb",
            "value": "100%"
          }
        },
        "6f07c24009ad4a5f997d71e50fb29598": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_be922992ae564492b8b4e8acadbc50e2",
            "max": 500,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_809ae75b578e428f925904c8ecfeeaf2",
            "value": 500
          }
        },
        "6f2005f559e445b7aa3fd305eac4a52f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_677e4b4d7aeb49f998a746a0724db4db",
              "IPY_MODEL_6f07c24009ad4a5f997d71e50fb29598",
              "IPY_MODEL_90c2c3532b6649369f97c46129b1380b"
            ],
            "layout": "IPY_MODEL_76c92f6ecdd44b14940f660f527fc2ff"
          }
        },
        "76c92f6ecdd44b14940f660f527fc2ff": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "809ae75b578e428f925904c8ecfeeaf2": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "90c2c3532b6649369f97c46129b1380b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a85a7668d4364ae18542a4918bd3862e",
            "placeholder": "​",
            "style": "IPY_MODEL_664b182a42344b0d98558acac8786ed2",
            "value": " 500/500 [07:02&lt;00:00,  1.19it/s, loss=0.483, q_magnitude=0.992]"
          }
        },
        "a85a7668d4364ae18542a4918bd3862e": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "be922992ae564492b8b4e8acadbc50e2": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d4be30bfc76940afba9b81c54e322986": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}