Running inference against DeepFloyd's IF on RunPod

inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!grep MemTotal /proc/meminfo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "notebook_login()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import json\n",
    "import os\n",
    "import random\n",
    "import torch\n",
    "\n",
    "os.environ['FORCE_MEM_EFFICIENT_ATTN'] = \"1\"\n",
    "from cog import BasePredictor, Path, Input\n",
    "from datetime import datetime\n",
    "from deepfloyd_if.modules import IFStageI, IFStageII, StableStageIII\n",
    "from deepfloyd_if.modules.t5 import T5Embedder\n",
    "from deepfloyd_if.pipelines import dream\n",
    "from typing import List\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# memory flush\n",
    "def flush():\n",
    "  gc.collect()\n",
    "  torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "flush()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup\n",
    "device = 'cuda:0'\n",
    "if_I = IFStageI('IF-I-XL-v1.0', device=device)\n",
    "if_II = IFStageII('IF-II-L-v1.0', device=device)\n",
    "if_III = StableStageIII('stable-diffusion-x4-upscaler', device=device)\n",
    "t5 = T5Embedder(device=device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = int((datetime.utcnow().timestamp() * 10 ** 6) % (2 ** 32 - 1))\n",
    "print(f'seed is {seed}')\n",
    "aspect_ratio = '2:3'\n",
    "\n",
    "prompt = (\"a portrait of a cat goddess wearing a mourning gown, she is going to a memorial service, a black and white photo by Frieke Janssens, stark, beautiful, coherent, bastet\")\n",
    "negative_prompt = 'photo frame'\n",
    "\n",
    "result = dream(\n",
    "    t5=t5, if_I=if_I, if_II=if_II, if_III=if_III,\n",
    "    disable_watermark=True,\n",
    "    negative_prompt=negative_prompt,\n",
    "    # style_prompt=style_prompt,\n",
    "    aspect_ratio=aspect_ratio,\n",
    "    prompt=prompt,\n",
    "    seed=seed,\n",
    "    if_I_kwargs={\n",
    "        \"guidance_scale\": 7.0,\n",
    "        \"sample_timestep_respacing\": \"smart100\",\n",
    "    },\n",
    "    if_II_kwargs={\n",
    "        \"guidance_scale\": 4.0,\n",
    "        \"sample_timestep_respacing\": \"smart50\",\n",
    "    },\n",
    "    if_III_kwargs={\n",
    "        \"guidance_scale\": 9.0,\n",
    "        \"noise_level\": 20,\n",
    "        \"sample_timestep_respacing\": \"75\",\n",
    "    },\n",
    ")\n",
    "\n",
    "if_III.show(result['III'], size=14)\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.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

setup.sh

#!/usr/bin/env sh
# Run this script against a running instance of runpod/pytorch:3.10-2.0.0-117 to
# get CUDA 11.8 and update torch 2 for running inference with DeepFloyd's IF model
# (via https://github.com/0x7o/IF-replicate which did the work of lining up the
# dependencies etc.)

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2004-11-8-local_11.8.0-520.61.05-1_amd64.deb
mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
dpkg -i cuda-repo-ubuntu2004-11-8-local_11.8.0-520.61.05-1_amd64.deb
cp /var/cuda-repo-ubuntu2004-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/

apt update
apt -y upgrade
apt -y install cuda less vim

python3 -m pip install --upgrade pip
python3 -m pip install --extra-index-url https://download.pytorch.org/whl/cu118 \
  git+https://github.com/0x7o/IF.git \
  git+https://github.com/openai/CLIP.git \
  cog \
  huggingface_hub==0.14.1 \
  Pillow==9.5.0 \
  protobuf==3.20.1 \
  safetensors==0.3.0 \
  scikit-image \
  sentencepiece==0.1.98 \
  tokenizers==0.13.3 \
  torch==2.0+cu118 \
  torchvision==0.15.1+cu118 \
  tqdm==4.65.0 \
  xformers==0.0.19