Fix loading more recent VaeKl checkpoints (#305)

Author: jonatanklosko

lib/bumblebee/conversion/pytorch.ex

@@ -106,35 +106,42 @@ defmodule Bumblebee.Conversion.PyTorch do
           Enum.reduce(layer.parameters, {[], diff}, fn param, {params, diff} ->
             param_expr = params_expr[layer_name][param.name]
 
-            {sources, source_fun} =
+            {sources, builder_fun} =
               case params_source do
-                %{} = layer_params_mapping ->
-                  if info = layer_params_mapping[param.name] do
-                    info
+                %{} = param_builders ->
+                  if param_builder = param_builders[param.name] do
+                    param_builder
                   else
-                    raise "no matching mapping found for parameter #{inspect(param.name)} in #{inspect(layer_params_mapping)}"
+                    raise "no matching mapping found for parameter #{inspect(param.name)} in #{inspect(param_builders)}"
                   end
 
-                source_layer_name when is_binary(source_layer_name) ->
-                  default_layer_param_source(layer, param.name, source_layer_name)
+                source_layer_name
+                when is_binary(source_layer_name) or
+                       is_list(source_layer_name) ->
+                  default_layer_param_builder(layer, param.name, source_layer_name)
               end
 
             {all_sources_found?, source_values, source_keys} =
-              for {source_layer_name, source_param_name} <- sources, reduce: {true, [], []} do
+              for source <- sources, reduce: {true, [], []} do
                 {all_found?, values, keys} ->
-                  source_param_names = List.wrap(source_param_name)
-
-                  case lookup_param(pytorch_state, source_layer_name, source_param_names) do
-                    {:ok, value, key} -> {all_found?, [value | values], [key | keys]}
-                    :error -> {false, values, keys}
+                  # Source can be either {layer_name, param_name}, or
+                  # a list of these, to find any match
+                  source
+                  |> List.wrap()
+                  |> Enum.find_value(fn {source_layer_name, source_param_name} ->
+                    lookup_param(pytorch_state, source_layer_name, source_param_name)
+                  end)
+                  |> case do
+                    {value, key} -> {all_found?, [value | values], [key | keys]}
+                    nil -> {false, values, keys}
                   end
               end
 
             diff = prepend(diff, :used_keys, source_keys)
 
             {value, diff} =
               if all_sources_found? do
-                value = source_fun.(Enum.reverse(source_values))
+                value = builder_fun.(Enum.reverse(source_values))
 
                 case verify_param_shape(param_expr, value) do
                   :ok ->
@@ -186,11 +193,15 @@ defmodule Bumblebee.Conversion.PyTorch do
 
     source_templates =
       Enum.flat_map(params_mapping, fn
-        {_target_template, %{} = params_source} ->
-          for {_target_param_name, {sources, _source_fun}} <- params_source,
-              {source_template, _source_param_name} <- sources,
+        {_target_template, %{} = param_builders} ->
+          for {_target_param_name, {sources, _builder_fun}} <- param_builders,
+              ref_or_refs <- sources,
+              {source_template, _source_param_name} <- List.wrap(ref_or_refs),
               do: source_template
 
+        {_target_template, source_templates} when is_list(source_templates) ->
+          source_templates
+
         {_target_template, source_template} when is_binary(source_template) ->
           [source_template]
       end)
@@ -339,17 +350,17 @@ defmodule Bumblebee.Conversion.PyTorch do
 
   defp format_list(items), do: Enum.map_join(items, "\n", &("  * " <> &1))
 
-  defp default_layer_param_source(%{op_name: :dense}, "kernel", layer_name) do
-    {[{layer_name, "weight"}],
+  defp default_layer_param_builder(%{op_name: :dense}, "kernel", layer_name) do
+    {[param_refs(layer_name, "weight")],
      fn [kernel] ->
        [out_features, in_features] = Nx.axes(kernel)
        Nx.transpose(kernel, axes: [in_features, out_features])
      end}
   end
 
-  defp default_layer_param_source(layer, "kernel", layer_name)
+  defp default_layer_param_builder(layer, "kernel", layer_name)
        when layer.op_name in [:conv, :depthwise_conv] do
-    {[{layer_name, "weight"}],
+    {[param_refs(layer_name, "weight")],
      fn [kernel] ->
        [out_channels, in_channels | kernel_spatials] = Nx.axes(kernel)
 
@@ -360,8 +371,8 @@ defmodule Bumblebee.Conversion.PyTorch do
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :conv_transpose} = layer, "kernel", layer_name) do
-    {[{layer_name, "weight"}],
+  defp default_layer_param_builder(%{op_name: :conv_transpose} = layer, "kernel", layer_name) do
+    {[param_refs(layer_name, "weight")],
      fn [kernel] ->
        [in_channels, out_channels | kernel_spatials] = Nx.axes(kernel)
 
@@ -372,57 +383,57 @@ defmodule Bumblebee.Conversion.PyTorch do
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :lstm}, "bias", layer_name) do
-    {[{layer_name, "bias_hh"}, {layer_name, "bias_ih"}],
+  defp default_layer_param_builder(%{op_name: :lstm}, "bias", layer_name) do
+    {[param_refs(layer_name, "bias_hh"), param_refs(layer_name, "bias_ih")],
      fn [bias_hh, bias_ih] ->
        bias = Nx.add(bias_ih, bias_hh)
        bias = Nx.reshape(bias, {4, :auto})
        {bias[0], bias[1], bias[2], bias[3]}
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :lstm}, "input_kernel", layer_name) do
-    {[{layer_name, "weight_ih"}],
+  defp default_layer_param_builder(%{op_name: :lstm}, "input_kernel", layer_name) do
+    {[param_refs(layer_name, "weight_ih")],
      fn [weight_ih] ->
        weight_ih = weight_ih |> unflatten_leading(4) |> Nx.transpose(axes: [0, 2, 1])
        {weight_ih[0], weight_ih[1], weight_ih[2], weight_ih[3]}
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :lstm}, "hidden_kernel", layer_name) do
-    {[{layer_name, "weight_hh"}],
+  defp default_layer_param_builder(%{op_name: :lstm}, "hidden_kernel", layer_name) do
+    {[param_refs(layer_name, "weight_hh")],
      fn [weight_hh] ->
        weight_hh = weight_hh |> unflatten_leading(4) |> Nx.transpose(axes: [0, 2, 1])
        {weight_hh[0], weight_hh[1], weight_hh[2], weight_hh[3]}
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :gru}, "bias", layer_name) do
-    {[{layer_name, "bias_hh"}, {layer_name, "bias_ih"}],
+  defp default_layer_param_builder(%{op_name: :gru}, "bias", layer_name) do
+    {[param_refs(layer_name, "bias_hh"), param_refs(layer_name, "bias_ih")],
      fn [bias_hh, bias_ih] ->
        bias_hh = unflatten_leading(bias_hh, 3)
        bias_ih = unflatten_leading(bias_ih, 3)
        {Nx.add(bias_ih[0], bias_hh[0]), Nx.add(bias_ih[1], bias_hh[1]), bias_ih[2], bias_hh[2]}
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :gru}, "input_kernel", layer_name) do
-    {[{layer_name, "weight_ih"}],
+  defp default_layer_param_builder(%{op_name: :gru}, "input_kernel", layer_name) do
+    {[param_refs(layer_name, "weight_ih")],
      fn [weight_ih] ->
        weight_ih = weight_ih |> unflatten_leading(4) |> Nx.transpose(axes: [0, 2, 1])
        {weight_ih[0], weight_ih[1], weight_ih[2]}
      end}
   end
 
-  defp default_layer_param_source(%{op_name: :gru}, "hidden_kernel", layer_name) do
-    {[{layer_name, "weight_hh"}],
+  defp default_layer_param_builder(%{op_name: :gru}, "hidden_kernel", layer_name) do
+    {[param_refs(layer_name, "weight_hh")],
      fn [weight_hh] ->
        weight_hh = weight_hh |> unflatten_leading(3) |> Nx.transpose(axes: [0, 2, 1])
        {weight_hh[0], weight_hh[1], weight_hh[2]}
      end}
   end
 
-  defp default_layer_param_source(_layer, param_name, layer_name) do
+  defp default_layer_param_builder(_layer, param_name, layer_name) do
     pytorch_names =
       case param_name do
         # PyTorch uses "weight" instead of "kernel" everywhere
@@ -440,18 +451,26 @@ defmodule Bumblebee.Conversion.PyTorch do
         name -> [name]
       end
 
-    {[{layer_name, pytorch_names}], fn [value] -> value end}
+    param_source = Enum.flat_map(pytorch_names, &param_refs(layer_name, &1))
+
+    {[param_source], fn [value] -> value end}
+  end
+
+  defp param_refs(layer_name, param_name) do
+    for layer_name <- List.wrap(layer_name) do
+      {layer_name, param_name}
+    end
   end
 
-  defp lookup_param(pytorch_state, layer_name, pytorch_names) do
+  defp lookup_param(pytorch_state, layer_name, pytorch_name) do
     # Note: the PyTorch model may have some root-level parameters that
-    # we need to namespace under a layer in Axon, so after trying params
-    # within layer_name, we also try the parameter name directly
-    pytorch_keys = Enum.map(pytorch_names, &(layer_name <> "." <> &1)) ++ pytorch_names
+    # we need to namespace under a layer in Axon, so after trying the
+    # param within layer_name, we also try the param name directly
+    pytorch_keys = [layer_name <> "." <> pytorch_name, pytorch_name]
 
-    Enum.find_value(pytorch_keys, :error, fn pytorch_key ->
+    Enum.find_value(pytorch_keys, fn pytorch_key ->
       if value = pytorch_state[pytorch_key] do
-        {:ok, value, pytorch_key}
+        {value, pytorch_key}
       end
     end)
   end

lib/bumblebee/diffusion/unet_2d_conditional.ex

@@ -407,6 +407,8 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
   end
 
   defimpl Bumblebee.HuggingFace.Transformers.Model do
+    alias Bumblebee.HuggingFace.Transformers
+
     def params_mapping(_spec) do
       block_mapping = %{
         "transformers.{m}.norm" => "attentions.{m}.norm",
@@ -449,10 +451,9 @@ defmodule Bumblebee.Diffusion.UNet2DConditional do
       }
 
       blocks_mapping =
-        for {target, source} <- block_mapping,
-            prefix <- ["down_blocks.{n}", "mid_block", "up_blocks.{n}"],
-            do: {prefix <> "." <> target, prefix <> "." <> source},
-            into: %{}
+        ["down_blocks.{n}", "mid_block", "up_blocks.{n}"]
+        |> Enum.map(&Transformers.Utils.prefix_params_mapping(block_mapping, &1, &1))
+        |> Enum.reduce(&Map.merge/2)
 
       %{
         "time_embedding.intermediate" => "time_embedding.linear_1",

lib/bumblebee/diffusion/vae_kl.ex

@@ -443,13 +443,16 @@ defmodule Bumblebee.Diffusion.VaeKl do
   end
 
   defimpl Bumblebee.HuggingFace.Transformers.Model do
+    alias Bumblebee.HuggingFace.Transformers
+
     def params_mapping(_spec) do
       block_mapping = %{
         "attentions.{m}.norm" => "attentions.{m}.group_norm",
-        "attentions.{m}.query" => "attentions.{m}.query",
-        "attentions.{m}.key" => "attentions.{m}.key",
-        "attentions.{m}.value" => "attentions.{m}.value",
-        "attentions.{m}.output" => "attentions.{m}.proj_attn",
+        # The layer name has been renamed upstream, so we try both
+        "attentions.{m}.query" => ["attentions.{m}.to_q", "attentions.{m}.query"],
+        "attentions.{m}.key" => ["attentions.{m}.to_k", "attentions.{m}.key"],
+        "attentions.{m}.value" => ["attentions.{m}.to_v", "attentions.{m}.value"],
+        "attentions.{m}.output" => ["attentions.{m}.to_out.0", "attentions.{m}.proj_attn"],
         "residual_blocks.{m}.norm_1" => "resnets.{m}.norm1",
         "residual_blocks.{m}.conv_1" => "resnets.{m}.conv1",
         "residual_blocks.{m}.norm_2" => "resnets.{m}.norm2",
@@ -460,15 +463,14 @@ defmodule Bumblebee.Diffusion.VaeKl do
       }
 
       blocks_mapping =
-        for {target, source} <- block_mapping,
-            prefix <- [
-              "encoder.down_blocks.{n}",
-              "encoder.mid_block",
-              "decoder.mid_block",
-              "decoder.up_blocks.{n}"
-            ],
-            do: {prefix <> "." <> target, prefix <> "." <> source},
-            into: %{}
+        [
+          "encoder.down_blocks.{n}",
+          "encoder.mid_block",
+          "decoder.mid_block",
+          "decoder.up_blocks.{n}"
+        ]
+        |> Enum.map(&Transformers.Utils.prefix_params_mapping(block_mapping, &1, &1))
+        |> Enum.reduce(&Map.merge/2)
 
       %{
         "encoder.input_conv" => "encoder.conv_in",

lib/bumblebee/huggingface/transformers/model.ex

@@ -4,16 +4,17 @@ defprotocol Bumblebee.HuggingFace.Transformers.Model do
   # This protocol defines details related to loading Bumblebee model
   # from huggingface/transformers model.
 
-  @type params_mapping :: %{
-          layer_name() => layer_name() | params_source()
-        }
+  @type params_mapping :: %{layer_name() => params_source()}
 
-  @type params_source :: %{
-          param_name() =>
-            {list(source()), (list(Nx.tensor()) -> Nx.Tensor.t() | Nx.Container.t())}
-        }
+  @type params_source :: layer_name() | list(layer_name()) | param_builders()
 
-  @type source :: {layer_name(), param_name() | list(param_name())}
+  @type param_builders :: %{param_name() => param_builder()}
+
+  @type param_builder ::
+          {list(param_source()), (list(Nx.tensor()) -> Nx.Tensor.t() | Nx.Container.t())}
+
+  @type param_source :: param_ref() | list(param_ref())
+  @type param_ref :: {layer_name(), param_name()}
 
   @type layer_name :: String.t()
   @type param_name :: String.t()
@@ -53,8 +54,8 @@ defprotocol Bumblebee.HuggingFace.Transformers.Model do
   automatically.
 
   In some cases, particularly with model-specific layers/parameters,
-  we may need more control over the parameter mapping. In such cases, instead
-  of source layer name, a map with parameter-level transformations
+  we may need more control over the parameter mapping. In such cases,
+  instead of source layer name, a map with parameter-level transformations
   may be specified:
 
       %{
@@ -69,9 +70,10 @@ defprotocol Bumblebee.HuggingFace.Transformers.Model do
 
   For each parameter, we specify a list of source parameters in the
   form of `{source_layer_name, source_param_name}`, then a function
-  to build our parameter value. Multiple source parameter names to
-  try may be specified. With the explicit transformation we can
-  handle arbitrary parameter name and value transformations.
+  to build our parameter value. Instead of a single tuple, we can
+  specify a list of those to try one by one. With the explicit
+  transformation we can handle arbitrary parameter name and value
+  transformations.
   """
   @spec params_mapping(t()) :: params_mapping()
   def params_mapping(spec)

lib/bumblebee/huggingface/transformers/utils.ex

@@ -28,14 +28,28 @@ defmodule Bumblebee.HuggingFace.Transformers.Utils do
   @spec map_params_source_layer_names(
           Transformers.Model.params_source(),
           (String.t() -> String.t())
-        ) :: Transformers.Model.layer_name() | Transformers.Model.params_source()
-  def map_params_source_layer_names(%{} = params_source, fun) do
-    Map.new(params_source, fn {param_name, {sources, source_fun}} ->
-      sources = for {layer_name, param_name} <- sources, do: {fun.(layer_name), param_name}
-      {param_name, {sources, source_fun}}
+        ) :: Transformers.Model.params_source()
+  def map_params_source_layer_names(%{} = param_builders, fun) do
+    Map.new(param_builders, fn {param_name, {sources, builder_fun}} ->
+      sources =
+        for ref_or_refs <- sources do
+          case ref_or_refs do
+            {layer_name, param_name} ->
+              {fun.(layer_name), param_name}
+
+            refs ->
+              for {layer_name, param_name} <- refs, do: {fun.(layer_name), param_name}
+          end
+        end
+
+      {param_name, {sources, builder_fun}}
     end)
   end
 
+  def map_params_source_layer_names(layer_names, fun) when is_list(layer_names) do
+    Enum.map(layer_names, fun)
+  end
+
   def map_params_source_layer_names(layer_name, fun) when is_binary(layer_name) do
     fun.(layer_name)
   end