TreeDiffusion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6102e316",
   "metadata": {},
   "source": [
    "# Making a Video from Prompts with Stable Diffusion\n",
    "\n",
    "_by John Robinson_\n",
    "\n",
    "I created this notebook while taking Jeremy Howards's fantastic course, [\"From Deep Learning Foundations to Stable Diffusion\"](https://www.fast.ai/posts/part2-2022.html).\n",
    "\n",
    "This notebook demonstrates using [StableDiffusion](https://stability.ai/blog/stable-diffusion-public-release) to generate a movie from nothing more than a seed image and a sequence of text prompts.\n",
    "\n",
    "[@jeremyphoward](https://twitter.com/jeremyphoward) gives a great [explanation of how it works here.](https://twitter.com/jeremyphoward/status/1583667503091548161)\n",
    "\n",
    "![Snowy Tree](/~https://github.com/johnrobinsn/diffusion_experiments/blob/main/images/tree_snow.png?raw=true)\n",
    "\n",
    "Follow me on twitter [johnrobinsn](https://twitter.com/johnrobinsn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ee1674c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I used conda with python 3.9\n",
    "\n",
    "def install_dependencies():\n",
    "    !pip install -qq numpy\n",
    "    !pip install -qq matplotlib\n",
    "    !pip install -qq fastai\n",
    "    !pip install -qq --upgrade transformers diffusers ftfy\n",
    "    !conda install -y -qq ffmpeg\n",
    "\n",
    "# Uncomment this line if you'd like to install the dependencies. \n",
    "#install_dependencies()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcffbce4",
   "metadata": {},
   "source": [
    "## Imports and Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d7da38",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy\n",
    "\n",
    "# For video display:\n",
    "from IPython.display import HTML\n",
    "from base64 import b64encode\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from PIL import Image\n",
    "import torch, logging\n",
    "from torch import autocast\n",
    "from torchvision import transforms as tfms\n",
    "\n",
    "from fastcore.all import concat\n",
    "from pathlib import Path\n",
    "\n",
    "from huggingface_hub import notebook_login\n",
    "from transformers import CLIPTextModel, CLIPTokenizer\n",
    "from transformers import logging\n",
    "from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler\n",
    "\n",
    "# Set device\n",
    "torch_device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0314552",
   "metadata": {},
   "source": [
    "## Authenticate with Hugging Face"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1cc43d2",
   "metadata": {},
   "source": [
    "To run Stable Diffusion on your computer you have to accept the model license. It's an open CreativeML OpenRail-M license that claims no rights on the outputs you generate and prohibits you from deliberately producing illegal or harmful content. The [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) provides more details. If you do accept the license, you need to be a registered user in 🤗 Hugging Face Hub and use an access token for the code to work. You have two options to provide your access token:\n",
    "\n",
    "* Use the `huggingface-cli login` command-line tool in your terminal and paste your token when prompted. It will be saved in a file in your computer.\n",
    "* Or use `notebook_login()` in a notebook, which does the same thing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d01f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1)\n",
    "if not (Path.home()/'.huggingface'/'token').exists(): notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "229427fa",
   "metadata": {},
   "source": [
    "## Load Pretrained Hugging Face Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b6cce99",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load the autoencoder model which will be used to decode the latents into image space. \n",
    "vae = AutoencoderKL.from_pretrained(\"CompVis/stable-diffusion-v1-4\", subfolder=\"vae\")\n",
    "\n",
    "# Load the tokenizer and text encoder to tokenize and encode the text. \n",
    "tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "text_encoder = CLIPTextModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "\n",
    "# The UNet model for generating the latents.\n",
    "unet = UNet2DConditionModel.from_pretrained(\"CompVis/stable-diffusion-v1-4\", subfolder=\"unet\")\n",
    "\n",
    "# The noise scheduler\n",
    "# hyper parameters match those used during training the model\n",
    "scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", num_train_timesteps=1000)\n",
    "\n",
    "# To the GPU we go!\n",
    "vae = vae.to(torch_device)\n",
    "text_encoder = text_encoder.to(torch_device)\n",
    "unet = unet.to(torch_device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07e18ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_magic = 0.18215 # vae model trained with a scale term to get closer to unit variance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b38bf983",
   "metadata": {},
   "source": [
    "## Functions to Convert between Latents and Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc52d09",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def image2latent(im):\n",
    "    im = tfms.ToTensor()(im).unsqueeze(0)\n",
    "    with torch.no_grad():\n",
    "        latent = vae.encode(im.to(torch_device)*2-1);\n",
    "    latent = latent.latent_dist.sample() * vae_magic      \n",
    "    return latent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f8769a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def latents2images(latents):\n",
    "    latents = latents/vae_magic\n",
    "    with torch.no_grad():\n",
    "        imgs = vae.decode(latents).sample\n",
    "    imgs = (imgs / 2 + 0.5).clamp(0,1)\n",
    "    imgs = imgs.detach().cpu().permute(0,2,3,1).numpy()\n",
    "    imgs = (imgs * 255).round().astype(\"uint8\")\n",
    "    imgs = [Image.fromarray(i) for i in imgs]\n",
    "    return imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ea3dc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clamp(n, smallest, largest): return max(smallest, min(n, largest))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da574df0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_image_from_embedding(text_embeddings, im_latents, seed=32):\n",
    "    height = 512                        # default height of Stable Diffusion\n",
    "    width = 512                         # default width of Stable Diffusion\n",
    "    num_inference_steps = 50  #30           # Number of denoising steps\n",
    "    guidance_scale = 7.5                # Scale for classifier-free guidance\n",
    "    generator = torch.manual_seed(seed)   # Seed generator to create the inital latent noise\n",
    "\n",
    "    max_length = tokenizer.model_max_length\n",
    "    uncond_input = tokenizer(\n",
    "      [\"\"], padding=\"max_length\", max_length=max_length, return_tensors=\"pt\"\n",
    "    )\n",
    "    with torch.no_grad():\n",
    "        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] \n",
    "    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])\n",
    "\n",
    "    # Prep Scheduler\n",
    "    scheduler.set_timesteps(num_inference_steps)\n",
    "\n",
    "    # Prep latents\n",
    "    \n",
    "    if im_latents != None:\n",
    "        # img2img\n",
    "        start_step = 10\n",
    "        noise = torch.randn_like(im_latents)\n",
    "        latents = scheduler.add_noise(im_latents,noise,timesteps=torch.tensor([scheduler.timesteps[start_step]]))\n",
    "        latents = latents.to(torch_device).float()\n",
    "    else:\n",
    "        # just text prompts\n",
    "        start_step = -1 # disable branching below\n",
    "        latents = torch.randn((1,unet.in_channels,height//8,width//8))#,generator=generator)\n",
    "        latents = latents.to(torch_device)\n",
    "        latents = latents * scheduler.init_noise_sigma # scale to initial amount of noise for t0\n",
    "\n",
    "    # Loop\n",
    "    for i, t in tqdm(enumerate(scheduler.timesteps),total=50):\n",
    "        if i > start_step:\n",
    "            # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.\n",
    "            latent_model_input = torch.cat([latents] * 2)\n",
    "            latent_model_input = scheduler.scale_model_input(latent_model_input, t)\n",
    "\n",
    "            # predict the noise residual\n",
    "            with torch.no_grad():\n",
    "                noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)[\"sample\"]\n",
    "\n",
    "            # perform guidance\n",
    "            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
    "            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
    "\n",
    "            # compute the previous noisy sample x_t -> x_t-1\n",
    "            latents = scheduler.step(noise_pred, t, latents).prev_sample\n",
    "\n",
    "    return latents2images(latents)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f30368c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embedding_for_prompt(prompt):\n",
    "    tokens = tokenizer([prompt], padding=\"max_length\", max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\")\n",
    "    with torch.no_grad():\n",
    "        embeddings = text_encoder(tokens.input_ids.to(torch_device))[0]\n",
    "    return embeddings\n",
    "\n",
    "def generate_image_from_embeddings(embeddings,im_latents,pos = 0,seed=32):\n",
    "    # integer part of pos is used for prompt index;\n",
    "    # fractional part of pos is used to \"lerp\" between the embeddings\n",
    "    l = len(embeddings)\n",
    "    if l > 1:\n",
    "        index = clamp(int(pos),0,len(embeddings)-2)\n",
    "        mix = clamp(pos-index,0,1)\n",
    "        mixed_embeddings = (embeddings[index]*(1-mix)+embeddings[index+1]*mix)\n",
    "        return generate_image_from_embedding(mixed_embeddings,im_latents,seed=seed)\n",
    "    elif l == 1:\n",
    "        return generate_image_from_embedding(embeddings[0],im_latents,seed=seed)\n",
    "    else:\n",
    "        raise Exception(\"Must provide at least one embedding\")\n",
    "        \n",
    "def generate_movie_from_prompts(prompts,im_latents,outdir,fps=12,seconds_per_prompt=2,seed=32):\n",
    "    if not os.path.exists(outdir): os.mkdir(outdir)\n",
    "    num_prompts = len(prompts)\n",
    "    num_frames = (num_prompts-1) * seconds_per_prompt * fps\n",
    "    embeddings = [get_embedding_for_prompt(p) for p in prompts]\n",
    "    for f in tqdm(range(0,num_frames)):\n",
    "        im = generate_image_from_embeddings(embeddings,im_latents,(f/num_frames)*num_prompts,seed=seed)\n",
    "        im.save(f'{outdir}/{f:04}.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "327c0d16",
   "metadata": {},
   "source": [
    "## Create Video from Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a18afad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_movie(dir,movie_name,fps=12):\n",
    "    !ffmpeg -v 1 -y -f image2 -framerate {fps} -i {dir}/%04d.jpg -c:v libx264 -preset slow -qp 18 -pix_fmt yuv420p {movie_name}\n",
    "\n",
    "def embed_movie(movie_name):\n",
    "    mp4 = open(movie_name,'rb').read()\n",
    "    data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "    return\"\"\"\n",
    "    <video width=600 controls>\n",
    "          <source src=\"%s\" type=\"video/mp4\">\n",
    "    </video>\n",
    "    \"\"\" % data_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd7cc8f3",
   "metadata": {},
   "source": [
    "## A Tree in Four Seasons... \n",
    "\n",
    "Demonstrate how to create a little video of a tree using Stable Diffusion and the following text prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68d547da",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_prompts = [\n",
    "    \"An oak tree with bare branches in the winter snowing blizzard bleak\",\n",
    "    \"A barren oak tree with no leaves and grass on the ground\",\n",
    "    \"An oak tree in the spring with bright green leaves\",\n",
    "    \"An oak tree in the summer with dark green leaves with a squirrel on the trunk\",\n",
    "    \"An oak tree in the fall with colorful leaves on the ground\",\n",
    "    \"An barren oak tree with no leaves in the fall leaves on the ground long shadows\",\n",
    "    \"An oak tree with bare branches in the winter snowing blizzard bleak\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7edc48e1",
   "metadata": {},
   "source": [
    "### From Prompts Alone\n",
    "First let's create a little movie about our tree from text prompts alone.  \n",
    "\n",
    "In order to get a more stable video, it's best to create a reference image and use *stable diffusion* in an image-to-image mode.  Here we'll generate an image using Stable Diffusion from prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c78d0b23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find a prompt that gives us a nice image for our tree\n",
    "tree_embedding = get_embedding_for_prompt('A magestic oak tree with bright green leaves on top of a hill')\n",
    "\n",
    "tree_image = generate_image_from_embeddings([tree_embedding],None,0,seed=17390125398225616219)\n",
    "tree_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d22b70a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert our generated image into SD latent-space\n",
    "tree_encoded = image2latent(tree_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "578b50df",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Let's make a quick movie to help debug the prompts\n",
    "img_dir = 'tree_frames_prompt'\n",
    "movie = 'tree_prompt.mp4'\n",
    "\n",
    "generate_movie_from_prompts(tree_prompts,tree_encoded,img_dir,fps=2,seconds_per_prompt=1,seed=17390125398225616219)\n",
    "create_movie(img_dir,movie,fps=1)\n",
    "HTML(embed_movie(movie))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce0d1647",
   "metadata": {},
   "source": [
    "### From an Existing Reference Image\n",
    "Now let's show how to use a already available image to guide the creation of our video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "689b8d44",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Load the image\n",
    "img = Image.open('./images/oak_tree.jpg').resize((512,512));img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "add3ae57",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoded = image2latent(img); encoded.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2460c41",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Let's make a quick movie to help debug the prompts\n",
    "\n",
    "img_dir = 'tree_frames_quick'\n",
    "movie = 'tree_quick.mp4'\n",
    "\n",
    "generate_movie_from_prompts(tree_prompts,encoded,img_dir,fps=2,seconds_per_prompt=1,seed=17390125398225616219)\n",
    "create_movie(img_dir,movie,fps=1)\n",
    "HTML(embed_movie(movie))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82eca17a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# WARNING This will take a while.\n",
    "# Make a \"full length\" movie; two seconds per prompt; 12 fps\n",
    "\n",
    "img_dir = 'tree_frames'\n",
    "movie = 'tree.mp4'\n",
    "\n",
    "generate_movie_from_prompts(tree_prompts,encoded,img_dir,seed=17390125398225616219)\n",
    "create_movie(img_dir,movie)\n",
    "HTML(embed_movie(movie))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43008aeb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.9.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {
    "height": "428.993px",
    "width": "279.983px"
   },
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}