Add support for controlnet to LCM converter

coreml_suite/lcm/lcm_converter.py

@@ -41,10 +41,7 @@ def get_unets():
     cml_unet = CoreMLUNet2DConditionModel().eval()
     cml_unet.load_state_dict(ref_unet.state_dict(), strict=False)
 
-    del ref_unet
-    gc.collect()
-
-    return cml_unet, ref_config
+    return cml_unet, ref_unet
 
 
 def get_encoder_hidden_states_shape(unet_config, batch_size):
@@ -101,7 +98,7 @@ def load_coreml_model(out_path):
 
 
 def convert_to_coreml(
-    submodule_name, torchscript_module, sample_inputs, output_names, out_path
+        submodule_name, torchscript_module, sample_inputs, output_names, out_path
 ):
     if os.path.exists(out_path):
         logger.info(f"Skipping export because {out_path} already exists")
@@ -162,37 +159,92 @@ def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape, sche
             ("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)),
         ]
     )
+    return sample_unet_inputs
+
+
+def get_unet_inputs_spec(sample_unet_inputs):
     sample_unet_inputs_spec = {
         k: (v.shape, v.dtype) for k, v in sample_unet_inputs.items()
     }
-    return sample_unet_inputs, sample_unet_inputs_spec
+    return sample_unet_inputs_spec
+
+
+def add_cnet_support(sample_shape, reference_unet):
+    from python_coreml_stable_diffusion.unet import calculate_conv2d_output_shape
+    additional_residuals_shapes = []
+
+    batch_size = sample_shape[0]
+    h, w = sample_shape[2:]
+
+    # conv_in
+    out_h, out_w = calculate_conv2d_output_shape(
+        h,
+        w,
+        reference_unet.conv_in,
+    )
+    additional_residuals_shapes.append(
+        (batch_size, reference_unet.conv_in.out_channels, out_h, out_w))
+
+    # down_blocks
+    for down_block in reference_unet.down_blocks:
+        additional_residuals_shapes += [
+            (batch_size, resnet.out_channels, out_h, out_w) for resnet in
+            down_block.resnets
+        ]
+        if hasattr(down_block,
+                   "downsamplers") and down_block.downsamplers is not None:
+            for downsampler in down_block.downsamplers:
+                out_h, out_w = calculate_conv2d_output_shape(out_h, out_w,
+                                                             downsampler.conv)
+            additional_residuals_shapes.append(
+                (batch_size, down_block.downsamplers[-1].conv.out_channels, out_h,
+                 out_w))
+
+    # mid_block
+    additional_residuals_shapes.append(
+        (
+            batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w)
+    )
+
+    additional_inputs = {}
+    for i, shape in enumerate(additional_residuals_shapes):
+        sample_residual_input = torch.rand(*shape)
+        additional_inputs[f"additional_residual_{i}"] = sample_residual_input
+
+    return additional_inputs
 
 
 def convert(
-    out_path: str, batch_size: int = 1, sample_size: tuple[int, int] = (64, 64)
+        out_path: str, batch_size: int = 1, sample_size: tuple[int, int] = (64, 64),
+        controlnet_support: bool = False
 ):
-    coreml_unet, unet_config = get_unets()
+    coreml_unet, ref_unet = get_unets()
 
     sample_shape = (
         batch_size,  # B
-        unet_config.in_channels,  # C
+        ref_unet.config.in_channels,  # C
         sample_size[0],  # H
         sample_size[1],  # W
     )
 
     encoder_hidden_states_shape = get_encoder_hidden_states_shape(
-        unet_config, batch_size
+        ref_unet.config, batch_size
     )
 
     scheduler = get_scheduler()
 
-    sample_inputs, sample_inputs_spec = get_sample_input(
+    sample_inputs = get_sample_input(
         batch_size, encoder_hidden_states_shape, sample_shape, scheduler
     )
 
+    if controlnet_support:
+        sample_inputs |= add_cnet_support(sample_shape, ref_unet)
+
+    sample_inputs_spec = get_unet_inputs_spec(sample_inputs)
+
     logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}")
     logger.info("JIT tracing..")
-    traced_unet = torch.jit.trace(coreml_unet, example_kwarg_inputs=sample_inputs)
+    traced_unet = torch.jit.trace(coreml_unet, example_inputs=list(sample_inputs.values()))
     logger.info("Done.")
 
     coreml_sample_inputs = get_coreml_inputs(sample_inputs)
@@ -223,7 +275,9 @@ if __name__ == "__main__":
     sample_size = (h // 8, w // 8)
     batch_size = 4
 
-    out_name = f"{MODEL_NAME}_{w}x{h}_batch{batch_size}"
+    cn_support_str = "_cn" if True else ""
+
+    out_name = f"{MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
 
     out_path = get_out_path("unet", f"{out_name}")
     if not os.path.exists(out_path):

coreml_suite/lcm/nodes.py

@@ -21,7 +21,8 @@ class CoreMLConverterLCM:
                                      ComputeUnit.CPU_AND_GPU.name,
                                      ComputeUnit.ALL.name,
                                      ComputeUnit.CPU_ONLY.name,
-                                 ],)
+                                 ],),
+                "controlnet_support": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -29,7 +30,7 @@ class CoreMLConverterLCM:
     RETURN_NAMES = ("coreml_model",)
     FUNCTION = "convert"
 
-    def convert(self, height, width, batch_size, compute_unit):
+    def convert(self, height, width, batch_size, compute_unit, controlnet_support):
         """Converts a LCM model to Core ML.
 
         Args:
@@ -48,14 +49,15 @@ class CoreMLConverterLCM:
         w = width
         sample_size = (h // 8, w // 8)
         batch_size = batch_size
+        cn_support_str = "_cn" if controlnet_support else ""
 
-        out_name = f"{lcm_converter.MODEL_NAME}_{w}x{h}_batch{batch_size}"
+        out_name = f"{lcm_converter.MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
 
         out_path = lcm_converter.get_out_path("unet", f"{out_name}")
 
         if not os.path.exists(out_path):
             lcm_converter.convert(
-                out_path=out_path, sample_size=sample_size, batch_size=batch_size
+                out_path=out_path, sample_size=sample_size, batch_size=batch_size, controlnet_support=controlnet_support
             )
         target_path = lcm_converter.compile_model(out_path=out_path, out_name=out_name)