elixir-nx · jonatanklosko · Dec 15, 2022 · Dec 15, 2022
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -625,10 +625,10 @@ defmodule Bumblebee do
   @doc """
   Predicts sample at the previous timestep using the given scheduler.
 
-  Takes the current `sample` and the `noise` predicted by the model at
-  the current timestep. Returns `{state, prev_sample}`, where `state`
-  is the updated scheduler loop state and `prev_sample` is the predicted
-  sample at the previous timestep.
+  Takes the current `sample` and `prediction` (usually noise) returned
+  by the model at the current timestep. Returns `{state, prev_sample}`,
+  where `state` is the updated scheduler loop state and `prev_sample`
+  is the predicted sample at the previous timestep.
 
   Note that some schedulers require several forward passes of the model
   (and a couple calls to this function) to make an actual prediction for
@@ -641,8 +641,8 @@ defmodule Bumblebee do
           Nx.Tensor.t(),
           Nx.Tensor.t()
         ) :: {Bumblebee.Scheduler.state(), Nx.Tensor.t()}
-  def scheduler_step(%module{} = scheduler, state, sample, noise) do
-    module.step(scheduler, state, sample, noise)
+  def scheduler_step(%module{} = scheduler, state, sample, prediction) do
+    module.step(scheduler, state, sample, prediction)
   end
 
   @doc """

diff --git a/lib/bumblebee/diffusion/ddim_scheduler.ex b/lib/bumblebee/diffusion/ddim_scheduler.ex
@@ -19,6 +19,20 @@ defmodule Bumblebee.Diffusion.DdimScheduler do
       default: 0.02,
       doc: "the end value for the beta schedule"
     ],
