Add Stable Diffusion ControlNet

lib/bumblebee.ex

@@ -116,6 +116,7 @@ defmodule Bumblebee do
     "CLIPModel" => {Bumblebee.Multimodal.Clip, :base},
     "CLIPTextModel" => {Bumblebee.Text.ClipText, :base},
     "CLIPVisionModel" => {Bumblebee.Vision.ClipVision, :base},
+    "ControlNetModel" => {Bumblebee.Diffusion.ControlNet, :base},
     "ConvNextForImageClassification" => {Bumblebee.Vision.ConvNext, :for_image_classification},
     "ConvNextModel" => {Bumblebee.Vision.ConvNext, :base},
     "DeiTForImageClassification" => {Bumblebee.Vision.Deit, :for_image_classification},

lib/bumblebee/conversion/pytorch_params.ex

@@ -300,8 +300,18 @@ defmodule Bumblebee.Conversion.PyTorchParams do
   defp match_template(name, template), do: match_template(name, template, %{})
 
   defp match_template(<<_, _::binary>> = name, <<"{", template::binary>>, substitutes) do
-    [value, name] = String.split(name, ".", parts: 2)
-    [key, template] = String.split(template, "}.", parts: 2)
+    {value, name} =
+      case String.split(name, ".", parts: 2) do
+        [value] -> {value, ""}
+        [value, name] -> {value, name}
+      end
+
+    {key, template} =
+      case String.split(template, "}", parts: 2) do
+        [key, ""] -> {key, ""}
+        [key, "." <> template] -> {key, template}
+      end
+
     match_template(name, template, put_in(substitutes[key], value))
   end