diff --git a/coreml_suite/lcm/lcm_converter.py b/coreml_suite/lcm/lcm_converter.py index 6356586..be3b4cc 100644 --- a/coreml_suite/lcm/lcm_converter.py +++ b/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): diff --git a/coreml_suite/lcm/nodes.py b/coreml_suite/lcm/nodes.py index 9028668..296d8a9 100644 --- a/coreml_suite/lcm/nodes.py +++ b/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)