5min-demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "845dbe61-4889-45e0-ad8a-76340bba7a21",
      "metadata": {
        "id": "845dbe61-4889-45e0-ad8a-76340bba7a21"
      },
      "source": [
        "# Exploiting the Signal-Leak Bias in Diffusion Models\n",
        "[![Open This Demo In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/IVRL/signal-leak-bias/blob/main/5min-demo.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0df56820-d059-4a1b-8c52-6eab7ff291e7",
      "metadata": {
        "id": "0df56820-d059-4a1b-8c52-6eab7ff291e7"
      },
      "source": [
        "**Project page and Research paper:** https://ivrl.github.io/signal-leak-bias\n",
        "\n",
        "**Link to this demo:** https://ivrl.github.io/signal-leak-bias/demo\n",
        "\n",
        "[![arXiv](https://img.shields.io/badge/arXiv-2309.15842-red)](https://arxiv.org/abs/2309.15842)\n",
        "[![Project Page](https://img.shields.io/badge/Project%20Page-IVRL-blue)](https://ivrl.github.io/signal-leak-bias/)\n",
        "[![Proceedings](https://img.shields.io/badge/WACV%20Proceedings-CVF-blue)](https://openaccess.thecvf.com/content/WACV2024/html/Everaert_Exploiting_the_Signal-Leak_Bias_in_Diffusion_Models_WACV_2024_paper.html)\n",
        "[![Code](https://img.shields.io/badge/Code-Github-black)](https://github.com/IVRL/signal-leak-bias)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9d0edfae-d449-48f6-b915-f9a3c57143d4",
      "metadata": {
        "id": "9d0edfae-d449-48f6-b915-f9a3c57143d4",
        "jp-MarkdownHeadingCollapsed": true
      },
      "source": [
        "## Who am I?\n",
        "\n",
        "**Martin Nicolas Everaert**\n",
        "\n",
        "PhD student at EPFL - IVRL\n",
        "\n",
        "https://martin-ev.github.io/"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6dd2fe68-d212-45ec-8c8d-5b79620f8264",
      "metadata": {
        "id": "6dd2fe68-d212-45ec-8c8d-5b79620f8264",
        "jp-MarkdownHeadingCollapsed": true
      },
      "source": [
        "## Intro\n",
        "\n",
        "In this demo, we will explore the key findings from my paper titled\n",
        "<strong>\"Exploiting the Signal-Leak Bias in Diffusion Models\"</strong> (WACV 2024).\n",
        "\n",
        "Feel free to reach out to me if you have any questions or feedback!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d436e682-0b10-4ff9-99fd-b046c98b408d",
      "metadata": {
        "id": "d436e682-0b10-4ff9-99fd-b046c98b408d"
      },
      "source": [
        "## Setup and Imports\n",
        "\n",
        "Here, we install the **🤗 Diffusers** library, and import the necessary components for the demo.\n",
        "\n",
        "This demo borrows code / is built on top of the **🤗 Diffusers** library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2ba76054-e69c-41b9-9ff2-6b2eed9729ec",
      "metadata": {
        "id": "2ba76054-e69c-41b9-9ff2-6b2eed9729ec"
      },
      "outputs": [],
      "source": [
        "!pip install diffusers==0.29.0\n",
        "!pip install transformers==4.41.2\n",
        "!pip install accelerate==0.31.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2d59da9c-d9bb-4dba-a049-27af18be7fb5",
      "metadata": {
        "id": "2d59da9c-d9bb-4dba-a049-27af18be7fb5"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "from PIL import Image\n",
        "from torchvision import transforms\n",
        "from diffusers import DiffusionPipeline\n",
        "from diffusers.utils import load_image\n",
        "from diffusers.utils.torch_utils import randn_tensor\n",
        "\n",
        "_ = torch.set_grad_enabled(False)\n",
        "gpu = torch.device('cuda')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "87d65974-1a2b-471b-8cc1-b87102d8fbe7",
      "metadata": {
        "id": "87d65974-1a2b-471b-8cc1-b87102d8fbe7"
      },
      "source": [
        "## Load Stable Diffusion 2.1\n",
        "\n",
        "We load the model from HuggingFace."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ab62fd0c-f069-4b41-94c2-5db955bd0dd3",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "6123af5cf8dc44b0bb6772c75f466aa7"
          ]
        },
        "id": "ab62fd0c-f069-4b41-94c2-5db955bd0dd3",
        "outputId": "7423b239-db5e-4903-f27a-71fe8456a907"
      },
      "outputs": [],
      "source": [
        "pipeline = DiffusionPipeline.from_pretrained(\n",
        "    \"stabilityai/stable-diffusion-2-1\",\n",
        "    torch_dtype=torch.float16\n",
        ")\n",
        "pipeline = pipeline.to(gpu)\n",
        "pipeline.scheduler.config.timestep_spacing = \"trailing\""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3a2d8488-5856-479b-ae5d-81486e379f91",
      "metadata": {
        "id": "3a2d8488-5856-479b-ae5d-81486e379f91"
      },
      "source": [
        "## Success and Failures of Stable Diffusion"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3155de78-bac8-42f6-98a0-ad9d6207509c",
      "metadata": {
        "id": "3155de78-bac8-42f6-98a0-ad9d6207509c"
      },
      "source": [
        "### It often works well"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3ba0da23-3663-4a40-85e9-d75b753b392d",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "909bb2e031264adeb59782cef67bbdbb"
          ]
        },
        "id": "3ba0da23-3663-4a40-85e9-d75b753b392d",
        "outputId": "42cb5f41-cb62-4431-97f6-a6417731f66f"
      },
      "outputs": [],
      "source": [
        "prompt = \"A professional photograph of an astronaut mowing the lawn. Moon in the background. Colorful image, hyperrealistic, fantasy, dark art.\"\n",
        "print(prompt)\n",
        "\n",
        "image = pipeline(prompt).images[0]\n",
        "\n",
        "display(image.resize((500,500)))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "737b2b3d-d2bd-4d51-8357-5dfeda4dba35",
      "metadata": {
        "id": "737b2b3d-d2bd-4d51-8357-5dfeda4dba35"
      },
      "source": [
        " ### But it fails on simple prompts such as \"solid red background\" on \"drawing on a white background\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4bb6fcdf-4172-4f59-9727-62578783081b",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "d60b2057322f4e168282c59ee2da4c79"
          ]
        },
        "id": "4bb6fcdf-4172-4f59-9727-62578783081b",
        "outputId": "854134ab-1bc2-4fc6-8b5a-cd65738e1f33"
      },
      "outputs": [],
      "source": [
        "prompt = \"A solid red background.\"\n",
        "print(prompt)\n",
        "\n",
        "image = pipeline(prompt).images[0]\n",
        "\n",
        "display(image.resize((500,500)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "db197a16-3503-4b8d-afe6-4c8bad88e75c",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "1f0d65c801454dd6a50a70a7090ad0f9"
          ]
        },
        "id": "db197a16-3503-4b8d-afe6-4c8bad88e75c",
        "outputId": "8e0d3385-ed39-45bd-e98f-57b9cb7d7a02"
      },
      "outputs": [],
      "source": [
        "prompt = \"A drawing of a squirrel on a completely white background.\"\n",
        "print(prompt)\n",
        "\n",
        "image = pipeline(prompt).images[0]\n",
        "\n",
        "display(image.resize((500,500)))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "80ac83d0-dd85-4659-ad43-4bcce0a05dd5",
      "metadata": {
        "id": "80ac83d0-dd85-4659-ad43-4bcce0a05dd5"
      },
      "source": [
        "### The diffusers library lets us control the initial noise\n",
        "\n",
        "We observe that the same initial noise (same seed) generates similar images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d17b3d43-7298-439d-9d59-958116f7d668",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "32ac09e22c6a45bd8c649ba3bdd6c014",
            "28c7b0feb51b4e2ba438c4f66854fa38",
            "3bdb546e477649179c42fea1a298953e"
          ]
        },
        "id": "d17b3d43-7298-439d-9d59-958116f7d668",
        "outputId": "59fc3593-f87d-43ba-a660-9c72f5b361e7"
      },
      "outputs": [],
      "source": [
        "noise = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(123456))\n",
        "\n",
        "prompt1 = \"A drawing of a squirrel on a completely white background.\"\n",
        "prompt2 = \"A colorful drawing of a squirrel on a completely white background.\"\n",
        "prompt3 = \"A realistic photograph of a squirrel on a completely white background.\"\n",
        "\n",
        "print(prompt1)\n",
        "image1 = pipeline(\n",
        "    prompt1,\n",
        "    latents=noise\n",
        ").images[0]\n",
        "display(image1.resize((500,500)))\n",
        "\n",
        "print(prompt2)\n",
        "image2 = pipeline(\n",
        "    prompt2,\n",
        "    latents=noise\n",
        ").images[0]\n",
        "display(image2.resize((500,500)))\n",
        "\n",
        "print(prompt3)\n",
        "image3 = pipeline(\n",
        "    prompt3,\n",
        "    latents=noise\n",
        ").images[0]\n",
        "display(image3.resize((500,500)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "435a5cab-782f-4e71-b7df-0bb61a5b4b52",
      "metadata": {
        "id": "435a5cab-782f-4e71-b7df-0bb61a5b4b52"
      },
      "source": [
        "See also: https://huggingface.co/docs/diffusers/v0.13.0/en/using-diffusers/reusing_seeds"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a1e21772-474c-4347-bda5-129a4baf6d8d",
      "metadata": {
        "id": "a1e21772-474c-4347-bda5-129a4baf6d8d"
      },
      "source": [
        "## Our solution\n",
        "\n",
        "We fix the issues by exploiting the initial noise:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9e4951ff-82fa-4de7-bb4b-38abaf3ce7e2",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "2a11cc0d9a084befa84d9ffe0b2d7daa"
          ]
        },
        "id": "9e4951ff-82fa-4de7-bb4b-38abaf3ce7e2",
        "outputId": "a78f8f1d-e0ed-4ba9-a147-b16dedf5175f"
      },
      "outputs": [],
      "source": [
        "noise = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "signal_leak_white_background = torch.tensor(\n",
        "    [[[[1.9]], [[1.3]], [[0.0]], [[-1.0]]]],\n",
        "    device=gpu,\n",
        "    dtype=torch.float16\n",
        ")\n",
        "\n",
        "initial_noise = 0.06826 * signal_leak_white_background + 0.99767 * noise\n",
        "\n",
        "prompt = \"A drawing of a squirrel on a white background.\"\n",
        "print(prompt)\n",
        "\n",
        "image = pipeline(\n",
        "    prompt,\n",
        "    latents=initial_noise\n",
        ").images[0]\n",
        "display(image.resize((500,500)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9a1e99f7-96ce-46f1-b4d8-c87308f94228",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "7bfe07fe7f1d4f0b92a85d4bb2a7e5aa"
          ]
        },
        "id": "9a1e99f7-96ce-46f1-b4d8-c87308f94228",
        "outputId": "1e6a9402-014b-4f7e-b2e2-cfd9dd6d5381"
      },
      "outputs": [],
      "source": [
        "noise = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "signal_leak_red_background = torch.tensor(\n",
        "    [[[[0.0]], [[0.0]], [[-2.5]], [[0.0]]]],\n",
        "    device=gpu,\n",
        "    dtype=torch.float16\n",
        ")\n",
        "\n",
        "initial_noise = 0.06826 * signal_leak_red_background + 0.99767 * noise\n",
        "\n",
        "prompt = \"A solid red background.\"\n",
        "print(prompt)\n",
        "\n",
        "image = pipeline(\n",
        "    prompt,\n",
        "    latents=initial_noise\n",
        ").images[0]\n",
        "display(image.resize((500,500)))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "02e48a9b-d3b6-4595-ba80-4c844d6e772b",
      "metadata": {
        "id": "02e48a9b-d3b6-4595-ba80-4c844d6e772b"
      },
      "source": [
        "## How are images generated?\n",
        "\n",
        "Stable Diffusion generates images by iteratively denoising an initial noise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b35d72fa-be37-42dd-937f-9ad8a5081b60",
      "metadata": {
        "id": "b35d72fa-be37-42dd-937f-9ad8a5081b60",
        "outputId": "f92e13c0-1692-400e-ecba-637ef15b8754"
      },
      "outputs": [],
      "source": [
        "noise = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "def decode_and_show(noisy_image):\n",
        "    image = pipeline.vae.decode(noisy_image/pipeline.vae.config.scaling_factor, return_dict=False)[0]\n",
        "    image = pipeline.image_processor.postprocess(image, output_type=\"pil\")[0]\n",
        "    display(image.resize((500,500)))\n",
        "\n",
        "decode_and_show(noise)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "414191f5-6af3-4de8-b7fe-7192d9e423bd",
      "metadata": {
        "id": "414191f5-6af3-4de8-b7fe-7192d9e423bd"
      },
      "source": [
        "Stable Diffusion uses timesteps to represent the current level of noise:\n",
        "- $t=1000$: (almost) indistinguishable from noise\n",
        "- $t=1$: (almost) clean original image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e432ccca-57e4-4297-abec-7f1a2986929e",
      "metadata": {
        "id": "e432ccca-57e4-4297-abec-7f1a2986929e",
        "outputId": "cf32a6c8-e778-46c5-f95b-731937a253b7"
      },
      "outputs": [],
      "source": [
        "num_inference_steps = 10\n",
        "pipeline.scheduler.config.timestep_spacing = \"trailing\"\n",
        "pipeline.scheduler.set_timesteps(num_inference_steps, device=gpu)\n",
        "timesteps = pipeline.scheduler.timesteps\n",
        "print(timesteps)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a894d8ee-070c-4b7a-b172-e6c0bb13b738",
      "metadata": {
        "id": "a894d8ee-070c-4b7a-b172-e6c0bb13b738"
      },
      "source": [
        "The textual prompt is preprocessed (encoded):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ba76b38b-bf3b-41fb-8c08-9791f16a9104",
      "metadata": {
        "id": "ba76b38b-bf3b-41fb-8c08-9791f16a9104",
        "outputId": "3cb3d1a7-af46-4d24-f615-c54986cae1ca"
      },
      "outputs": [],
      "source": [
        "prompt = \"A drawing of a squirrel on completely white background.\"\n",
        "prompt_embeds, negative_prompt_embeds = pipeline.encode_prompt(\n",
        "    prompt=prompt,\n",
        "    device=gpu,\n",
        "    num_images_per_prompt=1,\n",
        "    do_classifier_free_guidance=True\n",
        ")\n",
        "print(prompt_embeds)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dd7c0320-0ffa-4d0b-9d44-efb0af978cb3",
      "metadata": {
        "id": "dd7c0320-0ffa-4d0b-9d44-efb0af978cb3"
      },
      "source": [
        "At each timestep, the diffusion model predicts the clean image and denoise one step towards the predicted clean image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b88c194d-f6a9-4d7b-b287-6716c56038b6",
      "metadata": {
        "id": "b88c194d-f6a9-4d7b-b287-6716c56038b6"
      },
      "outputs": [],
      "source": [
        "def predict_original_image(noisy_image, t, guidance_scale=7.5):\n",
        "\n",
        "    model_input = torch.cat([noisy_image] * 2)\n",
        "    model_input = pipeline.scheduler.scale_model_input(model_input, t)\n",
        "\n",
        "    model_output = pipeline.unet(\n",
        "        model_input,\n",
        "        t,\n",
        "        encoder_hidden_states=torch.cat([negative_prompt_embeds, prompt_embeds]),\n",
        "        return_dict=False,\n",
        "    )[0]\n",
        "\n",
        "    alpha_prod_t = pipeline.scheduler.alphas_cumprod[t]\n",
        "    beta_prod_t = 1 - alpha_prod_t\n",
        "\n",
        "    pred_original_sample = (alpha_prod_t**0.5) * noisy_image - (beta_prod_t**0.5) * model_output #velocity-prediction\n",
        "    pred_uncond, pred_text = pred_original_sample.chunk(2)\n",
        "    pred_original_sample = pred_uncond + guidance_scale * (pred_text - pred_uncond)\n",
        "\n",
        "    return pred_original_sample\n",
        "\n",
        "def denoise_one_step(noisy_image, predicted_clean_image, t, next_t):\n",
        "\n",
        "    alpha_prod_t = pipeline.scheduler.alphas_cumprod[t]\n",
        "    beta_prod_t = 1 - alpha_prod_t\n",
        "    alpha_prod_t_next = pipeline.scheduler.alphas_cumprod[next_t]\n",
        "    beta_prod_t_next = 1 - alpha_prod_t_next\n",
        "\n",
        "    factor_noisy = (beta_prod_t_next/beta_prod_t)**0.5\n",
        "    factor_prediction = alpha_prod_t_next**0.5-(beta_prod_t_next/beta_prod_t*alpha_prod_t)**0.5\n",
        "\n",
        "    denoised_one_step = factor_noisy * noisy_image + factor_prediction * predicted_clean_image\n",
        "\n",
        "    return denoised_one_step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "76b8336e-0499-4715-8ab7-9d25cd0bd903",
      "metadata": {
        "id": "76b8336e-0499-4715-8ab7-9d25cd0bd903",
        "outputId": "c990d5f7-a8bf-40b8-a1be-068926b77d30"
      },
      "outputs": [],
      "source": [
        "noisy_image = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "for i in range(len(timesteps)):\n",
        "    t = timesteps[i]\n",
        "    print(t)\n",
        "\n",
        "    decode_and_show(noisy_image)\n",
        "    predicted_original_image = predict_original_image(noisy_image, t)\n",
        "    decode_and_show(predicted_original_image)\n",
        "\n",
        "    next_t = timesteps[i+1] if i+1<len(timesteps) else 0\n",
        "    noisy_image = denoise_one_step(noisy_image, predicted_original_image, t, next_t)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ad183877-1b4d-48e3-b5a5-e79c255c8981",
      "metadata": {
        "id": "ad183877-1b4d-48e3-b5a5-e79c255c8981"
      },
      "source": [
        "## Which noise levels were used to train the model?"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "02879842-4dd4-4f1a-b386-6c347bb5bad1",
      "metadata": {
        "id": "02879842-4dd4-4f1a-b386-6c347bb5bad1"
      },
      "source": [
        "At timestep $t \\in \\{1, 2, 3, ..., 1000\\}$, a noisy version $x_t$ of a original image $x_0$ is generated by the formula:\n",
        "\n",
        "$$x_t = \\sqrt{\\bar{\\alpha}_t} \\, x_0 + \\sqrt{1-\\bar{\\alpha}_t} \\, \\epsilon$$\n",
        "\n",
        "where $\\epsilon \\sim \\mathcal{N}(0, 1)$.\n",
        "\n",
        "The values of $\\sqrt{\\bar{\\alpha}_t}$ and $\\sqrt{1-\\bar{\\alpha}_t}$ are as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "26e5cd55-3715-41ca-b068-fcb9e1d7a361",
      "metadata": {
        "id": "26e5cd55-3715-41ca-b068-fcb9e1d7a361",
        "outputId": "40f22b59-0232-4f5a-a688-3697ffc601b4"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(10, 5))\n",
        "plt.plot( (pipeline.scheduler.alphas_cumprod) **0.5, label=r'$\\sqrt{\\bar{\\alpha}_t}$ = amount of clean image in the noisy image');\n",
        "plt.plot( (1-pipeline.scheduler.alphas_cumprod) **0.5, label=r'$\\sqrt{1-\\bar{\\alpha}_t}$ = amount of noise in the noisy image');\n",
        "plt.legend();"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9ea7cd8c-c5cb-47a8-bb3c-f7f7c07ada68",
      "metadata": {
        "id": "9ea7cd8c-c5cb-47a8-bb3c-f7f7c07ada68"
      },
      "source": [
        "Note that at the highest timestep ($t=1000$), the noisy image $x_t$ still contains $7\\%$ of the original image $x_0$:\n",
        "\n",
        "$$x_{1000} = \\sqrt{\\bar{\\alpha}_{1000}} \\, x_0 + \\sqrt{1-\\bar{\\alpha}_{1000}} \\, \\epsilon$$\n",
        "$$x_{1000} = 0.0683 \\, x_0 + 0.9977 \\, \\epsilon$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6956c23e-4d29-4b91-be92-0c323872bc01",
      "metadata": {
        "id": "6956c23e-4d29-4b91-be92-0c323872bc01",
        "outputId": "0697ee12-bc06-4492-cff6-83e18866fa2c"
      },
      "outputs": [],
      "source": [
        "alpha_prod_T = pipeline.scheduler.alphas_cumprod[999].item()\n",
        "beta_prod_T = 1 - alpha_prod_T\n",
        "\n",
        "print( alpha_prod_T**0.5 )\n",
        "print( beta_prod_T**0.5 )"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bf49617a-a44b-4c84-9c9e-f1b83418a158",
      "metadata": {
        "id": "bf49617a-a44b-4c84-9c9e-f1b83418a158"
      },
      "source": [
        "**Instead of starting denoising from complete noise, we should start from a mix of noise and $7\\%$ of the image we want.**"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9b135b1b-9e03-4027-a7ea-94bcb3a63679",
      "metadata": {
        "id": "9b135b1b-9e03-4027-a7ea-94bcb3a63679"
      },
      "source": [
        "## Details: Noise doesn't affect all frequencies equally"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "53ae96d0-7ce8-4a5a-912d-f1dedfaf40ee",
      "metadata": {
        "id": "53ae96d0-7ce8-4a5a-912d-f1dedfaf40ee",
        "outputId": "4fdf95e8-0a5a-4366-945a-f56c25a3d36d"
      },
      "outputs": [],
      "source": [
        "image = load_image(\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png\")\n",
        "image = image.crop((150, 0, 150+704, 704))\n",
        "image = image.resize((128, 128))\n",
        "display(image)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6cfb6a41-6c33-4e39-8f3c-3f9b649a782e",
      "metadata": {
        "id": "6cfb6a41-6c33-4e39-8f3c-3f9b649a782e",
        "outputId": "325fcf87-3458-472e-861c-39f4fbeac9d4"
      },
      "outputs": [],
      "source": [
        "# Code heavily borrowed/inspired from https://www.youtube.com/watch?v=3C6yEYXDijM \"Diffusion with Offset Noise\" Nicholas Guttenberg 2023\n",
        "np.random.seed(42)\n",
        "\n",
        "image_orig = np.array(image).astype(np.float32)\n",
        "image_orig = image_orig/255\n",
        "\n",
        "image_orig -= 0.5\n",
        "image_orig /= [0.229, 0.224, 0.225] #Scale to unit variance\n",
        "\n",
        "kx,ky = np.meshgrid(np.arange(128), np.arange(128))\n",
        "\n",
        "kx[kx>64] -= 128\n",
        "ky[ky>64] -= 128\n",
        "\n",
        "k2 = np.sqrt(kx*kx+ky*ky)[:,:,np.newaxis]\n",
        "mask1 = k2 < 1\n",
        "mask2 = (k2 >= 1)*(k2 < 4)\n",
        "mask3 = (k2 >= 4)*(k2 < 8)\n",
        "mask4 = (k2 >= 8)*(k2 < 16)\n",
        "mask5 = (k2 >= 16)\n",
        "\n",
        "\n",
        "for t in [0, 249, 499, 749, 999, -1]:\n",
        "\n",
        "    alpha_prod_t = pipeline.scheduler.alphas_cumprod[t].item()\n",
        "    beta_prod_t = 1 - alpha_prod_t\n",
        "\n",
        "    if t==-1:\n",
        "        noise = np.random.randn(128,128,3)\n",
        "        image = noise\n",
        "    else:\n",
        "        noise = np.random.randn(128,128,3)\n",
        "        image = alpha_prod_t**0.5 * image_orig + beta_prod_t**0.5 * noise\n",
        "\n",
        "\n",
        "    fft_im = np.fft.fft2(image, axes=[0,1])\n",
        "\n",
        "    all_compnts = image / (alpha_prod_t**0.5)\n",
        "    lp1 = np.real(np.fft.ifft2(fft_im*mask1, axes=[0,1])) / (alpha_prod_t**0.5)\n",
        "    lp2 = np.real(np.fft.ifft2(fft_im*mask2, axes=[0,1])) / (alpha_prod_t**0.5)\n",
        "    lp3 = np.real(np.fft.ifft2(fft_im*mask3, axes=[0,1])) / (alpha_prod_t**0.5)\n",
        "    lp4 = np.real(np.fft.ifft2(fft_im*mask4, axes=[0,1])) / (alpha_prod_t**0.5)\n",
        "    lp5 = np.real(np.fft.ifft2(fft_im*mask5, axes=[0,1])) / (alpha_prod_t**0.5)\n",
        "\n",
        "    images = []\n",
        "\n",
        "    for i in [all_compnts, lp1, lp2, lp3, lp4, lp5]:\n",
        "\n",
        "\n",
        "        total_image = i.copy()\n",
        "\n",
        "        total_image *= [0.229, 0.224, 0.225]\n",
        "        total_image += 0.5\n",
        "\n",
        "        total_image = total_image*255\n",
        "        total_image = np.clip(total_image, 0, 255)\n",
        "\n",
        "\n",
        "        im = Image.fromarray(total_image.astype(np.uint8))\n",
        "        images.append(im)\n",
        "\n",
        "\n",
        "    fig, axs = plt.subplots(1, 6, figsize=(15, 3))\n",
        "    titles = ['Noisy Image', 'DC', 'Low Freq', 'Mid Freq', 'High Freq', 'Very High Freq']\n",
        "\n",
        "    for ax, im, title in zip(axs, images, titles):\n",
        "        ax.imshow(im)\n",
        "        ax.set_title(f'{title}')\n",
        "        ax.axis('off')\n",
        "\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dab55478-14f0-4371-b8fb-6b65cd4e3356",
      "metadata": {
        "id": "dab55478-14f0-4371-b8fb-6b65cd4e3356"
      },
      "source": [
        "The diffusion model learns not to change the lowest-frequencies / DC components when denoising. If we do not include a signal leak at inference time, it generates greyish images."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d7a4aa63-e69b-4a91-a0cc-67bf43f09970",
      "metadata": {
        "id": "d7a4aa63-e69b-4a91-a0cc-67bf43f09970"
      },
      "source": [
        "## How do we sample the signal leak at inference time?"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8065c58d-948b-43e7-9da1-a3277fee6a3b",
      "metadata": {
        "id": "8065c58d-948b-43e7-9da1-a3277fee6a3b"
      },
      "source": [
        "We can model very roughly the distribution of desired images by computing the mean and std (pixel-wise) of a small set of images of the desired style/LF. Here we use 4 images with white background (https://huggingface.co/sd-concepts-library/one-line-drawing)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "036730e8-1977-4e3a-aa6d-2d12725f1ecb",
      "metadata": {
        "id": "036730e8-1977-4e3a-aa6d-2d12725f1ecb",
        "outputId": "145c38da-0320-43a6-90a4-321b50455a53"
      },
      "outputs": [],
      "source": [
        "resolution = 768\n",
        "train_transforms = transforms.Compose(\n",
        "    [\n",
        "        transforms.Resize(\n",
        "            resolution, interpolation=transforms.InterpolationMode.BILINEAR\n",
        "        ),\n",
        "        transforms.CenterCrop(resolution),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize([0.5], [0.5]),\n",
        "    ]\n",
        ")\n",
        "all_latents = []\n",
        "for n in range(4):\n",
        "    image = load_image(f\"https://huggingface.co/sd-concepts-library/one-line-drawing/resolve/main/concept_images/{n}.jpeg\").convert(\"RGB\")\n",
        "    pixel_values = train_transforms(image)\n",
        "    # Collate\n",
        "    pixel_values = torch.stack([pixel_values])\n",
        "    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).to(dtype=torch.float16)\n",
        "    pixel_values = pixel_values.to(pipeline.vae.device)\n",
        "    # Encode\n",
        "    latents = pipeline.vae.encode(pixel_values).latent_dist.mode()\n",
        "    latents *= pipeline.vae.config.scaling_factor\n",
        "    all_latents.append(latents)\n",
        "\n",
        "mean = torch.mean(torch.cat(all_latents), dim=0, keepdim=True)\n",
        "std = torch.std(torch.cat(all_latents), dim=0, keepdim=True)\n",
        "\n",
        "for s in range(4):\n",
        "    signal_leak = mean+std*randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(s))\n",
        "    decode_and_show(signal_leak)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "044a5939-fea3-4e65-a115-6d463f2b4b6e",
      "metadata": {
        "id": "044a5939-fea3-4e65-a115-6d463f2b4b6e"
      },
      "source": [
        "We use these mean and std to inject a signal leak in the initial noise at inference time.\n",
        "\n",
        "The generated images have a white background if we start denoising with a signal leak in the initial noise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ce387ca4-4d1c-49e8-b4b6-7b6a061e1cfe",
      "metadata": {
        "id": "ce387ca4-4d1c-49e8-b4b6-7b6a061e1cfe",
        "outputId": "e89f84e1-36f4-45eb-9c64-10fee95f0907"
      },
      "outputs": [],
      "source": [
        "initial_noise = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "signal_leak = mean+std*randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "alpha_prod_T = pipeline.scheduler.alphas_cumprod[999].item()\n",
        "beta_prod_T = 1 - alpha_prod_T\n",
        "\n",
        "noisy_image = alpha_prod_T**0.5 * signal_leak + beta_prod_T**0.5 * initial_noise\n",
        "\n",
        "for i in range(len(timesteps)):\n",
        "    t = timesteps[i]\n",
        "    print(t)\n",
        "\n",
        "    decode_and_show(noisy_image)\n",
        "    predicted_original_image = predict_original_image(noisy_image, t)\n",
        "    decode_and_show(predicted_original_image)\n",
        "\n",
        "    next_t = timesteps[i+1] if i+1<len(timesteps) else 0\n",
        "    noisy_image = denoise_one_step(noisy_image, predicted_original_image, t, next_t)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dfa21a60-9556-492d-a855-9d6af97b5d5d",
      "metadata": {
        "id": "dfa21a60-9556-492d-a855-9d6af97b5d5d"
      },
      "source": [
        "## Comparison"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "73d09b13-5259-4371-8628-6325e43abc17",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "5a4eb4be8bb547c59b28d96f3bb282e4",
            "f8e69e1053294612ad0095ea4dffbe9d"
          ]
        },
        "id": "73d09b13-5259-4371-8628-6325e43abc17",
        "outputId": "b0d59510-b9b0-493c-e462-d49afb139fec"
      },
      "outputs": [],
      "source": [
        "noise_without_signal_leak = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "signal_leak = mean+std*randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "noise_with_signal_leak = (\n",
        "    alpha_prod_T**0.5 * signal_leak\n",
        "    + beta_prod_T**0.5 * noise_without_signal_leak\n",
        ")\n",
        "\n",
        "prompt = \"A drawing of a squirrel on a white background.\"\n",
        "\n",
        "print(prompt)\n",
        "image1 = pipeline(\n",
        "    prompt,\n",
        "    latents=noise_without_signal_leak\n",
        ").images[0]\n",
        "display(image1.resize((500,500)))\n",
        "\n",
        "print(prompt)\n",
        "image2 = pipeline(\n",
        "    prompt,\n",
        "    latents=noise_with_signal_leak\n",
        ").images[0]\n",
        "display(image2.resize((500,500)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "de189e36-f2cc-46cf-9a14-6fd92f0e8824",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "c957f6f470df45a7b072b84d26371904",
            "64e2c9f3a8504a8088932b8a0820cef0"
          ]
        },
        "id": "de189e36-f2cc-46cf-9a14-6fd92f0e8824",
        "outputId": "3b0bc1dc-bed4-49eb-9d12-660a90b994f2"
      },
      "outputs": [],
      "source": [
        "noise_without_signal_leak = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "\n",
        "signal_leak = mean+std*randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(123456))\n",
        "\n",
        "noise_with_signal_leak = (\n",
        "    alpha_prod_T**0.5 * signal_leak\n",
        "    + beta_prod_T**0.5 * noise_without_signal_leak\n",
        ")\n",
        "\n",
        "prompt = \"A drawing of a Rubiks's cube on a white background.\"\n",
        "\n",
        "print(prompt)\n",
        "image1 = pipeline(\n",
        "    prompt,\n",
        "    latents=noise_without_signal_leak\n",
        ").images[0]\n",
        "display(image1.resize((500,500)))\n",
        "\n",
        "print(prompt)\n",
        "image2 = pipeline(\n",
        "    prompt,\n",
        "    latents=noise_with_signal_leak\n",
        ").images[0]\n",
        "display(image2.resize((500,500)))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8b25cac1-370d-4da1-b14f-39c62355daa0",
      "metadata": {
        "id": "8b25cac1-370d-4da1-b14f-39c62355daa0"
      },
      "source": [
        "We can also select the signal leak \"manually\":"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "17ee19aa-13e2-4eef-8ee5-62a7bc9d7d7a",
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "10009f0a54234636bad8c82fcc1b724b",
            "a5c804a5e3a84bc997addc36db607278",
            "cec081ac8af64fa4bb4abafe71a18f84",
            "4a4f2be12c4a448795f830c9d8c4ac07"
          ]
        },
        "id": "17ee19aa-13e2-4eef-8ee5-62a7bc9d7d7a",
        "outputId": "6e7162d3-ec49-4b46-9c5b-e9d7e3d650c6"
      },
      "outputs": [],
      "source": [
        "initial_noise = randn_tensor((1, 4, 96, 96), device=gpu, dtype=torch.float16, generator=torch.manual_seed(0))\n",
        "signal_leak_zero = torch.tensor([[[[0.0]], [[0.0]], [[0.0]], [[0.0]]]], device=gpu, dtype=torch.float16)\n",
        "signal_leak_white_background = torch.tensor([[[[1.9]], [[1.3]], [[0.0]], [[-1.0]]]], device=gpu, dtype=torch.float16)\n",
        "signal_leak_red_background = torch.tensor([[[[0.0]], [[0.0]], [[-2.5]], [[0.0]]]], device=gpu, dtype=torch.float16)\n",
        "signal_leak_red_background_weaker = torch.tensor([[[[0.0]], [[0.0]], [[-0.5]], [[0.0]]]], device=gpu, dtype=torch.float16)\n",
        "\n",
        "\n",
        "prompt = \"A professional photograph of an astronaut mowing the lawn. Moon in the background. Colorful image, hyperrealistic, fantasy, dark art.\"\n",
        "\n",
        "print(prompt)\n",
        "image1 = pipeline(\n",
        "    prompt,\n",
        "    latents=alpha_prod_T**0.5 * signal_leak_zero + beta_prod_T**0.5 * initial_noise\n",
        ").images[0]\n",
        "display(image1.resize((500,500)))\n",
        "\n",
        "print(prompt)\n",
        "image2 = pipeline(\n",
        "    prompt,\n",
        "    latents=alpha_prod_T**0.5 * signal_leak_white_background + beta_prod_T**0.5 * initial_noise\n",
        ").images[0]\n",
        "display(image2.resize((500,500)))\n",
        "\n",
        "print(prompt)\n",
        "image3 = pipeline(\n",
        "    prompt,\n",
        "    latents=alpha_prod_T**0.5 * signal_leak_red_background + beta_prod_T**0.5 * initial_noise\n",
        ").images[0]\n",
        "display(image3.resize((500,500)))\n",
        "\n",
        "print(prompt)\n",
        "image4 = pipeline(\n",
        "    prompt,\n",
        "    latents=alpha_prod_T**0.5 * signal_leak_red_background_weaker + beta_prod_T**0.5 * initial_noise\n",
        ").images[0]\n",
        "display(image4.resize((500,500)))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7a648409-8dad-49cf-8a80-25d4d1ccae35",
      "metadata": {
        "id": "7a648409-8dad-49cf-8a80-25d4d1ccae35"
      },
      "source": [
        "## Conclusion"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6fc9ac6d-e430-4a1c-b4e8-d8392db6f3de",
      "metadata": {
        "id": "6fc9ac6d-e430-4a1c-b4e8-d8392db6f3de"
      },
      "source": [
        "There is a discrepancy between inference and training in Stable Diffusion v1 and v2.\n",
        "They contain a signal leak during training but not during inference.\n",
        "**We propose to inject a signal leak in the initial noise at inference time** to fix this disrepency.\n",
        "We can bias the image generation **toward a desired specific color distribution or a specific style**.\n",
        "This simple step does not require any fine-tuning making it **much simpler** than existing approaches for style or color-specific image generation.\n",
        "We encourage future research to account for training and inference distribution gap when training or fine-tuning diffusion models, and to include a signal leak in the initial noise at inference time as well, in order to mirror the training process and achieve visually more pleasing results.\n",
        "\n",
        "Project page: https://ivrl.github.io/signal-leak-bias/\n",
        "\n",
        "Code: https://github.com/IVRL/signal-leak-bias\n",
        "\n",
        "License: Non-Commercial License https://github.com/IVRL/signal-leak-bias/blob/main/LICENSE"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}