+    prediction_type: [
+      default: :noise,
+      doc: """
+      prediction type of the denoising model. Either of:
+
+        * `:noise` (default) - the model predicts the noise of the diffusion process
+
+        * `:angular_velocity` - the model predicts velocity in angular parameterization.
+          See Section 2.4 in [Imagen Video: High Definition Video Generation with Diffusion Models](https://imagen.research.google/video/paper.pdf),
+          then Section 4 in [Progressive Distillation for Fast Sampling of Diffusion Models](https://arxiv.org/pdf/2202.00512.pdf)
+          and Appendix D
+
+      """
+    ],
     alpha_clip_strategy: [
       default: :one,
       doc: ~S"""
@@ -138,11 +152,11 @@ defmodule Bumblebee.Diffusion.DdimScheduler do
   end
 
   @impl true
-  def step(scheduler, state, sample, noise) do
-    do_step(scheduler, state, sample, noise)
+  def step(scheduler, state, sample, prediction) do
+    do_step(scheduler, state, sample, prediction)
   end
 
-  defnp do_step(scheduler \\ [], state, sample, noise) do
+  defnp do_step(scheduler \\ [], state, sample, prediction) do
     # See Equation (12)
 
     # Note that in the paper alpha_t represents a cumulative product,
@@ -164,7 +178,21 @@ defmodule Bumblebee.Diffusion.DdimScheduler do
         end
       end
 
-    pred_denoised_sample = (sample - Nx.sqrt(1 - alpha_bar_t) * noise) / Nx.sqrt(alpha_bar_t)
+    {pred_denoised_sample, noise} =
+      case scheduler.prediction_type do
+        :noise ->
+          pred_denoised_sample =
+            (sample - Nx.sqrt(1 - alpha_bar_t) * prediction) / Nx.sqrt(alpha_bar_t)
+
+          {pred_denoised_sample, prediction}
+
+        :angular_velocity ->
+          pred_denoised_sample =
+            Nx.sqrt(alpha_bar_t) * sample - Nx.sqrt(1 - alpha_bar_t) * prediction
+
+          noise = Nx.sqrt(alpha_bar_t) * prediction + Nx.sqrt(1 - alpha_bar_t) * sample
+          {pred_denoised_sample, noise}
+      end
 
     pred_denoised_sample =
       if scheduler.clip_denoised_sample do
@@ -219,6 +247,9 @@ defmodule Bumblebee.Diffusion.DdimScheduler do
           },
           beta_start: {"beta_start", number()},
           beta_end: {"beta_end", number()},
+          prediction_type:
+            {"prediction_type",
+             mapping(%{"epsilon" => :noise, "v_prediction" => :angular_velocity})},
           alpha_clip_strategy: {
             "set_alpha_to_one",
             mapping(%{true => :one, false => :alpha_zero})

diff --git a/lib/bumblebee/diffusion/layers/unet.ex b/lib/bumblebee/diffusion/layers/unet.ex
@@ -83,6 +83,7 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
     norm_epsilon = opts[:norm_epsilon] || 1.0e-6
     norm_num_groups = opts[:norm_num_groups] || 32
     num_attention_heads = opts[:num_attention_heads] || 1
+    use_linear_projection = Keyword.get(opts, :use_linear_projection, false)
     output_scale_factor = opts[:output_scale_factor] || 1.0
     downsample_padding = opts[:downsample_padding] || [{1, 1}, {1, 1}]
     add_downsample = Keyword.get(opts, :add_downsample, true)
@@ -114,6 +115,7 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
               spatial_transformer(hidden_state, encoder_hidden_state,
                 hidden_size: out_channels,
                 num_heads: num_attention_heads,
+                use_linear_projection: use_linear_projection,
                 depth: 1,
                 name: join(name, "attentions.#{idx}")
               )
@@ -157,6 +159,7 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
     norm_epsilon = opts[:norm_epsilon] || 1.0e-6
     norm_num_groups = opts[:norm_num_groups] || 32
     num_attention_heads = opts[:num_attention_heads] || 1
+    use_linear_projection = Keyword.get(opts, :use_linear_projection, false)
     output_scale_factor = opts[:output_scale_factor] || 1.0
     add_upsample = Keyword.get(opts, :add_upsample, true)
     name = opts[:name]
@@ -187,6 +190,7 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
             spatial_transformer(hidden_state, encoder_hidden_state,
               hidden_size: out_channels,
               num_heads: num_attention_heads,
+              use_linear_projection: use_linear_projection,
               depth: 1,
               name: join(name, "attentions.#{idx}")
             )
@@ -218,6 +222,7 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
     activation = opts[:activation] || :swish
     norm_num_groups = opts[:norm_num_groups] || 32
     num_attention_heads = opts[:num_attention_heads] || 1
+    use_linear_projection = Keyword.get(opts, :use_linear_projection, false)
     output_scale_factor = opts[:output_scale_factor] || 1.0
     name = opts[:name]
 
@@ -245,6 +250,7 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
           encoder_hidden_state,
           hidden_size: channels,
           num_heads: num_attention_heads,
+          use_linear_projection: use_linear_projection,
           depth: 1,
           name: join(name, "attentions.#{idx}")
         )
@@ -260,50 +266,71 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
   defp spatial_transformer(hidden_state, cross_hidden_state, opts) do
     hidden_size = opts[:hidden_size]
     num_heads = opts[:num_heads]
+    use_linear_projection = opts[:use_linear_projection]
     depth = opts[:depth] || 1
     dropout = opts[:dropout] || 0.0
     name = opts[:name]
 
     residual = hidden_state
 
-    hidden_state
-    |> Axon.group_norm(32, epsilon: 1.0e-6, name: join(name, "norm"))
-    |> Axon.conv(hidden_size,
-      kernel_size: 1,
-      strides: 1,
-      padding: :valid,
-      name: join(name, "proj_in")
-    )
-    |> then(
+    flatten_spatial =
       &Axon.layer(
         fn hidden_state, residual, _opts ->
           {b, h, w, c} = Nx.shape(residual)
           Nx.reshape(hidden_state, {b, h * w, c})
         end,
         [&1, residual]
       )
-    )
-    |> spatial_transformer_blocks(cross_hidden_state,
-      hidden_size: hidden_size,
-      num_heads: num_heads,
-      dropout: dropout,
-      depth: depth,
-      name: name
-    )
-    |> then(
+
+    unflatten_spatial =
       &Axon.layer(
         fn hidden_state, residual, _opts ->
           Nx.reshape(hidden_state, Nx.shape(residual))
         end,
         [&1, residual]
       )
+
+    hidden_state
+    |> Axon.group_norm(32, epsilon: 1.0e-6, name: join(name, "norm"))
+    |> then(fn hidden_state ->
+      if use_linear_projection do
+        hidden_state
+        |> flatten_spatial.()
+        |> Axon.dense(hidden_size, name: join(name, "proj_in"))
+      else
+        hidden_state
+        |> Axon.conv(hidden_size,
+          kernel_size: 1,
+          strides: 1,
+          padding: :valid,
+          name: join(name, "proj_in")
+        )
+        |> flatten_spatial.()
+      end
+    end)
+    |> spatial_transformer_blocks(cross_hidden_state,
+      hidden_size: hidden_size,
+      num_heads: num_heads,
+      dropout: dropout,
+      depth: depth,
+      name: name
     )
-    |> Axon.conv(hidden_size,
-      kernel_size: 1,
-      strides: 1,
-      padding: :valid,
-      name: join(name, "proj_out")
-    )
+    |> then(fn hidden_state ->
+      if use_linear_projection do
+        hidden_state
+        |> Axon.dense(hidden_size, name: join(name, "proj_out"))
+        |> unflatten_spatial.()
+      else
+        hidden_state
+        |> unflatten_spatial.()
+        |> Axon.conv(hidden_size,
+          kernel_size: 1,
+          strides: 1,
+          padding: :valid,
+          name: join(name, "proj_out")
+        )
+      end
+    end)
     |> Axon.add(residual)
   end
 

diff --git a/lib/bumblebee/diffusion/pndm_scheduler.ex b/lib/bumblebee/diffusion/pndm_scheduler.ex
@@ -183,8 +183,8 @@ defmodule Bumblebee.Diffusion.PndmScheduler do
   end
 
   @impl true
-  def step(scheduler, state, sample, noise) do
-    do_step(scheduler, state, sample, noise)
+  def step(scheduler, state, sample, prediction) do
+    do_step(scheduler, state, sample, prediction)
   end
 
   defnp do_step(scheduler \\ [], state, sample, noise) do

diff --git a/lib/bumblebee/diffusion/unet_2d_conditional.ex b/lib/bumblebee/diffusion/unet_2d_conditional.ex
@@ -69,12 +69,18 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
     ],
     num_attention_heads: [
       default: 8,
-      doc: "the number of attention heads for each attention layer"
+      doc:
+        "the number of attention heads for each attention layer. Optionally can be a list with one number per block"
     ],
     cross_attention_size: [
       default: 1280,
       doc: "the dimensionality of the cross attention features"
     ],
+    use_linear_projection: [
+      default: false,
+      doc:
+        "whether the input/output projection of the transformer block should be linear or convolutional"
+    ],
     activation: [
       default: :silu,
       doc: "the activation function"
@@ -228,15 +234,18 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
 
   defp down_blocks(sample, timestep_embedding, encoder_hidden_state, spec, opts) do
     name = opts[:name]
-    blocks = Enum.zip(spec.hidden_sizes, spec.down_block_types)
+
+    blocks =
+      Enum.zip([spec.hidden_sizes, spec.down_block_types, num_attention_heads_per_block(spec)])
 
     in_channels = hd(spec.hidden_sizes)
     down_block_residuals = [{sample, in_channels}]
 
     state = {sample, down_block_residuals, in_channels}
 
     {sample, down_block_residuals, _} =
-      for {{out_channels, block_type}, idx} <- Enum.with_index(blocks), reduce: state do
+      for {{out_channels, block_type, num_attention_heads}, idx} <- Enum.with_index(blocks),
+          reduce: state do
         {sample, down_block_residuals, in_channels} ->
           last_block? = idx == length(spec.hidden_sizes) - 1
 
@@ -254,7 +263,8 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
               activation: spec.activation,
               norm_epsilon: spec.group_norm_epsilon,
               norm_num_groups: spec.group_norm_num_groups,
-              num_attention_heads: spec.num_attention_heads,
+              num_attention_heads: num_attention_heads,
+              use_linear_projection: spec.use_linear_projection,
               name: join(name, idx)
             )
 
@@ -274,7 +284,8 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
       norm_epsilon: spec.group_norm_epsilon,
       norm_num_groups: spec.group_norm_num_groups,
       output_scale_factor: spec.mid_block_scale_factor,
-      num_attention_heads: spec.num_attention_heads,
+      num_attention_heads: spec |> num_attention_heads_per_block() |> List.last(),
+      use_linear_projection: spec.use_linear_projection,
       name: opts[:name]
     )
   end
@@ -298,13 +309,23 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
     reversed_hidden_sizes = Enum.reverse(spec.hidden_sizes)
     in_channels = hd(reversed_hidden_sizes)
 
+    num_attention_heads_per_block =
+      spec
+      |> num_attention_heads_per_block()
+      |> Enum.reverse()
+
     blocks_and_chunks =
-      [reversed_hidden_sizes, spec.up_block_types, down_block_residuals]
+      [
+        reversed_hidden_sizes,
+        spec.up_block_types,
+        num_attention_heads_per_block,
+        down_block_residuals
+      ]
       |> Enum.zip()
       |> Enum.with_index()
 
     {sample, _} =
-      for {{out_channels, block_type, residuals}, idx} <- blocks_and_chunks,
+      for {{out_channels, block_type, num_attention_heads, residuals}, idx} <- blocks_and_chunks,
           reduce: {sample, in_channels} do
         {sample, in_channels} ->
           last_block? = idx == length(spec.hidden_sizes) - 1
@@ -323,7 +344,8 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
               norm_epsilon: spec.group_norm_epsilon,
               norm_num_groups: spec.group_norm_num_groups,
               activation: spec.activation,
-              num_attention_heads: spec.num_attention_heads,
+              num_attention_heads: num_attention_heads,
+              use_linear_projection: spec.use_linear_projection,
               name: join(name, idx)
             )
 
@@ -333,6 +355,15 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
     sample
   end
 
+  defp num_attention_heads_per_block(spec) when is_list(spec.num_attention_heads) do
+    spec.num_attention_heads
+  end
+
+  defp num_attention_heads_per_block(spec) when is_integer(spec.num_attention_heads) do
+    num_blocks = length(spec.down_block_types)
+    List.duplicate(spec.num_attention_heads, num_blocks)
+  end
+
   defimpl Bumblebee.HuggingFace.Transformers.Config do
     def load(spec, data) do
       import Shared.Converters
@@ -367,8 +398,9 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
           },
           downsample_padding: {"downsample_padding", padding(2)},
           mid_block_scale_factor: {"mid_block_scale_factor", number()},
-          num_attention_heads: {"attention_head_dim", number()},
+          num_attention_heads: {"attention_head_dim", one_of([number(), list(number())])},
           cross_attention_size: {"cross_attention_dim", number()},
+          use_linear_projection: {"use_linear_projection", boolean()},
           activation: {"act_fn", atom()},
           group_norm_num_groups: {"norm_num_groups", number()},
           group_norm_epsilon: {"norm_eps", number()}

diff --git a/lib/bumblebee/scheduler.ex b/lib/bumblebee/scheduler.ex
@@ -56,15 +56,15 @@ defmodule Bumblebee.Scheduler do
   @doc """
   Predicts sample at the previous timestep.
 
-  Takes the current `sample` and the `noise` predicted by the model at
-  the current timestep. Returns `{state, prev_sample}`, where `state`
-  is the updated state and `prev_sample` is the predicted sample at the
-  previous timestep.
+  Takes the current `sample` and `prediction` (usually noise) returned
+  by the model at the current timestep. Returns `{state, prev_sample}`,
+  where `state` is the updated state and `prev_sample` is the predicted
+  sample at the previous timestep.
   """
   @callback step(
               t(),
               state(),
               sample :: Nx.Tensor.t(),
-              noise :: Nx.Tensor.t()
+              prediction :: Nx.Tensor.t()
             ) :: {state :: map(), prev_sample :: Nx.Tensor.t()}
 end