{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "355UKMUQJxFd"
   },
   "source": [
    "# Scalable Diffusion Models with Transformer (DiT)\n",
    "\n",
    "This notebook samples from pre-trained DiT models. DiTs are class-conditional latent diffusion models trained on ImageNet that use transformers in place of U-Nets as the DDPM backbone. DiT outperforms all prior diffusion models on the ImageNet benchmarks.\n",
    "\n",
    "[Project Page](https://www.wpeebles.com/DiT) | [HuggingFace Space](https://huggingface.co/spaces/wpeebles/DiT) | [Paper](http://arxiv.org/abs/2212.09748) | [GitHub](github.com/facebookresearch/DiT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zJlgLkSaKn7u"
   },
   "source": [
    "# 1. Setup\n",
    "\n",
    "We recommend using GPUs (Runtime > Change runtime type > Hardware accelerator > GPU). Run this cell to clone the DiT GitHub repo and setup PyTorch. You only have to run this once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/facebookresearch/DiT.git\n",
    "import DiT, os\n",
    "os.chdir('DiT')\n",
    "os.environ['PYTHONPATH'] = '/env/python:/content/DiT'\n",
    "!pip install diffusers timm --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-02-21T02:55:56.417045Z",
     "iopub.status.busy": "2024-02-21T02:55:56.416754Z",
     "iopub.status.idle": "2024-02-21T02:56:06.911052Z",
     "shell.execute_reply": "2024-02-21T02:56:06.910591Z",
     "shell.execute_reply.started": "2024-02-21T02:55:56.417025Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "正克隆到 'DiT'...\n",
      "remote: Enumerating objects: 102, done.\u001b[K\n",
      "remote: Counting objects: 100% (78/78), done.\u001b[K\n",
      "remote: Compressing objects: 100% (43/43), done.\u001b[K\n",
      "remote: Total 102 (delta 55), reused 35 (delta 35), pack-reused 24\u001b[K\n",
      "接收对象中: 100% (102/102), 6.37 MiB | 4.06 MiB/s, 完成.\n",
      "处理 delta 中: 100% (56/56), 完成.\n",
      "Looking in indexes: https://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: diffusers in /opt/conda/lib/python3.10/site-packages (0.26.3)\n",
      "Requirement already satisfied: timm in /opt/conda/lib/python3.10/site-packages (0.9.16)\n",
      "Requirement already satisfied: importlib-metadata in /opt/conda/lib/python3.10/site-packages (from diffusers) (7.0.1)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from diffusers) (3.13.1)\n",
      "Requirement already satisfied: huggingface-hub>=0.20.2 in /opt/conda/lib/python3.10/site-packages (from diffusers) (0.20.3)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from diffusers) (1.26.3)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.10/site-packages (from diffusers) (2023.12.25)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from diffusers) (2.31.0)\n",
      "Requirement already satisfied: safetensors>=0.3.1 in /opt/conda/lib/python3.10/site-packages (from diffusers) (0.4.1)\n",
      "Requirement already satisfied: Pillow in /opt/conda/lib/python3.10/site-packages (from diffusers) (10.2.0)\n",
      "Requirement already satisfied: torch in /opt/conda/lib/python3.10/site-packages (from timm) (2.1.2+cu121)\n",
      "Requirement already satisfied: torchvision in /opt/conda/lib/python3.10/site-packages (from timm) (0.16.2+cu121)\n",
      "Requirement already satisfied: pyyaml in /opt/conda/lib/python3.10/site-packages (from timm) (6.0.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.20.2->diffusers) (2023.10.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.20.2->diffusers) (4.65.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.20.2->diffusers) (4.9.0)\n",
      "Requirement already satisfied: packaging>=20.9 in /opt/conda/lib/python3.10/site-packages (from huggingface-hub>=0.20.2->diffusers) (23.1)\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.10/site-packages (from importlib-metadata->diffusers) (3.17.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->diffusers) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->diffusers) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->diffusers) (1.26.16)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->diffusers) (2023.11.17)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch->timm) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch->timm) (2.8.4)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch->timm) (3.1.2)\n",
      "Requirement already satisfied: triton==2.1.0 in /opt/conda/lib/python3.10/site-packages (from torch->timm) (2.1.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch->timm) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch->timm) (1.3.0)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2024-02-21 10:56:06,878 - modelscope - INFO - PyTorch version 2.1.2+cu121 Found.\n",
      "2024-02-21 10:56:06,880 - modelscope - INFO - TensorFlow version 2.14.0 Found.\n",
      "2024-02-21 10:56:06,881 - modelscope - INFO - Loading ast index from /mnt/workspace/.cache/modelscope/ast_indexer\n",
      "2024-02-21 10:56:06,907 - modelscope - INFO - Loading done! Current index file version is 1.12.0, with md5 509123dba36c5e70a95f6780df348471 and a total number of 964 components indexed\n"
     ]
    }
   ],
   "source": [
    "# DiT imports:\n",
    "import torch\n",
    "from torchvision.utils import save_image\n",
    "from diffusion import create_diffusion\n",
    "from diffusers.models import AutoencoderKL\n",
    "from download import find_model\n",
    "from models import DiT_XL_2\n",
    "from PIL import Image\n",
    "from IPython.display import display\n",
    "from modelscope import snapshot_download\n",
    "torch.set_grad_enabled(False)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "if device == \"cpu\":\n",
    "    print(\"GPU not found. Using CPU instead.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AXpziRkoOvV9"
   },
   "source": [
    "# Download DiT-XL/2 Models\n",
    "\n",
    "You can choose between a 512x512 model and a 256x256 model. You can swap-out the LDM VAE, too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-02-21T02:58:20.338677Z",
     "iopub.status.busy": "2024-02-21T02:58:20.338356Z",
     "iopub.status.idle": "2024-02-21T02:58:31.246188Z",
     "shell.execute_reply": "2024-02-21T02:58:31.245600Z",
     "shell.execute_reply.started": "2024-02-21T02:58:20.338656Z"
    },
    "id": "EWG-WNimO59K",
    "tags": []
   },
   "outputs": [],
   "source": [
    "image_size = 256 #@param [256, 512]\n",
    "vae_model = snapshot_download(\"AI-ModelScope/sd-vae-ft-ema\") #@param [\"stabilityai/sd-vae-ft-mse\", \"stabilityai/sd-vae-ft-ema\"]\n",
    "latent_size = int(image_size) // 8\n",
    "# Load model:\n",
    "model = DiT_XL_2(input_size=latent_size).to(device)\n",
    "DiT_model = snapshot_download(f\"AI-ModelScope/DiT-XL-2-{image_size}x{image_size}\")\n",
    "state_dict = find_model(f\"{DiT_model}/DiT-XL-2-{image_size}x{image_size}.pt\")\n",
    "model.load_state_dict(state_dict)\n",
    "model.eval() # important!\n",
    "vae = AutoencoderKL.from_pretrained(vae_model).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5JTNyzNZKb9E"
   },
   "source": [
    "# 2. Sample from Pre-trained DiT Models\n",
    "\n",
    "You can customize several sampling options. For the full list of ImageNet classes, [check out this](https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-21T02:58:36.546161Z",
     "iopub.status.busy": "2024-02-21T02:58:36.545823Z",
     "iopub.status.idle": "2024-02-21T03:00:26.517853Z",
     "shell.execute_reply": "2024-02-21T03:00:26.517365Z",
     "shell.execute_reply.started": "2024-02-21T02:58:36.546137Z"
    },
    "id": "-Hw7B5h4Kk4p",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 250/250 [01:49<00:00,  2.29it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 130796 characters]",
      "image/png": "[encoded data, 1312816 characters]",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1034x518>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Set user inputs:\n",
    "seed = 0 #@param {type:\"number\"}\n",
    "torch.manual_seed(seed)\n",
    "num_sampling_steps = 250 #@param {type:\"slider\", min:0, max:1000, step:1}\n",
    "cfg_scale = 4 #@param {type:\"slider\", min:1, max:10, step:0.1}\n",
    "class_labels = 207, 360, 387, 974, 88, 979, 417, 279 #@param {type:\"raw\"}\n",
    "samples_per_row = 4 #@param {type:\"number\"}\n",
    "\n",
    "# Create diffusion object:\n",
    "diffusion = create_diffusion(str(num_sampling_steps))\n",
    "\n",
    "# Create sampling noise:\n",
    "n = len(class_labels)\n",
    "z = torch.randn(n, 4, latent_size, latent_size, device=device)\n",
    "y = torch.tensor(class_labels, device=device)\n",
    "\n",
    "# Setup classifier-free guidance:\n",
    "z = torch.cat([z, z], 0)\n",
    "y_null = torch.tensor([1000] * n, device=device)\n",
    "y = torch.cat([y, y_null], 0)\n",
    "model_kwargs = dict(y=y, cfg_scale=cfg_scale)\n",
    "\n",
    "# Sample images:\n",
    "samples = diffusion.p_sample_loop(\n",
    "    model.forward_with_cfg, z.shape, z, clip_denoised=False, \n",
    "    model_kwargs=model_kwargs, progress=True, device=device\n",
    ")\n",
    "samples, _ = samples.chunk(2, dim=0)  # Remove null class samples\n",
    "samples = vae.decode(samples / 0.18215).sample\n",
    "\n",
    "# Save and display images:\n",
    "save_image(samples, \"sample.png\", nrow=int(samples_per_row), \n",
    "           normalize=True, value_range=(-1, 1))\n",
    "samples = Image.open(\"sample.png\")\n",
    "display(samples)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
