chore: import upstream snapshot with attribution

python_coreml_stable_diffusion/__init__.py

@@ -0,0 +1 @@
+from ._version import __version__

python_coreml_stable_diffusion/_version.py

@@ -0,0 +1 @@
+__version__ = "1.1.0"

python_coreml_stable_diffusion/activation_quantization.py

@@ -0,0 +1,503 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+import logging
+import operator
+
+import torch
+
+logging.basicConfig()
+logger = logging.getLogger()
+logger.setLevel('INFO')
+
+import argparse
+import gc
+import json
+import os
+import pickle
+from copy import deepcopy
+
+import coremltools as ct
+import numpy as np
+from coremltools.optimize.torch.quantization import (
+    LinearQuantizer, LinearQuantizerConfig, ModuleLinearQuantizerConfig)
+from diffusers import StableDiffusionPipeline
+from tqdm import tqdm
+
+from python_coreml_stable_diffusion import attention
+from python_coreml_stable_diffusion import unet
+from python_coreml_stable_diffusion.layer_norm import LayerNormANE
+from python_coreml_stable_diffusion.torch2coreml import compute_psnr
+from python_coreml_stable_diffusion.unet import Einsum
+
+attention.SPLIT_SOFTMAX = True
+
+CALIBRATION_DATA = [
+    "image of a transparent tall glass with ice, fruits and mint, photograph, commercial, food, warm background, beautiful image, detailed",
+    "picture of dimly lit living room, minimalist furniture, vaulted ceiling, huge room, floor to ceiling window with an ocean view, nighttime, 3D render, high quality, detailed",
+    "modern office building, 8 stories tall, glass and steel, 3D render style, wide angle view, very detailed, sharp photographic image, in an office park, bright sunny day, clear blue skies, trees and landscaping",
+    "cute small cat sitting in a movie theater eating popcorn, watching a movie, cozy indoor lighting, detailed, digital painting, character design",
+    "a highly detailed matte painting of a man on a hill watching a rocket launch in the distance by studio ghibli, volumetric lighting, octane render, 4K resolution, hyperrealism, highly detailed, insanely detailed, cinematic lighting, depth of field",
+    "an undersea world with several of fish, rocks, detailed, realistic, photograph, amazing, beautiful, high resolution",
+    "large ocean wave hitting a beach at sunset, photograph, detailed",
+    "pocket watch on a table, close up. macro, sharp, high gloss, brass, gears, sharp, detailed",
+    "pocket watch in the style of pablo picasso, painting",
+    "majestic royal tall ship on a calm sea, realistic painting, cloudy blue sky, in the style of edward hopper",
+    "german castle on a mountain, blue sky, realistic, photograph, dramatic, wide angle view",
+    "artificial intelligence, AI, concept art, blue line sketch",
+    "a humanoid robot, concept art, 3D render, high quality, detailed",
+    "donut with sprinkles and a cup of coffee on a wood table, detailed, photograph",
+    "orchard at sunset, beautiful, photograph, great composition, detailed, realistic, HDR",
+    "image of a map of a country, tattered, old, styled, illustration, for a video game style",
+    "blue and green woven fibers, nano fiber material, detailed, concept art, micro photography",
+]
+
+RANDOM_TEST_DATA = [
+    "a black and brown dog standing outside a door.",
+    "a person on a motorcycle makes a turn on the track.",
+    "inflatable boats sit on the arizona river, and on the bank",
+    "a white cat sitting under a white umbrella",
+    "black bear standing in a field of grass under a tree.",
+    "a train that is parked on tracks and has graffiti writing on it, with a mountain range in the background.",
+    "a cake inside of a pan sitting in an oven.",
+    "a table with paper plates and flowers in a home",
+]
+
+def get_coreml_inputs(sample_inputs):
+    return [
+        ct.TensorType(
+            name=k,
+            shape=v.shape,
+            dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
+        ) for k, v in sample_inputs.items()
+    ]
+
+def convert_to_coreml(torchscript_module, sample_inputs):
+    logger.info("Converting model to CoreML..")
+    coreml_model = ct.convert(
+        torchscript_module,
+        convert_to="mlprogram",
+        minimum_deployment_target=ct.target.macOS14,
+        inputs=get_coreml_inputs(sample_inputs),
+        outputs=[ct.TensorType(name="noise_pred", dtype=np.float32)],
+        compute_units=ct.ComputeUnit.ALL,
+        skip_model_load=True,
+    )
+
+    return coreml_model
+
+
+def unet_data_loader(data_dir, device='cpu', calibration_nsamples=None):
+    """
+    Load calibration data from specified path.
+    Limit number of samples to calibration_nsamples, if specified.
+    """
+    dataloader = []
+    skip_load = False
+    for file in sorted(os.listdir(data_dir)):
+        if file.endswith('.pkl'):
+            filepath = os.path.join(data_dir, file)
+            with open(filepath, 'rb') as data:
+                try:
+                    while not skip_load:
+                        unet_data = pickle.load(data)
+                        for input in unet_data:
+                            dataloader.append([x.to(torch.float).to(device) for x in input])
+
+                            if calibration_nsamples:
+                                if len(dataloader) >= calibration_nsamples:
+                                    skip_load = True
+                                    break
+                except EOFError:
+                    pass
+        if skip_load: 
+            break
+
+    logger.info(f"Total calibration samples: {len(dataloader)}")
+    return dataloader
+
+def quantize_module_config(module_name):
+    """
+    Generate quantization config to apply W8A8 quantization for specified module.
+    Rest of the model is kept in FP32 precision.
+    """
+    config = LinearQuantizerConfig(
+        global_config=ModuleLinearQuantizerConfig(
+            milestones=[0, 1000, 1000, 0],
+            weight_dtype=torch.float32,
+            activation_dtype=torch.float32,
+        ),
+        module_name_configs={
+            module_name: ModuleLinearQuantizerConfig(
+                quantization_scheme="symmetric",
+                milestones=[0, 1000, 1000, 0],
+            ),
+        },
+    )
+    return config
+
+def quantize_cumulative_config(skip_conv_layers, skip_einsum_layers):
+    """
+    Generate quantization config to apply W8A8 quantization.
+    Skipped layers are kept in W8A32 precision.
+    """
+    logger.info(f"Skipping {len(skip_conv_layers)} conv layers and {len(skip_einsum_layers)} einsum layers")
+    w8config = ModuleLinearQuantizerConfig(
+            quantization_scheme="symmetric",
+            milestones=[0, 1000, 1000, 0],
+            activation_dtype=torch.float32)
+    conv_modules_config = {name: w8config for name in skip_conv_layers}
+    einsum_modules_config = {name: w8config for name in skip_einsum_layers}
+    module_name_config = {}
+    module_name_config.update(conv_modules_config)
+    module_name_config.update(einsum_modules_config)
+
+    config = LinearQuantizerConfig(
+        global_config=ModuleLinearQuantizerConfig(
+            quantization_scheme="symmetric",
+            milestones=[0, 1000, 1000, 0],
+        ),
+        module_name_configs=module_name_config,
+        module_type_configs={
+            torch.cat: None,
+            torch.nn.GroupNorm: None,
+            torch.nn.SiLU: None,
+            torch.nn.functional.gelu: None,
+            operator.add: None,
+        },
+    )
+    return config
+
+def quantize(model, config, calibration_data):
+    """
+    Apply post training activation quantization to specified model, using calibration data
+    """
+    submodules = dict(model.named_modules(remove_duplicate=True))
+    layer_norm_modules = [key for key, val in submodules.items() if isinstance(val, LayerNormANE)]
+    non_traceable_module_names = layer_norm_modules + [
+        "time_proj",
+        "time_embedding",
+    ]
+
+    # Mark certain modules as non-traceable to make the UNet model fx traceable
+    config.non_traceable_module_names = non_traceable_module_names
+    config.preserved_attributes = ['config', 'device']
+
+    sample_input = calibration_data[0]
+    quantizer = LinearQuantizer(model, config)
+    logger.info("Preparing model for quantization")
+    prepared_model = quantizer.prepare(example_inputs=(sample_input,))
+    prepared_model.eval()
+
+    quantizer.step()
+
+    logger.info("Calibrate")
+    for idx, data in enumerate(calibration_data):
+        logger.info(f"Calibration data sample: {idx}")
+        prepared_model(*data)
+
+    logger.info("Finalize model")
+    quantized_model = quantizer.finalize()
+    return quantized_model
+
+def get_quantizable_modules(unet):
+    quantizable_modules = []
+    for name, module in unet.named_modules():
+        if len(list(module.children())) > 0:
+            continue
+        if type(module) == torch.nn.modules.conv.Conv2d:
+            quantizable_modules.append(('conv', name))
+        if type(module) == Einsum:
+            quantizable_modules.append(('einsum', name))
+
+    return quantizable_modules
+
+def recipe_overrides_for_inference_speedup(conv_layers, skipped_conv):
+    """
+    Quantize the slowest conv layers, even if in skipped set based on PSNR, for good inference speedup
+    """
+    for layer in conv_layers:
+        if "up_blocks" in layer and "resnets" in layer and "conv1" in layer:
+            if layer in skipped_conv:
+                logger.info(f"removing {layer}")
+                skipped_conv.remove(layer)
+        if "upsamplers" in layer:
+            if layer in skipped_conv:
+                logger.info(f"removing {layer}")
+                skipped_conv.remove(layer)
+
+def recipe_overrides_for_quality(conv_layers, skipped_conv):
+    """
+    Do not quantize out projection layers to avoid quantizing outputs of preceding concat layers.
+    Quantizing output of concat layers can lead to quality degradation, due to sharing of scales
+    across concat inputs, which can have varied ranges. Since this is a constraint enforced during
+    model conversion, it may not be captured in layer-wise PSNR analysis of PyTorch model.
+    """
+    out_proj_layers = [layer for layer in conv_layers if "to_out" in layer]
+    for layer in out_proj_layers:
+        if layer not in skipped_conv:
+            logger.info(f"adding {layer}")
+            skipped_conv.add(layer)
+
+def register_input_log_hook(unet, inputs):
+    """
+    Register forward pre hook to save model inputs 
+    """
+    def hook(_, input):
+        input_copy = deepcopy(input)
+        input_copy = tuple(i.to('cpu') for i in input_copy)
+        inputs.append(input_copy)
+
+        # Return inputs unmodified
+        return input
+
+    return unet.register_forward_pre_hook(hook)
+
+def generate_calibration_data(pipe, args, calibration_dir):
+    # Register forward pre hook to record unet inputs
+    unet_inputs = []
+    handle = register_input_log_hook(pipe.unet, unet_inputs)
+
+    # If directory doesn't exist, create it
+    os.makedirs(calibration_dir, exist_ok=True)
+
+    # Run calibration prompts through the pipeline and
+    # serialize recorded UNet model inputs
+    for prompt in CALIBRATION_DATA:
+        gen = torch.manual_seed(args.seed)
+        # run forward pass
+        pipe(prompt=prompt, generator=gen)
+        # save unet inputs
+        filename = "_".join(prompt.split(" ")) + "_" + str(args.seed) + ".pkl"
+        filepath = os.path.join(calibration_dir, filename)
+        with open(filepath, 'wb') as f:
+            pickle.dump(unet_inputs, f)
+        # clear
+        unet_inputs.clear()
+
+    handle.remove()
+
+def register_input_preprocessing_hook(pipe):
+    """
+    Register forward pre hook to convert UNet inputs from HuggingFace StableDiffusionPipeline
+    to match expected model inputs in UNet2DConditionModel defined in unet.py
+    """
+    def hook(_, args, kwargs):
+        sample = args[0]
+        timestep = args[1]
+        if len(timestep.shape) == 0:
+            timestep = timestep[None]
+        timestep = timestep.expand(sample.shape[0])
+        encoder_hidden_states = kwargs["encoder_hidden_states"]
+        encoder_hidden_states = encoder_hidden_states.permute((0, 2, 1)).unsqueeze(2)
+        modified_args = (sample, timestep, encoder_hidden_states)
+        return (modified_args, {})
+
+    return pipe.unet.register_forward_pre_hook(hook, with_kwargs=True)
+
+def prepare_pipe(pipe, unet):
+    """
+    Create a new pipeline from `pipe` with `unet` as the noise predictor
+    """
+    new_pipe = deepcopy(pipe)
+    unet.to(new_pipe.unet.device)
+    new_pipe.unet = unet
+    pre_hook_handle = register_input_preprocessing_hook(new_pipe)
+    return new_pipe, pre_hook_handle
+
+def run_pipe(pipe):
+    gen = torch.manual_seed(args.seed)
+    kwargs = dict(
+        prompt=RANDOM_TEST_DATA,
+        output_type="latent",
+        generator=gen,
+    )
+    return np.array([latent.cpu().numpy() for latent in pipe(**kwargs).images])
+
+
+def get_reference_pipeline(model_version):
+    # Initialize pipe
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_version,
+        use_safetensors=True,
+        use_auth_token=True,
+    )
+    DEFAULT_NUM_INFERENCE_STEPS = 50
+    pipe.scheduler.set_timesteps(DEFAULT_NUM_INFERENCE_STEPS)
+
+    # Initialize reference unet
+    unet_cls = unet.UNet2DConditionModel
+    reference_unet = unet_cls(**pipe.unet.config).eval()
+    reference_unet.load_state_dict(pipe.unet.state_dict())
+
+    # Initialize reference pipeline
+    ref_pipe, _ = prepare_pipe(pipe, reference_unet)
+
+    del pipe
+    gc.collect()
+    return ref_pipe
+
+def main(args):
+    # Initialize reference pipeline
+    ref_pipe = get_reference_pipeline(args.model_version)
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    logger.debug(f"Placing pipe in {device}")
+    ref_pipe.to(device)
+    # Generate baseline outputs
+    ref_out = run_pipe(ref_pipe)
+
+    # Setup artifact file paths
+    os.makedirs(args.o, exist_ok=True)
+    recipe_json_path = os.path.join(args.o, f"{args.model_version.replace('/', '_')}_quantization_recipe.json")
+    calibration_dir = os.path.join(args.o, f"calibration_data_{args.model_version.replace('/', '_')}")
+
+    # Generate calibration data 
+    if args.generate_calibration_data:
+        generate_calibration_data(ref_pipe, args, calibration_dir)
+
+    # Compute layer-wise PSNR
+    if args.layerwise_sensitivity:
+        logger.info("Compute Layer-wise PSNR")
+        quantizable_modules = get_quantizable_modules(ref_pipe.unet)
+
+        results = {
+            'conv': {},
+            'einsum': {},
+            'model_version': args.model_version
+        }
+        dataloader = unet_data_loader(calibration_dir, device, args.calibration_nsamples)
+
+        for module_type, module_name in tqdm(quantizable_modules):
+            logger.info(f"Quantizing UNet Layer: {module_name}")
+            config = quantize_module_config(module_name)
+            quantized_unet = quantize(ref_pipe.unet, config, dataloader)
+
+            # Generate outputs from quantized model
+            q_pipe, _ = prepare_pipe(ref_pipe, quantized_unet)
+            test_out = run_pipe(q_pipe)
+
+            psnr = [float(f"{compute_psnr(r, t):.1f}") for r, t in zip(ref_out, test_out)]
+            logger.info(f"PSNR: {psnr}")
+            avg_psnr = sum(psnr) / len(psnr)
+            logger.info(f"AVG PSNR: {avg_psnr}")
+            results[module_type][module_name] = avg_psnr
+
+            del quantized_unet
+            del q_pipe
+            gc.collect()
+
+        with open(recipe_json_path, 'w') as f:
+            json.dump(results, f, indent=2)
+
+    if args.quantize_pytorch:
+        logger.info("Quantizing UNet PyTorch model")
+        dataloader = unet_data_loader(calibration_dir, device, args.calibration_nsamples)
+
+        with open(recipe_json_path, "r") as f:
+            results = json.load(f)
+
+        logger.info(f"Conv PSNR threshold: {args.conv_psnr}, Attn PSNR threshold: {args.attn_psnr}")
+        skipped_conv = set([layer for layer, psnr in results['conv'].items() if psnr < args.conv_psnr])
+        skipped_einsum = set([layer for layer, psnr in results['einsum'].items() if psnr < args.attn_psnr])
+
+        # Apply some overrides on PSNR based recipe for inference and quality improvements
+        # Users can disable these selectively based on specific targets
+        recipe_overrides_for_inference_speedup(results['conv'].keys(), skipped_conv)
+        recipe_overrides_for_quality(results['conv'].keys(), skipped_conv)
+
+        config = quantize_cumulative_config(skipped_conv, skipped_einsum)
+        quantized_unet = quantize(ref_pipe.unet, config, dataloader)
+
+        # Generate outputs from quantized model
+        q_pipe, handle = prepare_pipe(ref_pipe, quantized_unet)
+        test_out = run_pipe(q_pipe)
+
+        psnr = [float(f"{compute_psnr(r, t):.1f}") for r, t in zip(ref_out, test_out)]
+        logger.info(f"PSNR: {psnr}")
+        avg_psnr = sum(psnr) / len(psnr)
+        logger.info(f"AVG PSNR: {avg_psnr}")
+
+        handle.remove()
+        quantized_unet.to('cpu')
+        sample_unet_input = {
+            "sample": dataloader[0][0].to('cpu'),
+            "timestep": dataloader[0][1].to('cpu'),
+            "encoder_hidden_states": dataloader[0][2].to('cpu'),
+        }
+
+        logger.info("JIT tracing quantized model")
+        traced_model = torch.jit.trace(quantized_unet, example_inputs=list(sample_unet_input.values()))
+
+        logger.info("Converting to CoreML")
+        coreml_sample_unet_input = {
+            k: v.numpy().astype(np.float16)
+            for k, v in sample_unet_input.items()
+        }
+        coreml_model = convert_to_coreml(traced_model, coreml_sample_unet_input)
+        coreml_filename = f"Stable_Diffusion_version_{args.model_version.replace('/', '_')}_unet.mlpackage"
+        coreml_model.save(os.path.join(args.o, coreml_filename))
+
+        del q_pipe
+
+    del ref_pipe
+    gc.collect()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        required=True,
+        help="Output directory to save calibration data and quantization artifacts"
+    )
+    parser.add_argument(
+        "--model-version",
+        required=True,
+        choices=("runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1-base"),
+        help=
+        ("The pre-trained model checkpoint and configuration to restore"
+    ))
+    parser.add_argument(
+        "--generate-calibration-data",
+        action="store_true",
+        help="Generate calibration data for UNet model"
+    )
+    parser.add_argument(
+        "--layerwise-sensitivity",
+        action="store_true",
+        help="Compute compression sensitivity per-layer, by quantizing one layer at a time"
+    )
+    parser.add_argument(
+        "--quantize-pytorch",
+        action="store_true",
+        help="Generate activation quantized UNet model by quantizing layers above specified PSNR threshold"
+    )
+    parser.add_argument(
+        "--calibration-nsamples",
+        type=int,
+        help="Number of samples to use for calibrating UNet model"
+    )
+    parser.add_argument("--seed",
+                        "-s",
+                        default=50,
+                        type=int,
+                        help="Random seed to be able to reproduce results"
+    )
+    parser.add_argument("--conv-psnr",
+                        default=40.0,
+                        type=float,
+                        help="PSNR threshold for convolutional layers (default for stabilityai/stable-diffusion-2-1-base)"
+    )
+    parser.add_argument("--attn-psnr",
+                        default=30.0,
+                        type=float,
+                        help="PSNR threshold for attention (Einsum) layers (default for stabilityai/stable-diffusion-2-1-base)"
+    )
+    args = parser.parse_args()
+
+    main(args)

python_coreml_stable_diffusion/attention.py

@@ -0,0 +1,168 @@
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+import torch
+import math
+
+SPLIT_SOFTMAX = False
+
+def softmax(x, dim):
+    # Reduction max
+    max_x = x.max(dim=dim, keepdim=True).values
+    # EW sub
+    x -= max_x
+    # Scale for EXP to EXP2, Activation EXP2
+    scaled_x = x * (1 / math.log(2))
+    exp_act = torch.exp2(scaled_x)
+    # Reduction Sum + Inv
+    exp_sum_inv = 1 / exp_act.sum(dim=dim, keepdims=True)
+    # EW Mult
+    return exp_act * exp_sum_inv
+
+def split_einsum(q, k, v, mask, heads, dim_head):
+    """ Attention Implementation backing AttentionImplementations.SPLIT_EINSUM
+
+    - Implements https://machinelearning.apple.com/research/neural-engine-transformers
+    - Recommended for ANE
+    - Marginally slower on GPU
+    """
+    mh_q = [
+        q[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    k = k.transpose(1, 3)
+    mh_k = [
+        k[:, :, :,
+          head_idx * dim_head:(head_idx + 1) * dim_head]
+        for head_idx in range(heads)
+    ]  # (bs, max_seq_length, 1, dim_head) * heads
+
+    mh_v = [
+        v[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    attn_weights = [
+        torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * (dim_head**-0.5)
+        for qi, ki in zip(mh_q, mh_k)
+    ]  # (bs, max_seq_length, 1, max_seq_length) * heads
+
+    if mask is not None:
+        for head_idx in range(heads):
+            attn_weights[head_idx] = attn_weights[head_idx] + mask
+
+    if SPLIT_SOFTMAX:
+        attn_weights = [
+            softmax(aw, dim=1) for aw in attn_weights
+        ]  # (bs, max_seq_length, 1, max_seq_length) * heads
+    else:
+        attn_weights = [
+            aw.softmax(dim=1) for aw in attn_weights
+        ]  # (bs, max_seq_length, 1, max_seq_length) * heads
+
+    attn = [
+        torch.einsum("bkhq,bchk->bchq", wi, vi)
+        for wi, vi in zip(attn_weights, mh_v)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    attn = torch.cat(attn, dim=1)  # (bs, dim, 1, max_seq_length)
+    return attn
+
+
+CHUNK_SIZE = 512
+
+def split_einsum_v2(q, k, v, mask, heads, dim_head):
+    """ Attention Implementation backing AttentionImplementations.SPLIT_EINSUM_V2
+
+    - Implements https://machinelearning.apple.com/research/neural-engine-transformers
+    - Recommended for ANE
+    - Marginally slower on GPU
+    - Chunks the query sequence to avoid large intermediate tensors and improves ANE performance
+    """
+    query_seq_length = q.size(3)
+    num_chunks = query_seq_length // CHUNK_SIZE
+    
+    if num_chunks == 0:
+        logger.info(
+            "AttentionImplementations.SPLIT_EINSUM_V2: query sequence too short to chunk "
+            f"({query_seq_length}<{CHUNK_SIZE}), fall back to AttentionImplementations.SPLIT_EINSUM (safe to ignore)")
+        return split_einsum(q, k, v, mask, heads, dim_head)
+    
+    logger.info(
+        "AttentionImplementations.SPLIT_EINSUM_V2: Splitting query sequence length of "
+        f"{query_seq_length} into {num_chunks} chunks")
+
+    mh_q = [
+        q[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    # Chunk the query sequence for each head
+    mh_q_chunked = [
+        [h_q[..., chunk_idx * CHUNK_SIZE:(chunk_idx + 1) * CHUNK_SIZE] for chunk_idx in range(num_chunks)]
+        for h_q in mh_q
+    ]  # ((bs, dim_head, 1, QUERY_SEQ_CHUNK_SIZE) * num_chunks) * heads
+
+    k = k.transpose(1, 3)
+    mh_k = [
+        k[:, :, :,
+          head_idx * dim_head:(head_idx + 1) * dim_head]
+        for head_idx in range(heads)
+    ]  # (bs, max_seq_length, 1, dim_head) * heads
+
+    mh_v = [
+        v[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    attn_weights = [
+        [
+            torch.einsum("bchq,bkhc->bkhq", [qi_chunk, ki]) * (dim_head**-0.5)
+            for qi_chunk in h_q_chunked
+        ] for h_q_chunked, ki in zip(mh_q_chunked, mh_k)
+    ]  # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
+
+    attn_weights = [
+        [aw_chunk.softmax(dim=1) for aw_chunk in aw_chunked]
+        for aw_chunked in attn_weights
+    ]  # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
+
+    attn = [
+        [
+            torch.einsum("bkhq,bchk->bchq", wi_chunk, vi)
+            for wi_chunk in wi_chunked
+        ] for wi_chunked, vi in zip(attn_weights, mh_v)
+    ]  # ((bs, dim_head, 1, chunk_size) * num_chunks) * heads
+
+    attn = torch.cat([
+        torch.cat(attn_chunked, dim=3) for attn_chunked in attn
+    ], dim=1)  # (bs, dim, 1, max_seq_length)
+
+    return attn
+
+
+def original(q, k, v, mask, heads, dim_head):
+    """ Attention Implementation backing AttentionImplementations.ORIGINAL
+
+    - Not recommended for ANE
+    - Recommended for GPU
+    """
+    bs = q.size(0)
+    mh_q = q.view(bs, heads, dim_head, -1)
+    mh_k = k.view(bs, heads, dim_head, -1)
+    mh_v = v.view(bs, heads, dim_head, -1)
+
+    attn_weights = torch.einsum("bhcq,bhck->bhqk", [mh_q, mh_k])
+    attn_weights.mul_(dim_head**-0.5)
+
+    if mask is not None:
+        attn_weights = attn_weights + mask
+
+    attn_weights = attn_weights.softmax(dim=3)
+
+    attn = torch.einsum("bhqk,bhck->bhcq", [attn_weights, mh_v])
+    attn = attn.contiguous().view(bs, heads * dim_head, 1, -1)
+    return attn

python_coreml_stable_diffusion/chunk_mlprogram.py

@@ -0,0 +1,410 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+import argparse
+from collections import OrderedDict
+
+import coremltools as ct
+from coremltools.converters.mil import Block, Program, Var
+from coremltools.converters.mil.frontend.milproto.load import load as _milproto_to_pymil
+from coremltools.converters.mil.mil import Builder as mb
+from coremltools.converters.mil.mil import Placeholder
+from coremltools.converters.mil.mil import types as types
+from coremltools.converters.mil.mil.passes.helper import block_context_manager
+from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
+from coremltools.converters.mil.testing_utils import random_gen_input_feature_type
+
+import gc
+
+import logging
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+import numpy as np
+import os
+from python_coreml_stable_diffusion import torch2coreml
+import shutil
+import time
+
+
+def _verify_output_correctness_of_chunks(full_model,
+                                         first_chunk_model=None,
+                                         second_chunk_model=None,
+                                         pipeline_model=None,):
+    """ Verifies the end-to-end output correctness of full (original) model versus chunked models
+    """
+    # Generate inputs for first chunk and full model
+    input_dict = {}
+    for input_desc in full_model._spec.description.input:
+        input_dict[input_desc.name] = random_gen_input_feature_type(input_desc)
+
+    # Generate outputs for full model
+    outputs_from_full_model = full_model.predict(input_dict)
+
+    if pipeline_model is not None:
+        outputs_from_pipeline_model = pipeline_model.predict(input_dict)
+        final_outputs = outputs_from_pipeline_model
+
+    elif first_chunk_model is not None and second_chunk_model is not None:
+        # Generate outputs for first chunk
+        outputs_from_first_chunk_model = first_chunk_model.predict(input_dict)
+
+        # Prepare inputs for second chunk model from first chunk's outputs and regular inputs
+        second_chunk_input_dict = {}
+        for input_desc in second_chunk_model._spec.description.input:
+            if input_desc.name in outputs_from_first_chunk_model:
+                second_chunk_input_dict[
+                    input_desc.name] = outputs_from_first_chunk_model[
+                        input_desc.name]
+            else:
+                second_chunk_input_dict[input_desc.name] = input_dict[
+                    input_desc.name]
+
+        # Generate output for second chunk model
+        outputs_from_second_chunk_model = second_chunk_model.predict(
+            second_chunk_input_dict)
+        final_outputs = outputs_from_second_chunk_model
+    else:
+        raise ValueError
+
+    # Verify correctness across all outputs from second chunk and full model
+    for out_name in outputs_from_full_model.keys():
+        torch2coreml.report_correctness(
+            original_outputs=outputs_from_full_model[out_name],
+            final_outputs=final_outputs[out_name],
+            log_prefix=f"{out_name}")
+
+
+def _load_prog_from_mlmodel(model):
+    """ Load MIL Program from an MLModel
+    """
+    model_spec = model.get_spec()
+    start_ = time.time()
+    logger.info(
+        "Loading MLModel object into a MIL Program object (including the weights).."
+    )
+    prog = _milproto_to_pymil(
+        model_spec=model_spec,
+        specification_version=model_spec.specificationVersion,
+        file_weights_dir=model.weights_dir,
+    )
+    logger.info(f"Program loaded in {time.time() - start_:.1f} seconds")
+
+    return prog
+
+
+def _get_op_idx_split_location(prog: Program):
+    """ Find the op that approximately bisects the graph as measure by weights size on each side
+    """
+    main_block = prog.functions["main"]
+    main_block.operations = list(main_block.operations)
+    total_size_in_mb = 0
+
+    for op in main_block.operations:
+        if op.op_type == "const" and isinstance(op.val.val, np.ndarray):
+            size_in_mb = op.val.val.size * op.val.val.itemsize / (1024 * 1024)
+            total_size_in_mb += size_in_mb
+    half_size = total_size_in_mb / 2
+
+    # Find the first non const op (single child), where the total cumulative size exceeds
+    # the half size for the first time
+    cumulative_size_in_mb = 0
+    for op in main_block.operations:
+        if op.op_type == "const" and isinstance(op.val.val, np.ndarray):
+            size_in_mb = op.val.val.size * op.val.val.itemsize / (1024 * 1024)
+            cumulative_size_in_mb += size_in_mb
+
+        # Note: The condition "not op.op_type.startswith("const")" is to make sure that the
+        # incision op is neither of type "const" nor "constexpr_*" ops that
+        # are used to store compressed weights
+        if (cumulative_size_in_mb > half_size and not op.op_type.startswith("const")
+                and len(op.outputs) == 1
+                and len(op.outputs[0].child_ops) == 1):
+            op_idx = main_block.operations.index(op)
+            return op_idx, cumulative_size_in_mb, total_size_in_mb
+
+
+def _get_first_chunk_outputs(block, op_idx):
+    # Get the list of all vars that go across from first program (all ops from 0 to op_idx (inclusive))
+    # to the second program (all ops from op_idx+1 till the end). These all vars need to be made the output
+    # of the first program and the input of the second program
+    boundary_vars = set()
+    block.operations = list(block.operations)
+    for i in range(op_idx + 1):
+        op = block.operations[i]
+        if not op.op_type.startswith("const"):
+            for var in op.outputs:
+                if var.val is None:  # only consider non const vars
+                    for child_op in var.child_ops:
+                        child_op_idx = block.operations.index(child_op)
+                        if child_op_idx > op_idx:
+                            boundary_vars.add(var)
+    return list(boundary_vars)
+
+
+@block_context_manager
+def _add_fp32_casts(block, boundary_vars):
+    new_boundary_vars = []
+    for var in boundary_vars:
+        if var.dtype != types.fp16:
+            new_boundary_vars.append(var)
+        else:
+            fp32_var = mb.cast(x=var, dtype="fp32", name=var.name)
+            new_boundary_vars.append(fp32_var)
+    return new_boundary_vars
+
+
+def _make_first_chunk_prog(prog, op_idx):
+    """ Build first chunk by declaring early outputs and removing unused subgraph
+    """
+    block = prog.functions["main"]
+    boundary_vars = _get_first_chunk_outputs(block, op_idx)
+
+    # Due to possible numerical issues, cast any fp16 var to fp32
+    new_boundary_vars = _add_fp32_casts(block, boundary_vars)
+
+    block.outputs.clear()
+    block.set_outputs(new_boundary_vars)
+    PASS_REGISTRY["common::dead_code_elimination"](prog)
+    return prog
+
+
+def _make_second_chunk_prog(prog, op_idx):
+    """ Build second chunk by rebuilding a pristine MIL Program from MLModel
+    """
+    block = prog.functions["main"]
+    block.opset_version = ct.target.iOS16
+
+    # First chunk outputs are second chunk inputs (e.g. skip connections)
+    boundary_vars = _get_first_chunk_outputs(block, op_idx)
+
+    # This op will not be included in this program. Its output var will be made into an input
+    block.operations = list(block.operations)
+    boundary_op = block.operations[op_idx]
+
+    # Add all boundary ops as inputs
+    with block:
+        for var in boundary_vars:
+            new_placeholder = Placeholder(
+                sym_shape=var.shape,
+                dtype=var.dtype if var.dtype != types.fp16 else types.fp32,
+                name=var.name,
+            )
+
+            block._input_dict[
+                new_placeholder.outputs[0].name] = new_placeholder.outputs[0]
+
+            block.function_inputs = tuple(block._input_dict.values())
+            new_var = None
+            if var.dtype == types.fp16:
+                new_var = mb.cast(x=new_placeholder.outputs[0],
+                                  dtype="fp16",
+                                  before_op=var.op)
+            else:
+                new_var = new_placeholder.outputs[0]
+
+            block.replace_uses_of_var_after_op(
+                anchor_op=boundary_op,
+                old_var=var,
+                new_var=new_var,
+                # This is needed if the program contains "constexpr_*" ops. In normal cases, there are stricter
+                # rules for removing them, and their presence may prevent replacing this var.
+                # However in this case, since we want to remove all the ops in chunk 1, we can safely
+                # set this to True.
+                force_replace=True,
+            )
+
+    PASS_REGISTRY["common::dead_code_elimination"](prog)
+
+    # Remove any unused inputs
+    new_input_dict = OrderedDict()
+    for k, v in block._input_dict.items():
+        if len(v.child_ops) > 0:
+            new_input_dict[k] = v
+    block._input_dict = new_input_dict
+    block.function_inputs = tuple(block._input_dict.values())
+
+    return prog
+
+
+def _legacy_model_chunking(args):
+    # TODO: Remove this method after setting the coremltools dependency >= 8.0
+    os.makedirs(args.o, exist_ok=True)
+
+    # Check filename extension
+    mlpackage_name = os.path.basename(args.mlpackage_path)
+    name, ext = os.path.splitext(mlpackage_name)
+    assert ext == ".mlpackage", f"`--mlpackage-path` (args.mlpackage_path) is not an .mlpackage file"
+
+    # Load CoreML model
+    logger.info("Loading model from {}".format(args.mlpackage_path))
+    start_ = time.time()
+    model = ct.models.MLModel(
+        args.mlpackage_path,
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+    )
+    logger.info(
+        f"Loading {args.mlpackage_path} took {time.time() - start_:.1f} seconds"
+    )
+
+    # Load the MIL Program from MLModel
+    prog = _load_prog_from_mlmodel(model)
+
+    # Compute the incision point by bisecting the program based on weights size
+    op_idx, first_chunk_weights_size, total_weights_size = _get_op_idx_split_location(
+        prog)
+    main_block = prog.functions["main"]
+    incision_op = main_block.operations[op_idx]
+    logger.info(f"{args.mlpackage_path} will chunked into two pieces.")
+    logger.info(
+        f"The incision op: name={incision_op.name}, type={incision_op.op_type}, index={op_idx}/{len(main_block.operations)}"
+    )
+    logger.info(f"First  chunk size = {first_chunk_weights_size:.2f} MB")
+    logger.info(
+        f"Second chunk size = {total_weights_size - first_chunk_weights_size:.2f} MB"
+    )
+
+    # Build first chunk (in-place modifies prog by declaring early exits and removing unused subgraph)
+    prog_chunk1 = _make_first_chunk_prog(prog, op_idx)
+
+    # Build the second chunk
+    prog_chunk2 = _make_second_chunk_prog(_load_prog_from_mlmodel(model),
+                                          op_idx)
+
+    if not args.check_output_correctness:
+        # Original model no longer needed in memory
+        del model
+        gc.collect()
+
+    # Convert the MIL Program objects into MLModels
+    logger.info("Converting the two programs")
+    model_chunk1 = ct.convert(
+        prog_chunk1,
+        convert_to="mlprogram",
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+        minimum_deployment_target=ct.target.iOS16,
+    )
+    del prog_chunk1
+    gc.collect()
+    logger.info("Conversion of first chunk done.")
+
+    model_chunk2 = ct.convert(
+        prog_chunk2,
+        convert_to="mlprogram",
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+        minimum_deployment_target=ct.target.iOS16,
+    )
+    del prog_chunk2
+    gc.collect()
+    logger.info("Conversion of second chunk done.")
+
+    # Verify output correctness
+    if args.check_output_correctness:
+        logger.info("Verifying output correctness of chunks")
+        _verify_output_correctness_of_chunks(
+            full_model=model,
+            first_chunk_model=model_chunk1,
+            second_chunk_model=model_chunk2,
+        )
+
+    if args.merge_chunks_in_pipeline_model:
+        # Make a single pipeline model to manage the model chunks
+        pipeline_model = ct.utils.make_pipeline(model_chunk1, model_chunk2)
+        out_path_pipeline = os.path.join(args.o, name + "_chunked_pipeline.mlpackage")
+
+        # Save and reload to ensure CPU placement
+        pipeline_model.save(out_path_pipeline)
+        pipeline_model = ct.models.MLModel(out_path_pipeline, compute_units=ct.ComputeUnit.CPU_ONLY)
+
+        if args.check_output_correctness:
+            logger.info("Verifying output correctness of pipeline model")
+            _verify_output_correctness_of_chunks(
+                full_model=model,
+                pipeline_model=pipeline_model,
+            )
+    else:
+        # Save the chunked models to disk
+        out_path_chunk1 = os.path.join(args.o, name + "_chunk1.mlpackage")
+        out_path_chunk2 = os.path.join(args.o, name + "_chunk2.mlpackage")
+
+        logger.info(
+            f"Saved chunks in {args.o} with the suffix _chunk1.mlpackage and _chunk2.mlpackage"
+        )
+        model_chunk1.save(out_path_chunk1)
+        model_chunk2.save(out_path_chunk2)
+        logger.info("Done.")
+
+
+def main(args):
+    ct_version = ct.__version__
+
+    if ct_version != "8.0b2" and ct_version < "8.0":
+        # With coremltools version <= 8.0b1,
+        # we use the legacy implementation.
+        # TODO: Remove the logic after setting the coremltools dependency >= 8.0.
+        logger.info(
+            f"coremltools version {ct_version} detected. Recommended upgrading the package version to "
+            f"'8.0b2' when you running chunk_mlprogram.py script for the latest supports and bug fixes."
+        )
+        _legacy_model_chunking(args)
+    else:
+        # Starting from coremltools==8.0b2, there is this `bisect_model` API that
+        # we can directly call into.
+        from coremltools.models.utils import bisect_model
+        logger.info(f"Start chunking model {args.mlpackage_path} into two pieces.")
+        ct.models.utils.bisect_model(
+            model=args.mlpackage_path,
+            output_dir=args.o,
+            merge_chunks_to_pipeline=args.merge_chunks_in_pipeline_model,
+            check_output_correctness=args.check_output_correctness,
+        )
+        logger.info(f"Model chunking is done.")
+
+    # Remove original (non-chunked) model if requested
+    if args.remove_original:
+        logger.info(
+            "Removing original (non-chunked) model at {args.mlpackage_path}")
+        shutil.rmtree(args.mlpackage_path)
+        logger.info("Done.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--mlpackage-path",
+        required=True,
+        help=
+        "Path to the mlpackage file to be split into two mlpackages of approximately same file size.",
+    )
+    parser.add_argument(
+        "-o",
+        required=True,
+        help=
+        "Path to output directory where the two model chunks should be saved.",
+    )
+    parser.add_argument(
+        "--remove-original",
+        action="store_true",
+        help=
+        "If specified, removes the original (non-chunked) model to avoid duplicating storage."
+    )
+    parser.add_argument(
+        "--check-output-correctness",
+        action="store_true",
+        help=
+        ("If specified, compares the outputs of original Core ML model with that of pipelined CoreML model chunks and reports PSNR in dB. ",
+         "Enabling this feature uses more memory. Disable it if your machine runs out of memory."
+         ))
+    parser.add_argument(
+        "--merge-chunks-in-pipeline-model",
+        action="store_true",
+        help=
+        ("If specified, model chunks are managed inside a single pipeline model for easier asset maintenance"
+         ))
+
+    args = parser.parse_args()
+    main(args)

python_coreml_stable_diffusion/controlnet.py

@@ -0,0 +1,250 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers import ModelMixin
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .unet import Timesteps, TimestepEmbedding, get_down_block, UNetMidBlock2DCrossAttn, linear_to_conv2d_map
+
+class ControlNetConditioningEmbedding(nn.Module):
+
+    def __init__(
+        self,
+        conditioning_embedding_channels,
+        conditioning_channels=3,
+        block_out_channels=(16, 32, 96, 256),
+    ):
+        super().__init__()
+
+        self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)
+
+        self.blocks = nn.ModuleList([])
+
+        for i in range(len(block_out_channels) - 1):
+            channel_in = block_out_channels[i]
+            channel_out = block_out_channels[i + 1]
+            self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))
+            self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))
+
+        self.conv_out = nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
+
+    def forward(self, conditioning):
+        embedding = self.conv_in(conditioning)
+        embedding = F.silu(embedding)
+
+        for block in self.blocks:
+            embedding = block(embedding)
+            embedding = F.silu(embedding)
+
+        embedding = self.conv_out(embedding)
+
+        return embedding
+
+class ControlNetModel(ModelMixin, ConfigMixin):
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels=4,
+        flip_sin_to_cos=True,
+        freq_shift=0,
+        down_block_types=(
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "DownBlock2D",
+        ),
+        only_cross_attention=False,
+        block_out_channels=(320, 640, 1280, 1280),
+        layers_per_block=2,
+        downsample_padding=1,
+        mid_block_scale_factor=1,
+        act_fn="silu",
+        norm_num_groups=32,
+        norm_eps=1e-5,
+        cross_attention_dim=1280,
+        transformer_layers_per_block=1,
+        attention_head_dim=8,
+        use_linear_projection=False,
+        upcast_attention=False,
+        resnet_time_scale_shift="default",
+        conditioning_embedding_out_channels=(16, 32, 96, 256),
+        **kwargs,
+    ):
+        super().__init__()
+
+        # Check inputs
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        self._register_load_state_dict_pre_hook(linear_to_conv2d_map)
+
+        # input
+        conv_in_kernel = 3
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = nn.Conv2d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+
+        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+        )
+
+        # control net conditioning embedding
+        self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
+            conditioning_embedding_channels=block_out_channels[0],
+            block_out_channels=conditioning_embedding_out_channels,
+        )
+
+        self.down_blocks = nn.ModuleList([])
+        self.controlnet_down_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        # down
+        output_channel = block_out_channels[0]
+
+        controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+        self.controlnet_down_blocks.append(controlnet_block)
+
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim[i],
+                downsample_padding=downsample_padding,
+                add_downsample=not is_final_block,
+            )
+            self.down_blocks.append(down_block)
+
+            for _ in range(layers_per_block):
+                controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                self.controlnet_down_blocks.append(controlnet_block)
+
+            if not is_final_block:
+                controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                self.controlnet_down_blocks.append(controlnet_block)
+
+        # mid
+        mid_block_channel = block_out_channels[-1]
+
+        controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
+        self.controlnet_mid_block = controlnet_block
+
+        self.mid_block = UNetMidBlock2DCrossAttn(
+            in_channels=mid_block_channel,
+            temb_channels=time_embed_dim,
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            output_scale_factor=mid_block_scale_factor,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attention_head_dim[-1],
+            resnet_groups=norm_num_groups,
+            use_linear_projection=use_linear_projection,
+            upcast_attention=upcast_attention,
+        )
+
+    def get_num_residuals(self):
+        num_res = 2 # initial sample + mid block
+        for down_block in self.down_blocks:
+            num_res += len(down_block.resnets)
+            if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None:
+                num_res += len(down_block.downsamplers)
+        return num_res
+
+    def forward(
+        self,
+        sample,
+        timestep,
+        encoder_hidden_states,
+        controlnet_cond,
+    ):
+        # 1. time
+        t_emb = self.time_proj(timestep)
+        emb = self.time_embedding(t_emb)
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
+
+        sample += controlnet_cond
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "attentions") and downsample_block.attentions is not None:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+
+        # 5. Control net blocks
+        controlnet_down_block_res_samples = ()
+
+        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
+            down_block_res_sample = controlnet_block(down_block_res_sample)
+            controlnet_down_block_res_samples += (down_block_res_sample,)
+
+        down_block_res_samples = controlnet_down_block_res_samples
+
+        mid_block_res_sample = self.controlnet_mid_block(sample)
+
+        return down_block_res_samples, mid_block_res_sample

python_coreml_stable_diffusion/coreml_model.py

@@ -0,0 +1,225 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+import coremltools as ct
+
+import logging
+import json
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+import numpy as np
+
+import os
+import time
+import subprocess
+import sys
+
+
+def _macos_version():
+    """
+    Returns macOS version as a tuple of integers. On non-Macs, returns an empty tuple.
+    """
+    if sys.platform == "darwin":
+        try:
+            ver_str = subprocess.run(["sw_vers", "-productVersion"], stdout=subprocess.PIPE).stdout.decode('utf-8').strip('\n')
+            return tuple([int(v) for v in ver_str.split(".")])
+        except:
+            raise Exception("Unable to determine the macOS version")
+    return ()
+
+
+class CoreMLModel:
+    """ Wrapper for running CoreML models using coremltools
+    """
+
+    def __init__(self, model_path, compute_unit, sources='packages', optimization_hints=None):
+
+        logger.info(f"Loading {model_path}")
+
+        start = time.time()
+        if sources == 'packages':
+            assert os.path.exists(model_path) and model_path.endswith(".mlpackage")
+
+            self.model = ct.models.MLModel(
+                model_path,
+                compute_units=ct.ComputeUnit[compute_unit],
+                optimization_hints=optimization_hints,
+            )
+            DTYPE_MAP = {
+                65552: np.float16,
+                65568: np.float32,
+                131104: np.int32,
+            }
+            self.expected_inputs = {
+                input_tensor.name: {
+                    "shape": tuple(input_tensor.type.multiArrayType.shape),
+                    "dtype": DTYPE_MAP[input_tensor.type.multiArrayType.dataType],
+                }
+                for input_tensor in self.model._spec.description.input
+            }
+        elif sources == 'compiled':
+            assert os.path.exists(model_path) and model_path.endswith(".mlmodelc")
+
+            self.model = ct.models.CompiledMLModel(
+                model_path,
+                compute_units=ct.ComputeUnit[compute_unit],
+                optimization_hints=optimization_hints,
+            )
+
+            # Grab expected inputs from metadata.json
+            with open(os.path.join(model_path, 'metadata.json'), 'r') as f:
+                config = json.load(f)[0]
+
+            self.expected_inputs = {
+                input_tensor['name']: {
+                    "shape": tuple(eval(input_tensor['shape'])),
+                    "dtype": np.dtype(input_tensor['dataType'].lower()),
+                }
+                for input_tensor in config['inputSchema']
+            }
+        else:
+            raise ValueError(f'Expected `packages` or `compiled` for sources, received {sources}')
+
+        load_time = time.time() - start
+        logger.info(f"Done. Took {load_time:.1f} seconds.")
+
+        if load_time > LOAD_TIME_INFO_MSG_TRIGGER:
+            logger.info(
+                "Loading a CoreML model through coremltools triggers compilation every time. "
+                "The Swift package we provide uses precompiled Core ML models (.mlmodelc) to avoid compile-on-load."
+            )
+
+    def _verify_inputs(self, **kwargs):
+        for k, v in kwargs.items():
+            if k in self.expected_inputs:
+                if not isinstance(v, np.ndarray):
+                    raise TypeError(
+                        f"Expected numpy.ndarray, got {v} for input: {k}")
+
+                expected_dtype = self.expected_inputs[k]["dtype"]
+                if not v.dtype == expected_dtype:
+                    raise TypeError(
+                        f"Expected dtype {expected_dtype}, got {v.dtype} for input: {k}"
+                    )
+
+                expected_shape = self.expected_inputs[k]["shape"]
+                if not v.shape == expected_shape:
+                    raise TypeError(
+                        f"Expected shape {expected_shape}, got {v.shape} for input: {k}"
+                    )
+            else:
+                raise ValueError(f"Received unexpected input kwarg: {k}")
+
+    def __call__(self, **kwargs):
+        self._verify_inputs(**kwargs)
+        return self.model.predict(kwargs)
+
+
+LOAD_TIME_INFO_MSG_TRIGGER = 10  # seconds
+
+
+def get_resource_type(resources_dir: str) -> str:
+    """
+        Detect resource type based on filepath extensions.
+        returns:
+            `packages`: for .mlpackage resources
+            'compiled`: for .mlmodelc resources
+    """
+    directories = [f for f in os.listdir(resources_dir) if os.path.isdir(os.path.join(resources_dir, f))]
+
+    # consider directories ending with extension
+    extensions = set([os.path.splitext(e)[1] for e in directories if os.path.splitext(e)[1]])
+
+    # if one extension present we may be able to infer sources type
+    if len(set(extensions)) == 1:
+        extension = extensions.pop()
+    else:
+        raise ValueError(f'Multiple file extensions found at {resources_dir}.'
+                         f'Cannot infer resource type from contents.')
+
+    if extension == '.mlpackage':
+        sources = 'packages'
+    elif extension == '.mlmodelc':
+        sources = 'compiled'
+    else:
+        raise ValueError(f'Did not find .mlpackage or .mlmodelc at {resources_dir}')
+
+    return sources
+
+
+def _load_mlpackage(submodule_name,
+                    mlpackages_dir,
+                    model_version,
+                    compute_unit,
+                    sources=None):
+    """
+        Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
+
+    """
+
+    # if sources not provided, attempt to infer `packages` or `compiled` from the
+    # resources directory
+    if sources is None:
+        sources = get_resource_type(mlpackages_dir)
+
+    if sources == 'packages':
+        logger.info(f"Loading {submodule_name} mlpackage")
+        fname = f"Stable_Diffusion_version_{model_version}_{submodule_name}.mlpackage".replace(
+            "/", "_")
+        mlpackage_path = os.path.join(mlpackages_dir, fname)
+
+        if not os.path.exists(mlpackage_path):
+            raise FileNotFoundError(
+                f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
+
+    elif sources == 'compiled':
+        logger.info(f"Loading {submodule_name} mlmodelc")
+
+        # FixMe: Submodule names and compiled resources names differ. Can change if names match in the future.
+        submodule_names = ["text_encoder", "text_encoder_2", "unet", "vae_decoder", "vae_encoder", "safety_checker"]
+        compiled_names = ['TextEncoder', 'TextEncoder2', 'Unet', 'VAEDecoder', 'VAEEncoder', 'SafetyChecker']
+        name_map = dict(zip(submodule_names, compiled_names))
+
+        cname = name_map[submodule_name] + '.mlmodelc'
+        mlpackage_path = os.path.join(mlpackages_dir, cname)
+
+        if not os.path.exists(mlpackage_path):
+            raise FileNotFoundError(
+                f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
+
+    # On macOS 15+, set fast prediction optimization hint for the unet.
+    optimization_hints = None
+    if submodule_name == "unet" and _macos_version() >= (15, 0):
+        optimization_hints = {"specializationStrategy": ct.SpecializationStrategy.FastPrediction}
+
+    return CoreMLModel(mlpackage_path,
+                       compute_unit,
+                       sources=sources,
+                       optimization_hints=optimization_hints)
+
+
+def _load_mlpackage_controlnet(mlpackages_dir, model_version, compute_unit):
+    """ Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
+    """
+    model_name = model_version.replace("/", "_")
+
+    logger.info(f"Loading controlnet_{model_name} mlpackage")
+
+    fname = f"ControlNet_{model_name}.mlpackage"
+
+    mlpackage_path = os.path.join(mlpackages_dir, fname)
+
+    if not os.path.exists(mlpackage_path):
+        raise FileNotFoundError(
+            f"controlnet_{model_name} CoreML model doesn't exist at {mlpackage_path}")
+
+    return CoreMLModel(mlpackage_path, compute_unit)
+
+
+def get_available_compute_units():
+    return tuple(cu for cu in ct.ComputeUnit._member_names_)

python_coreml_stable_diffusion/layer_norm.py

@@ -0,0 +1,80 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+import torch
+import torch.nn as nn
+
+
+# Reference: https://github.com/apple/ml-ane-transformers/blob/main/ane_transformers/reference/layer_norm.py
+class LayerNormANE(nn.Module):
+    """ LayerNorm optimized for Apple Neural Engine (ANE) execution
+
+    Note: This layer only supports normalization over the final dim. It expects `num_channels`
+    as an argument and not `normalized_shape` which is used by `torch.nn.LayerNorm`.
+    """
+
+    def __init__(self,
+                 num_channels,
+                 clip_mag=None,
+                 eps=1e-5,
+                 elementwise_affine=True):
+        """
+        Args:
+            num_channels:       Number of channels (C) where the expected input data format is BC1S. S stands for sequence length.
+            clip_mag:           Optional float value to use for clamping the input range before layer norm is applied.
+                                If specified, helps reduce risk of overflow.
+            eps:                Small value to avoid dividing by zero
+            elementwise_affine: If true, adds learnable channel-wise shift (bias) and scale (weight) parameters
+        """
+        super().__init__()
+        # Principle 1: Picking the Right Data Format (machinelearning.apple.com/research/apple-neural-engine)
+        self.expected_rank = len("BC1S")
+
+        self.num_channels = num_channels
+        self.eps = eps
+        self.clip_mag = clip_mag
+        self.elementwise_affine = elementwise_affine
+
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.Tensor(num_channels))
+            self.bias = nn.Parameter(torch.Tensor(num_channels))
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        if self.elementwise_affine:
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
+
+    def forward(self, inputs):
+        input_rank = len(inputs.size())
+
+        # Principle 1: Picking the Right Data Format (machinelearning.apple.com/research/apple-neural-engine)
+        # Migrate the data format from BSC to BC1S (most conducive to ANE)
+        if input_rank == 3 and inputs.size(2) == self.num_channels:
+            inputs = inputs.transpose(1, 2).unsqueeze(2)
+            input_rank = len(inputs.size())
+
+        assert input_rank == self.expected_rank
+        assert inputs.size(1) == self.num_channels
+
+        if self.clip_mag is not None:
+            inputs.clamp_(-self.clip_mag, self.clip_mag)
+
+        channels_mean = inputs.mean(dim=1, keepdims=True)
+
+        zero_mean = inputs - channels_mean
+
+        zero_mean_sq = zero_mean * zero_mean
+
+        denom = (zero_mean_sq.mean(dim=1, keepdims=True) + self.eps).rsqrt()
+
+        out = zero_mean * denom
+
+        if self.elementwise_affine:
+            out = (out + self.bias.view(1, self.num_channels, 1, 1)
+                   ) * self.weight.view(1, self.num_channels, 1, 1)
+
+        return out

python_coreml_stable_diffusion/mixed_bit_compression_apply.py

@@ -0,0 +1,133 @@
+import argparse
+import gc
+import json
+import logging
+import os
+
+import coremltools as ct
+import coremltools.optimize.coreml as cto
+import numpy as np
+
+from python_coreml_stable_diffusion.torch2coreml import get_pipeline
+from python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis import (
+    NBITS,
+    PALETTIZE_MIN_SIZE as MIN_SIZE
+)
+
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def main(args):
+    # Load Core ML model
+    coreml_model = ct.models.MLModel(args.mlpackage_path, compute_units=ct.ComputeUnit.CPU_ONLY)
+    logger.info(f"Loaded {args.mlpackage_path}")
+
+    # Load palettization recipe
+    with open(args.pre_analysis_json_path, 'r') as f:
+        pre_analysis = json.load(f)
+
+    if args.selected_recipe not in list(pre_analysis["recipes"]):
+        raise KeyError(
+            f"--selected-recipe ({args.selected_recipe}) not found in "
+            f"--pre-analysis-json-path ({args.pre_analysis_json_path}). "
+            f" Available recipes: {list(pre_analysis['recipes'])}"
+        )
+
+
+    recipe = pre_analysis["recipes"][args.selected_recipe]
+    assert all(nbits in NBITS + [16] for nbits in recipe.values()), \
+        f"Some nbits values in the recipe are illegal. Allowed values: {NBITS}"
+
+    # Hash tensors to be able to match torch tensor names to mil tensors
+    def get_tensor_hash(tensor):
+        assert tensor.dtype == np.float16
+        return tensor.ravel()[0] + np.prod(tensor.shape)
+
+    args.model_version = pre_analysis["model_version"]
+    pipe = get_pipeline(args)
+    torch_model = pipe.unet
+
+    hashed_recipe = {}
+    for torch_module_name, nbits in recipe.items():
+        tensor = [
+            tensor.cpu().numpy().astype(np.float16) for name,tensor in torch_model.named_parameters()
+            if name == torch_module_name + '.weight'
+        ][0]
+        hashed_recipe[get_tensor_hash(tensor)] = nbits
+
+    del pipe
+    gc.collect()
+
+    op_name_configs = {}
+    weight_metadata = cto.get_weights_metadata(coreml_model, weight_threshold=MIN_SIZE)
+    hashes = np.array(list(hashed_recipe))
+    for name, metadata in weight_metadata.items():
+        # Look up target bits for this weight
+        tensor_hash = get_tensor_hash(metadata.val)
+        pdist = np.abs(hashes - tensor_hash)
+        assert(pdist.min() < 0.01)
+        matched = pdist.argmin()
+        target_nbits = hashed_recipe[hashes[matched]]
+
+        if target_nbits == 16:
+            continue
+
+        op_name_configs[name] = cto.OpPalettizerConfig(
+            mode="kmeans",
+            nbits=target_nbits,
+            weight_threshold=int(MIN_SIZE)
+        )
+
+    config = ct.optimize.coreml.OptimizationConfig(op_name_configs=op_name_configs)
+    coreml_model = ct.optimize.coreml.palettize_weights(coreml_model, config)
+
+    coreml_model.save(args.o)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        required=True,
+        help="Output directory to save the custom palettized model"
+    )
+    parser.add_argument(
+        "--mlpackage-path",
+        required=True,
+        help="Path to .mlpackage model to be palettized"
+    )
+    parser.add_argument(
+        "--pre-analysis-json-path",
+        required=True,
+        type=str,
+        help=("The JSON file generated by mixed_bit_compression_pre_analysis.py"
+    ))
+    parser.add_argument(
+        "--selected-recipe",
+        required=True,
+        type=str,
+        help=("The string key into --pre-analysis-json-path's baselines dict"
+    ))
+    parser.add_argument(
+        "--custom-vae-version",
+        type=str,
+        default=None,
+        help=
+        ("Custom VAE checkpoint to override the pipeline's built-in VAE. "
+            "If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
+            "No precision override is applied when using a custom VAE."
+    ))
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.mlpackage_path):
+        raise FileNotFoundError
+    if not os.path.exists(args.pre_analysis_json_path):
+        raise FileNotFoundError
+    if not args.pre_analysis_json_path.endswith('.json'):
+        raise ValueError("--recipe-json-path should end with '.json'")
+
+    main(args)

python_coreml_stable_diffusion/mixed_bit_compression_pre_analysis.py

@@ -0,0 +1,583 @@
+from collections import OrderedDict
+from copy import deepcopy
+from functools import partial
+import argparse
+import gc
+import json
+
+import logging
+logging.basicConfig()
+logger = logging.getLogger()
+logger.setLevel('INFO')
+
+import numpy as np
+import os
+from PIL import Image
+from python_coreml_stable_diffusion.torch2coreml import compute_psnr, get_pipeline
+import time
+
+import torch
+import torch.nn as nn
+import requests
+torch.set_grad_enabled(False)
+
+from tqdm import tqdm
+
+# Bit-widths the Neural Engine is capable of accelerating
+NBITS = [1, 2, 4, 6, 8]
+
+# Minimum number of elements in a weight tensor to be considered for palettization
+# (saves pre-analysis time)
+PALETTIZE_MIN_SIZE = 1e5
+
+# Signal integrity is computed based on these 4 random prompts
+RANDOM_TEST_DATA = [
+    "a black and brown dog standing outside a door.",
+    "a person on a motorcycle makes a turn on the track.",
+    "inflatable boats sit on the arizona river, and on the bank",
+    "a white cat sitting under a white umbrella",
+    "black bear standing in a field of grass under a tree.",
+    "a train that is parked on tracks and has graffiti writing on it, with a mountain range in the background.",
+    "a cake inside of a pan sitting in an oven.",
+    "a table with paper plates and flowers in a home",
+]
+
+TEST_RESOLUTION = 768
+
+RANDOM_TEST_IMAGE_DATA = [
+    Image.open(
+        requests.get(path, stream=True).raw).convert("RGB").resize(
+            (TEST_RESOLUTION, TEST_RESOLUTION), Image.LANCZOS
+    ) for path in [
+        "http://farm1.staticflickr.com/106/298138827_19bb723252_z.jpg",
+        "http://farm4.staticflickr.com/3772/9666116202_648cd752d6_z.jpg",
+        "http://farm3.staticflickr.com/2238/2472574092_f5534bb2f7_z.jpg",
+        "http://farm1.staticflickr.com/220/475442674_47d81fdc2c_z.jpg",
+        "http://farm8.staticflickr.com/7231/7359341784_4c5358197f_z.jpg",
+        "http://farm8.staticflickr.com/7283/8737653089_d0c77b8597_z.jpg",
+        "http://farm3.staticflickr.com/2454/3989339438_2f32b76ebb_z.jpg",
+        "http://farm1.staticflickr.com/34/123005230_13051344b1_z.jpg",
+]]
+
+
+# Copied from https://github.com/apple/coremltools/blob/7.0b1/coremltools/optimize/coreml/_quantization_passes.py#L602
+from coremltools.converters.mil.mil import types
+def fake_linear_quantize(val, axis=-1, mode='LINEAR', dtype=types.int8):
+    from coremltools.optimize.coreml._quantization_passes import AffineQuantParams
+    from coremltools.converters.mil.mil.types.type_mapping import nptype_from_builtin
+
+    val_dtype = val.dtype
+    def _ensure_numerical_range_and_cast(val, low, high, np_dtype):
+        '''
+        For some cases, the computed quantized data might exceed the data range.
+        For instance, after rounding and addition, we might get `128` for the int8 quantization.
+        This utility function ensures the val in the data range before doing the cast.
+        '''
+        val = np.minimum(val, high)
+        val = np.maximum(val, low)
+        return val.astype(np_dtype)
+
+    mode_dtype_to_range = {
+        (types.int8, "LINEAR"): (-128, 127),
+        (types.int8, "LINEAR_SYMMETRIC"): (-127, 127),
+        (types.uint8, "LINEAR"): (0, 255),
+        (types.uint8, "LINEAR_SYMMETRIC"): (0, 254),
+    }
+
+    if not isinstance(val, (np.ndarray, np.generic)):
+        raise ValueError("Only numpy arrays are supported")
+
+    params = AffineQuantParams()
+    axes = tuple([i for i in range(len(val.shape)) if i != axis])
+    val_min = np.amin(val, axis=axes, keepdims=True)
+    val_max = np.amax(val, axis=axes, keepdims=True)
+
+    if mode == "LINEAR_SYMMETRIC":
+        # For the linear_symmetric mode, the range is symmetrical to 0
+        max_abs = np.maximum(np.abs(val_min), np.abs(val_max))
+        val_min = -max_abs
+        val_max = max_abs
+    else:
+        assert mode == "LINEAR"
+        # For the linear mode, we need to make sure the data range contains `0`
+        val_min = np.minimum(0.0, val_min)
+        val_max = np.maximum(0.0, val_max)
+
+    q_val_min, q_val_max = mode_dtype_to_range[(dtype, mode)]
+
+    # Set the zero point to symmetric mode
+    np_dtype = nptype_from_builtin(dtype)
+    if mode == "LINEAR_SYMMETRIC":
+        if dtype == types.int8:
+            params.zero_point = (0 * np.ones(val_min.shape)).astype(np.int8)
+        else:
+            assert dtype == types.uint8
+            params.zero_point = (127 * np.ones(val_min.shape)).astype(np.uint8)
+    else:
+        assert mode == "LINEAR"
+        params.zero_point = (q_val_min * val_max - q_val_max * val_min) / (val_max - val_min)
+        params.zero_point = np.round(params.zero_point)
+        params.zero_point = _ensure_numerical_range_and_cast(params.zero_point, q_val_min, q_val_max, np_dtype)
+
+    # compute the params
+    params.scale = (val_max - val_min) / (q_val_max - q_val_min)
+    params.scale = params.scale.astype(val.dtype).squeeze()
+
+    params.quantized_data = np.round(
+        val * (q_val_max - q_val_min) / (val_max - val_min)
+    )
+    params.quantized_data = (params.quantized_data + params.zero_point)
+    params.quantized_data = _ensure_numerical_range_and_cast(params.quantized_data, q_val_min, q_val_max, np_dtype)
+
+    params.zero_point = params.zero_point.squeeze()
+    params.axis = axis
+
+    return (params.quantized_data.astype(val_dtype) - params.zero_point.astype(val_dtype)) * params.scale
+
+
+# Copied from https://github.com/apple/coremltools/blob/7.0b1/coremltools/optimize/coreml/_quantization_passes.py#L423
+def fake_palettize(module, nbits, in_ngroups=1, out_ngroups=1):
+    """ Simulate weight palettization
+    """
+    from coremltools.models.neural_network.quantization_utils import _get_kmeans_lookup_table_and_weight
+
+    def compress_kmeans(val, nbits):
+        lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val)
+        lut = lut.astype(val.dtype)
+        indices = indices.astype(np.uint8)
+        return lut, indices
+
+    dtype = module.weight.data.dtype
+    device = module.weight.data.device
+    val = module.weight.data.cpu().numpy().astype(np.float16)
+    
+    if out_ngroups == 1 and in_ngroups == 1:
+        lut, indices = compress_kmeans(val=val, nbits=nbits)
+        module.weight.data = torch.from_numpy(lut[indices]).reshape(val.shape).to(dtype)
+
+    elif out_ngroups > 1 and in_ngroups == 1:
+        assert val.shape[0] % out_ngroups == 0
+        rvals = [
+            compress_kmeans(val=chunked_val, nbits=nbits)
+            for chunked_val in np.split(val, out_ngroups, axis=0)
+        ]
+        shape = list(val.shape)
+        shape[0] = shape[0] // out_ngroups
+        module.weight.data = torch.cat([
+            torch.from_numpy(lut[indices]).reshape(shape)
+            for lut,indices in rvals
+        ], dim=0).to(dtype).to(device)
+ 
+    elif in_ngroups > 1 and out_ngroups == 1:
+        assert val.shape[1] % in_ngroups == 0
+        rvals = [
+            compress_kmeans(val=chunked_val, nbits=nbits)
+            for chunked_val in np.split(val, in_ngroups, axis=1)
+        ]
+        shape = list(val.shape)
+        shape[1] = shape[1] // in_ngroups
+        module.weight.data = torch.cat([
+            torch.from_numpy(lut[indices]).reshape(shape)
+            for lut,indices in rvals
+        ], dim=1).to(dtype).to(device)
+    else:
+        raise ValueError(f"in_ngroups={in_ngroups} & out_ngroups={out_ngroups} is illegal!!!")
+    
+    return torch.from_numpy(val).to(dtype)
+
+
+def restore_weight(module, value):
+    device = module.weight.data.device
+    module.weight.data = value.to(device)
+
+
+def get_palettizable_modules(unet, min_size=PALETTIZE_MIN_SIZE):
+    ret = [
+        (name, getattr(module, 'weight').data.numel()) for name, module in unet.named_modules()
+        if isinstance(module, (nn.Linear, nn.Conv2d))
+        if hasattr(module, 'weight') and getattr(module, 'weight').data.numel() > min_size
+    ]
+    candidates, sizes = [[a for a,b in ret], [b for a,b in ret]]
+    logger.info(f"{len(candidates)} candidate tensors with {sum(sizes)/1e6} M total params")
+    return candidates, sizes
+
+
+def fake_int8_quantize(module):
+    i = 0
+    for name, submodule in tqdm(module.named_modules()):
+        if hasattr(submodule, 'weight'):
+            i+=1
+            submodule.weight.data = torch.from_numpy(
+                fake_linear_quantize(submodule.weight.data.numpy()))
+    logger.info(f"{i} modules fake int8 quantized")
+    return module
+
+
+def fake_nbits_palette(module, nbits):
+    i = 0
+    for name, submodule in tqdm(module.named_modules()):
+        if hasattr(submodule, 'weight'):
+            i+=1
+            fake_palettize(submodule, nbits=nbits)
+    logger.info(f"{i} modules fake {nbits}-bits palettized")
+    return module
+
+
+def fake_palette_from_recipe(module, recipe):
+    tot_bits = 0
+    tot_numel = 0
+    for name, submodule in tqdm(module.named_modules()):
+        if hasattr(submodule, 'weight'):
+            tot_numel += submodule.weight.numel()
+            if name in recipe:
+                nbits = recipe[name]
+                assert nbits in NBITS + [16]
+                tot_bits += submodule.weight.numel() * nbits
+                if nbits == 16:
+                    continue
+                fake_palettize(submodule, nbits=nbits)
+            else:
+                tot_bits += submodule.weight.numel() * 16
+
+    logger.info(f"Palettized to {tot_bits/tot_numel:.2f}-bits mixed palette ({tot_bits/8e6} MB) ")
+
+# Globally synced RNG state
+rng = torch.Generator()
+rng_state = rng.get_state()
+
+def run_pipe(pipe):
+    if torch.backends.mps.is_available():
+        device = "mps"
+    elif torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    logger.debug(f"Placing pipe in {device}")
+
+    global rng, rng_state
+    rng.set_state(rng_state)
+    kwargs = dict(
+        prompt=RANDOM_TEST_DATA,
+        negative_prompt=[""] * len(RANDOM_TEST_DATA),
+        num_inference_steps=1,
+        height=TEST_RESOLUTION,
+        width=TEST_RESOLUTION,
+        output_type="latent",
+        generator=rng
+    )
+    if "Img2Img" in pipe.__class__.__name__:
+        kwargs["image"] = RANDOM_TEST_IMAGE_DATA
+        kwargs.pop("height")
+        kwargs.pop("width")
+
+        # Run a single denoising step
+        kwargs["num_inference_steps"] = 4
+        kwargs["strength"] = 0.25
+
+    return np.array([latent.cpu().numpy() for latent in pipe.to(device)(**kwargs).images])
+
+
+def benchmark_signal_integrity(pipe,
+                               candidates,
+                               nbits,
+                               cumulative,
+                               in_ngroups=1,
+                               out_ngroups=1,
+                               ref_out=None,
+                               ):
+    results = {}
+    results['metadata'] = {
+        'nbits': nbits,
+        'out_ngroups': out_ngroups,
+        'in_ngroups': in_ngroups,
+        'cumulative': cumulative,
+    }
+
+    # If reference outputs are not provided, treat current pipe as reference
+    if ref_out is None:
+        ref_out = run_pipe(pipe)
+
+    for candidate in tqdm(candidates):
+        palettized = False
+        for name, module in pipe.unet.named_modules():
+            if name == candidate:
+                orig_weight = fake_palettize(
+                    module,
+                    nbits,
+                    out_ngroups=out_ngroups,
+                    in_ngroups=in_ngroups,
+                )
+                palettized = True
+                break
+        if not palettized:
+            raise KeyError(name)
+
+        test_out = run_pipe(pipe)
+
+        if not cumulative:
+            restore_weight(module, orig_weight)
+
+        results[candidate] = [
+            float(f"{compute_psnr(r,t):.1f}")
+            for r,t in zip(ref_out, test_out)
+        ]
+        logger.info(f"{nbits}-bit: {candidate} = {results[candidate]}")
+
+    return results
+            
+
+def descending_psnr_order(results):
+    if 'metadata' in results:
+        results.pop('metadata')
+
+    return OrderedDict(sorted(results.items(), key=lambda items: -sum(items[1])))
+
+
+def simulate_quant_fn(ref_pipe, quantization_to_simulate):
+    simulated_pipe = deepcopy(ref_pipe.to('cpu'))
+    quantization_to_simulate(simulated_pipe.unet)
+    simulated_out = run_pipe(simulated_pipe)
+    del simulated_pipe
+    gc.collect()
+
+    ref_out = run_pipe(ref_pipe)
+    simulated_psnr = sum([
+        float(f"{compute_psnr(r, t):.1f}")
+        for r, t in zip(ref_out, simulated_out)
+    ]) / len(ref_out)
+
+    return simulated_out, simulated_psnr
+
+
+def build_recipe(results, sizes, psnr_threshold, default_nbits):
+    stats = {'nbits': 0}
+    recipe = {}
+
+    for key in results[str(NBITS[0])]:
+        if key == 'metadata':
+            continue
+
+        achieved_nbits = default_nbits
+        for nbits in NBITS:
+            avg_psnr = sum(results[str(nbits)][key])/len(RANDOM_TEST_DATA)
+            if avg_psnr > psnr_threshold:
+                achieved_nbits = nbits
+                break
+        recipe[key] = achieved_nbits
+        stats['nbits'] += achieved_nbits * sizes[key]
+
+    stats['size_mb'] = stats['nbits'] / (8*1e6)
+    tot_size = sum(list(sizes.values()))
+    stats['nbits'] /= tot_size
+
+    return recipe, stats
+
+
+def plot(results, args):
+    import matplotlib.pyplot as plt
+    max_model_size = sum(results['cumulative'][str(NBITS[0])]['metadata']['sizes'])
+    f, ax = plt.subplots(1, 1, figsize=(7, 5))
+
+    def compute_x_axis(sizes, nbits, default_nbits):
+        max_compression_percent = (default_nbits - nbits) / default_nbits
+        progress = np.cumsum(sizes)
+        normalized_progress = progress / progress.max()
+
+        return normalized_progress * max_compression_percent * 100
+
+    # Linear 8-bit baseline and the intercept points for mixed-bit recipes
+    linear8bit_baseline = results['baselines']['linear_8bit']
+
+    # Mark the linear 8-bit baseline
+    ax.plot(
+        8 / args.default_nbits * 100,
+        linear8bit_baseline,
+        'bx',
+        markersize=8,
+        label="8-bit (linear quant)")
+
+    # Plot the iso-dB line that matches the 8-bit baseline
+    ax.plot([0,100], [linear8bit_baseline]*2, '--b')
+
+    # Plot non-mixed-bit palettization curves
+    for idx, nbits in enumerate(NBITS):
+        size_keys = compute_x_axis(results['cumulative'][str(nbits)]['metadata']['sizes'], nbits, args.default_nbits)
+        psnr = [
+            sum(v) / len(RANDOM_TEST_DATA) # avg psnr
+            for k,v in results['cumulative'][str(nbits)].items() if k != 'metadata'
+        ]
+        ax.plot(
+            size_keys,
+            psnr,
+            label=f"{nbits}-bit")
+
+
+    # Plot mixed-bit results
+    mixed_palettes = [
+        (float(spec.rsplit('_')[1]), psnr)
+        for spec,psnr in results['baselines'].items()
+        if 'recipe' in spec
+    ]
+    mixedbit_sizes = [100. * (1. - a[0] / args.default_nbits) for a in mixed_palettes]
+    mixedbit_psnrs = [a[1] for a in mixed_palettes]
+    ax.plot(
+        mixedbit_sizes,
+        mixedbit_psnrs,
+        label="mixed-bit",
+    )
+
+    ax.set_xlabel("Model Size Reduction (%)")
+    ax.set_ylabel("Signal Integrity (PSNR in dB)")
+    ax.set_title(args.model_version)
+    ax.legend()
+
+    f.savefig(os.path.join(args.o, f"{args.model_version.replace('/','_')}_psnr_vs_size.png"))
+
+def main(args):
+
+    # Initialize pipe
+    pipe = get_pipeline(args)
+
+    # Preserve a pristine copy for reference outputs
+    ref_pipe = deepcopy(pipe)
+    if args.default_nbits != 16:
+        logger.info(f"Palettizing unet to default {args.default_nbits}-bit")
+        fake_nbits_palette(pipe.unet, args.default_nbits)
+        logger.info("Done.")
+
+    # Cache reference outputs
+    ref_out = run_pipe(pipe)
+
+    # Bookkeeping
+    os.makedirs(args.o, exist_ok=True)
+
+    results = {
+        'single_layer': {},
+        'cumulative': {},
+        'model_version': args.model_version,
+    }
+    json_name = f"{args.model_version.replace('/','-')}_palettization_recipe.json"
+    candidates, sizes = get_palettizable_modules(pipe.unet)
+
+    sizes_table = dict(zip(candidates, sizes))
+
+    if os.path.isfile(os.path.join(args.o, json_name)):
+        with open(os.path.join(args.o, json_name), "r") as f:
+            results = json.load(f)
+
+    # Analyze uniform-precision palettization impact on signal integrity
+    for nbits in NBITS:
+        if str(nbits) not in results['single_layer']:
+            # Measure the impact of palettization of each layer independently
+            results['single_layer'][str(nbits)] = benchmark_signal_integrity(
+                pipe,
+                candidates,
+                nbits,
+                cumulative=False,
+                ref_out=ref_out,
+            )
+            with open(os.path.join(args.o, json_name), 'w') as f:
+                json.dump(results, f, indent=2)
+
+        # Measure the cumulative impact of palettization based on ascending individual impact computed earlier
+        sorted_candidates = descending_psnr_order(results['single_layer'][str(nbits)])
+
+        if str(nbits) not in results['cumulative']:
+            results['cumulative'][str(nbits)] = benchmark_signal_integrity(
+                deepcopy(pipe),
+                sorted_candidates,
+                nbits,
+                cumulative=True,
+                ref_out=ref_out,
+            )
+            results['cumulative'][str(nbits)]['metadata'].update({
+                'candidates': list(sorted_candidates.keys()),
+                'sizes': [sizes_table[candidate] for candidate in sorted_candidates],
+            })
+
+            with open(os.path.join(args.o, json_name), 'w') as f:
+                json.dump(results, f, indent=2)
+
+    # Generate uniform-quantization baselines
+    results['baselines'] = {
+        "original": simulate_quant_fn(ref_pipe, lambda x: x)[1],
+        "linear_8bit": simulate_quant_fn(ref_pipe, fake_int8_quantize)[1],
+    }
+    with open(os.path.join(args.o, json_name), 'w') as f:
+        json.dump(results, f, indent=2)
+
+
+    # Generate mixed-bit recipes via decreasing PSNR thresholds
+    results['recipes'] = {}
+    recipe_psnr_thresholds = np.linspace(
+        results['baselines']['original'] - 1,
+        results['baselines']["linear_8bit"] + 5,
+        args.num_recipes,
+    )
+
+    for recipe_no, psnr_threshold in enumerate(recipe_psnr_thresholds):
+        logger.info(f"Building recipe #{recipe_no}")
+        recipe, stats = build_recipe(
+            results['cumulative'],
+            sizes_table,
+            psnr_threshold,
+            args.default_nbits,
+        )
+        achieved_psnr = simulate_quant_fn(ref_pipe, lambda x: partial(fake_palette_from_recipe, recipe=recipe)(x))[1]
+        logger.info(
+            f"Recipe #{recipe_no}: {stats['nbits']:.2f}-bits @ per-layer {psnr_threshold} dB, "
+            f"end-to-end {achieved_psnr} dB & "
+            f"{stats['size_mb']:.2f} MB"
+        )
+
+        # Save achieved PSNR and compressed size
+        recipe_key = f"recipe_{stats['nbits']:.2f}_bit_mixedpalette"
+        results['baselines'][recipe_key] = float(f"{achieved_psnr:.1f}")
+        results['recipes'][recipe_key] = recipe
+
+        with open(os.path.join(args.o, json_name), 'w') as f:
+            json.dump(results, f, indent=2)
+
+    # Plot model size vs signal integrity
+    plot(results, args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        required=True,
+        help="Output directory to save the palettization artifacts (recipe json, PSNR plots etc.)"
+    )
+    parser.add_argument(
+        "--model-version",
+        required=True,
+        help=
+        ("The pre-trained model checkpoint and configuration to restore. "
+         "For available versions: https://huggingface.co/models?search=stable-diffusion"
+     ))
+    parser.add_argument(
+        "--default-nbits",
+        help="Default number of bits to use for palettization",
+        choices=tuple(NBITS + [16]),
+        default=16,
+        type=int,
+    )
+    parser.add_argument(
+        "--num-recipes",
+        help="Maximum number of recipes to generate (with decreasing model size and signal integrity)",
+        default=7,
+        type=int,
+    )
+    parser.add_argument(
+        "--custom-vae-version",
+        type=str,
+        default=None,
+        help=
+        ("Custom VAE checkpoint to override the pipeline's built-in VAE. "
+         "If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
+         "No precision override is applied when using a custom VAE."
+         ))
+
+    args = parser.parse_args()
+    main(args)

python_coreml_stable_diffusion/multilingual_projection.py

@@ -0,0 +1,60 @@
+from python_coreml_stable_diffusion.torch2coreml import _compile_coreml_model
+
+import argparse
+import coremltools as ct
+import numpy as np
+import os
+import torch
+import torch.nn as nn
+
+# TODO: Read these values off of the NLContextualEmbedding API to enforce dimensions and track API versioning
+MAX_SEQUENCE_LENGTH = 256
+EMBED_DIM = 512
+BATCH_SIZE = 1
+
+def main(args):
+    # Layer that was trained to map NLContextualEmbedding to your text_encoder.hidden_size dimensionality
+    text_encoder_projection = torch.jit.load(args.input_path)
+
+    # Prepare random inputs for tracing the network before conversion
+    random_input = torch.randn(BATCH_SIZE, MAX_SEQUENCE_LENGTH, EMBED_DIM)
+
+    # Create a class to bake in the reshape operations required to fit the existing model interface
+    class TextEncoderProjection(nn.Module):
+        def __init__(self, proj):
+            super().__init__()
+            self.proj = proj
+
+        def forward(self, x):
+            return self.proj(x).transpose(1, 2).unsqueeze(2) # BSC, BC1S
+
+    # Trace the torch model
+    text_encoder_projection = torch.jit.trace(TextEncoderProjection(text_encoder_projection), (random_input,))
+
+    # Convert the model to Core ML
+    mlpackage_path = os.path.join(args.output_dir, "MultilingualTextEncoderProjection.mlpackage")
+    ct.convert(
+        text_encoder_projection,
+        inputs=[ct.TensorType('nlcontextualembeddings_output', shape=(1, MAX_SEQUENCE_LENGTH, EMBED_DIM), dtype=np.float32)],
+        outputs=[ct.TensorType('encoder_hidden_states', dtype=np.float32)],
+        minimum_deployment_target=ct.target.macOS14,  # NLContextualEmbedding minimum availability build
+        convert_to='mlprogram',
+    ).save()
+
+    # Compile the model and save it under the specified directory
+    _compile_coreml_model(mlpackage_path, args.output_dir, final_name="MultilingualTextEncoderProjection")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input-path",
+        help="Path to the torchscript file that contains the projection layer"
+    )
+    parser.add_argument(
+        "--output-dir",
+        help="Output directory in which the Core ML model should be saved",
+    )
+    args = parser.parse_args()
+
+    main(args)

python_coreml_stable_diffusion/pipeline.py

@@ -0,0 +1,858 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+import argparse
+
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.schedulers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+)
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+import gc
+import inspect
+
+import logging
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+import numpy as np
+import os
+
+from python_coreml_stable_diffusion.coreml_model import (
+    CoreMLModel,
+    _load_mlpackage,
+    _load_mlpackage_controlnet,
+    get_available_compute_units,
+)
+
+import time
+import torch  # Only used for `torch.from_tensor` in `pipe.scheduler.step()`
+from transformers import CLIPFeatureExtractor, CLIPTokenizer
+from typing import List, Optional, Union, Tuple
+from PIL import Image
+
+
+class CoreMLStableDiffusionPipeline(DiffusionPipeline):
+    """ Core ML version of
+    `diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline`
+    """
+
+    def __init__(
+            self,
+            text_encoder: CoreMLModel,
+            unet: CoreMLModel,
+            vae_decoder: CoreMLModel,
+            scheduler: Union[
+                DDIMScheduler,
+                DPMSolverMultistepScheduler,
+                EulerAncestralDiscreteScheduler,
+                EulerDiscreteScheduler,
+                LMSDiscreteScheduler,
+                PNDMScheduler
+            ],
+            tokenizer: CLIPTokenizer,
+            controlnet: Optional[List[CoreMLModel]],
+            xl: Optional[bool] = False,
+            force_zeros_for_empty_prompt: Optional[bool] = True,
+            feature_extractor: Optional[CLIPFeatureExtractor] = None,
+            safety_checker: Optional[CoreMLModel] = None,
+            text_encoder_2: Optional[CoreMLModel] = None,
+            tokenizer_2: Optional[CLIPTokenizer] = None
+
+    ):
+        super().__init__()
+
+        # Register non-Core ML components of the pipeline similar to the original pipeline
+        self.register_modules(
+            tokenizer=tokenizer,
+            scheduler=scheduler,
+            feature_extractor=feature_extractor,
+        )
+
+        if safety_checker is None:
+            # Reproduce original warning:
+            # https://github.com/huggingface/diffusers/blob/v0.9.0/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L119
+            logger.warning(
+                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
+                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
+                " results in services or applications open to the public. Both the diffusers team and Hugging Face"
+                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
+                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
+                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
+            )
+        self.xl = xl
+        self.force_zeros_for_empty_prompt = force_zeros_for_empty_prompt
+
+        # Register Core ML components of the pipeline
+        self.safety_checker = safety_checker
+        self.text_encoder = text_encoder
+        self.text_encoder_2 = text_encoder_2
+        self.tokenizer_2 = tokenizer_2
+        self.unet = unet
+        self.unet.in_channels = self.unet.expected_inputs["sample"]["shape"][1]
+
+        self.controlnet = controlnet
+
+        self.vae_decoder = vae_decoder
+
+        VAE_DECODER_UPSAMPLE_FACTOR = 8
+
+        # In PyTorch, users can determine the tensor shapes dynamically by default
+        # In CoreML, tensors have static shapes unless flexible shapes were used during export
+        # See https://coremltools.readme.io/docs/flexible-inputs
+        latent_h, latent_w = self.unet.expected_inputs["sample"]["shape"][2:]
+        self.height = latent_h * VAE_DECODER_UPSAMPLE_FACTOR
+        self.width = latent_w * VAE_DECODER_UPSAMPLE_FACTOR
+
+        logger.info(
+            f"Stable Diffusion configured to generate {self.height}x{self.width} images"
+        )
+
+    def _encode_prompt(self,
+                       prompt,
+                       prompt_2: Optional[str] = None,
+                       do_classifier_free_guidance: bool = True,
+                       negative_prompt: Optional[str] = None,
+                       negative_prompt_2: Optional[str] = None,
+                       ):
+
+        batch_size = len(prompt) if isinstance(prompt, list) else 1
+
+        if self.xl is True:
+            prompts = [prompt, prompt_2] if prompt_2 is not None else [prompt, prompt]
+
+            # refiner uses only one tokenizer and text encoder (tokenizer_2 and text_encoder_2)
+            tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
+
+            text_encoders = [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [
+                self.text_encoder_2]
+            hidden_state_key = 'hidden_embeds'
+        else:
+            prompts = [prompt]
+            tokenizers = [self.tokenizer]
+            text_encoders = [self.text_encoder]
+            hidden_state_key = 'last_hidden_state'
+
+        prompt_embeds_list = []
+        for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
+
+            text_inputs = tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="np",
+            )
+
+            text_input_ids = text_inputs.input_ids
+
+            # tokenize without max_length to catch any truncation
+            untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="np").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not np.equal(
+                    text_input_ids, untruncated_ids
+            ):
+                removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1: -1])
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            embeddings = text_encoder(input_ids=text_input_ids.astype(np.float32))
+
+            prompt_embeds_list.append(embeddings[hidden_state_key])
+
+            # We are only ALWAYS interested in the pooled output of the final text encoder
+            if self.xl:
+                pooled_prompt_embeds = embeddings['pooled_outputs']
+
+        prompt_embeds = np.concatenate(prompt_embeds_list, axis=-1)
+
+        if do_classifier_free_guidance and negative_prompt is None and self.force_zeros_for_empty_prompt:
+            negative_prompt_embeds = np.zeros_like(prompt_embeds)
+
+            if self.xl:
+                negative_pooled_prompt_embeds = np.zeros_like(pooled_prompt_embeds)
+        elif do_classifier_free_guidance:
+
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+
+            uncond_tokens: List[str]
+            if prompts is not None and type(prompts) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`.")
+            else:
+                uncond_tokens = [negative_prompt, negative_prompt_2]
+
+            negative_prompt_embeds_list = []
+            for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
+
+                max_length = prompt_embeds.shape[1]
+
+                uncond_input = tokenizer(
+                    negative_prompt,
+                    padding="max_length",
+                    max_length=max_length,
+                    truncation=True,
+                    return_tensors="np",
+                )
+                uncond_input_ids = uncond_input.input_ids
+
+                negative_embeddings = text_encoder(
+                    input_ids=uncond_input_ids.astype(np.float32)
+                )
+
+                negative_text_embeddings = negative_embeddings[hidden_state_key]
+
+                negative_prompt_embeds_list.append(negative_text_embeddings)
+
+                # We are only ALWAYS interested in the pooled output of the final text encoder
+                if self.xl:
+                    negative_pooled_prompt_embeds = negative_embeddings['pooled_outputs']
+
+            negative_prompt_embeds = np.concatenate(negative_prompt_embeds_list, axis=-1)
+
+        if do_classifier_free_guidance:
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            prompt_embeds = np.concatenate(
+                [negative_prompt_embeds, prompt_embeds])
+
+            if self.xl:
+                pooled_prompt_embeds = np.concatenate(
+                    [negative_pooled_prompt_embeds, pooled_prompt_embeds])
+
+        prompt_embeddings = prompt_embeds.transpose(0, 2, 1)[:, :, None, :]
+
+        if self.xl:
+            return prompt_embeddings, pooled_prompt_embeds
+        else:
+            return prompt_embeddings, None
+
+    def run_controlnet(self,
+                       sample,
+                       timestep,
+                       encoder_hidden_states,
+                       controlnet_cond,
+                       output_dtype=np.float16):
+        if not self.controlnet:
+            raise ValueError(
+                "Conditions for controlnet are given but the pipeline has no controlnet modules")
+
+        for i, (module, cond) in enumerate(zip(self.controlnet, controlnet_cond)):
+            module_outputs = module(
+                sample=sample.astype(np.float16),
+                timestep=timestep.astype(np.float16),
+                encoder_hidden_states=encoder_hidden_states.astype(np.float16),
+                controlnet_cond=cond.astype(np.float16),
+            )
+            if i == 0:
+                outputs = module_outputs
+            else:
+                for key in outputs.keys():
+                    outputs[key] += module_outputs[key]
+
+        outputs = {k: v.astype(output_dtype) for k, v in outputs.items()}
+
+        return outputs
+
+    def run_safety_checker(self, image):
+        if self.safety_checker is not None:
+            safety_checker_input = self.feature_extractor(
+                self.numpy_to_pil(image),
+                return_tensors="np",
+            )
+
+            safety_checker_outputs = self.safety_checker(
+                clip_input=safety_checker_input.pixel_values.astype(
+                    np.float16),
+                images=image.astype(np.float16),
+                adjustment=np.array([0.]).astype(
+                    np.float16),  # defaults to 0 in original pipeline
+            )
+
+            # Unpack dict
+            has_nsfw_concept = safety_checker_outputs["has_nsfw_concepts"]
+            image = safety_checker_outputs["filtered_images"]
+            concept_scores = safety_checker_outputs["concept_scores"]
+
+            logger.info(
+                f"Generated image has nsfw concept={has_nsfw_concept.any()}")
+        else:
+            has_nsfw_concept = None
+
+        return image, has_nsfw_concept
+
+    def decode_latents(self, latents):
+        latents = 1 / 0.18215 * latents
+        dtype = self.vae_decoder.expected_inputs['z']['dtype']
+        image = self.vae_decoder(z=latents.astype(dtype))["image"]
+        image = np.clip(image / 2 + 0.5, 0, 1)
+        image = image.transpose((0, 2, 3, 1))
+
+        return image
+
+    def prepare_latents(self,
+                        batch_size,
+                        num_channels_latents,
+                        height,
+                        width,
+                        latents=None):
+        latents_shape = (batch_size, num_channels_latents, self.height // 8,
+                         self.width // 8)
+        if latents is None:
+            latents = np.random.randn(*latents_shape).astype(np.float16)
+        elif latents.shape != latents_shape:
+            raise ValueError(
+                f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"
+            )
+
+        init_noise = self.scheduler.init_noise_sigma
+
+        if isinstance(init_noise, torch.Tensor):
+            init_noise = init_noise.numpy()
+
+        latents = latents * init_noise
+
+        return latents
+
+    def prepare_control_cond(self,
+                             controlnet_cond,
+                             do_classifier_free_guidance,
+                             batch_size,
+                             num_images_per_prompt):
+        processed_cond_list = []
+        for cond in controlnet_cond:
+            cond = np.stack([cond] * batch_size * num_images_per_prompt)
+            if do_classifier_free_guidance:
+                cond = np.concatenate([cond] * 2)
+            processed_cond_list.append(cond)
+        return processed_cond_list
+
+    def check_inputs(self, prompt, height, width, callback_steps):
+        if height != self.height or width != self.width:
+            logger.warning(
+                "`height` and `width` dimensions (of the output image tensor) are fixed when exporting the Core ML models " \
+                "unless flexible shapes are used during export (https://coremltools.readme.io/docs/flexible-inputs). " \
+                "This pipeline was provided with Core ML models that generate {self.height}x{self.width} images (user requested {height}x{width})"
+            )
+
+        if not isinstance(prompt, str) and not isinstance(prompt, list):
+            raise ValueError(
+                f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
+            )
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(
+                f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
+            )
+
+        if (callback_steps is None) or (callback_steps is not None and
+                                        (not isinstance(callback_steps, int)
+                                         or callback_steps <= 0)):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}.")
+
+    def prepare_extra_step_kwargs(self, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(
+            inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        return extra_step_kwargs
+
+    def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
+        add_time_ids = list(original_size + crops_coords_top_left + target_size)
+        add_time_ids = np.array(add_time_ids).astype(dtype)
+        return add_time_ids
+
+    def __call__(
+            self,
+            prompt,
+            height=512,
+            width=512,
+            num_inference_steps=50,
+            guidance_scale=7.5,
+            negative_prompt=None,
+            num_images_per_prompt=1,
+            eta=0.0,
+            latents=None,
+            output_type="pil",
+            return_dict=True,
+            callback=None,
+            callback_steps=1,
+            controlnet_cond=None,
+            original_size: Optional[Tuple[int, int]]=None,
+            crops_coords_top_left: Tuple[int, int]=(0, 0),
+            target_size: Optional[Tuple[int, int]]=None,
+            unet_batch_one=False,
+            **kwargs,
+    ):
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(prompt, height, width, callback_steps)
+        height = height or self.height
+        width = width or self.width
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 2. Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        if batch_size > 1 or num_images_per_prompt > 1:
+            raise NotImplementedError(
+                "For batched generation of multiple images and/or multiple prompts, please refer to the Swift package."
+            )
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_embeddings, pooled_prompt_embeds = self._encode_prompt(
+            prompt=prompt,
+            prompt_2=None,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=None
+        )
+
+        # 4. Prepare XL kwargs if needed
+        unet_additional_kwargs = {}
+
+        # we add pooled prompt embeds + time_ids to unet kwargs
+        if self.xl:
+            add_text_embeds = pooled_prompt_embeds
+            add_time_ids = self._get_add_time_ids(original_size, crops_coords_top_left, target_size,
+                                                  text_embeddings.dtype)
+            if do_classifier_free_guidance:
+
+                # TODO: This checks if the time_ids input is looking for time_ids.shape == (12,) or (2, 6)
+                # Remove once model input shapes are ubiquitous
+                if len(self.unet.expected_inputs['time_ids']['shape']) > 1:
+                    add_time_ids = [add_time_ids]
+
+                add_time_ids = np.concatenate([add_time_ids, add_time_ids])
+
+            unet_additional_kwargs.update({'text_embeds': add_text_embeds.astype(np.float16),
+                                           'time_ids': add_time_ids.astype(np.float16)})
+
+        # 5. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps)
+        timesteps = self.scheduler.timesteps
+
+        # 6. Prepare latent variables and controlnet cond
+        num_channels_latents = self.unet.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            latents,
+        )
+
+        if controlnet_cond:
+            controlnet_cond = self.prepare_control_cond(
+                controlnet_cond,
+                do_classifier_free_guidance,
+                batch_size,
+                num_images_per_prompt,
+            )
+
+        # 7. Prepare extra step kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(eta)
+
+        # 8. Denoising loop
+        for i, t in enumerate(self.progress_bar(timesteps)):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = np.concatenate(
+                [latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(
+                latent_model_input, t)
+
+            if isinstance(latent_model_input, torch.Tensor):
+                latent_model_input = latent_model_input.numpy()
+
+            if do_classifier_free_guidance:
+                timestep = np.array([t, t], np.float16)
+            else:
+                timestep = np.array([t,], np.float16)
+
+            # controlnet
+            if controlnet_cond:
+                control_net_additional_residuals = self.run_controlnet(
+                    sample=latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=text_embeddings,
+                    controlnet_cond=controlnet_cond,
+                )
+            else:
+                control_net_additional_residuals = {}
+
+            # predict the noise residual
+            unet_additional_kwargs.update(control_net_additional_residuals)
+
+            # get prediction from unet
+            if not (unet_batch_one and do_classifier_free_guidance):
+                noise_pred = self.unet(
+                    sample=latent_model_input.astype(np.float16),
+                    timestep=timestep,
+                    encoder_hidden_states=text_embeddings.astype(np.float16),
+                    **unet_additional_kwargs,
+                )["noise_pred"]
+
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
+            else:
+                # query unet sequentially
+                latent_model_input = latent_model_input.astype(np.float16)
+                text_embeddings = text_embeddings.astype(np.float16)
+                timestep = np.array([t,], np.float16)
+
+                noise_pred_uncond = self.unet(
+                    sample=np.expand_dims(latent_model_input[0], axis=0),
+                    timestep=timestep,
+                    encoder_hidden_states=np.expand_dims(text_embeddings[0], axis=0),
+                    **unet_additional_kwargs,
+                )["noise_pred"]
+                noise_pred_text = self.unet(
+                    sample=np.expand_dims(latent_model_input[1], axis=0),
+                    timestep=timestep,
+                    encoder_hidden_states=np.expand_dims(text_embeddings[1], axis=0),
+                    **unet_additional_kwargs,
+                )["noise_pred"]
+
+            # perform guidance
+            if do_classifier_free_guidance:
+                noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(torch.from_numpy(noise_pred),
+                                          t,
+                                          torch.from_numpy(latents),
+                                          **extra_step_kwargs,
+                                          ).prev_sample.numpy()
+
+            # call the callback, if provided
+            if callback is not None and i % callback_steps == 0:
+                callback(i, t, latents)
+
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+
+        # 9. Run safety checker
+        image, has_nsfw_concept = self.run_safety_checker(image)
+
+        # 10. Convert to PIL
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(
+            images=image, nsfw_content_detected=has_nsfw_concept)
+
+
+def get_available_schedulers():
+    schedulers = {}
+    for scheduler in [DDIMScheduler,
+                      DPMSolverMultistepScheduler,
+                      EulerAncestralDiscreteScheduler,
+                      EulerDiscreteScheduler,
+                      LMSDiscreteScheduler,
+                      PNDMScheduler]:
+        schedulers[scheduler().__class__.__name__.replace("Scheduler", "")] = scheduler
+    return schedulers
+
+
+SCHEDULER_MAP = get_available_schedulers()
+
+
+def get_coreml_pipe(pytorch_pipe,
+                    mlpackages_dir,
+                    model_version,
+                    compute_unit,
+                    delete_original_pipe=True,
+                    scheduler_override=None,
+                    controlnet_models=None,
+                    force_zeros_for_empty_prompt=True,
+                    sources=None):
+    """
+    Initializes and returns a `CoreMLStableDiffusionPipeline` from an original
+    diffusers PyTorch pipeline
+        sources: 'packages' or 'compiled' forces creation of model from specified sources. sources must be in mlpackages_dir
+    """
+
+    # Ensure `scheduler_override` object is of correct type if specified
+    if scheduler_override is not None:
+        assert isinstance(scheduler_override, SchedulerMixin)
+        logger.warning(
+            "Overriding scheduler in pipeline: "
+            f"Default={pytorch_pipe.scheduler}, Override={scheduler_override}")
+
+    # Gather configured tokenizer and scheduler attributes from the original pipe
+    if 'xl' in model_version:
+        coreml_pipe_kwargs = {
+            "tokenizer": pytorch_pipe.tokenizer,
+            'tokenizer_2': pytorch_pipe.tokenizer_2,
+            "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
+            'xl': True,
+        }
+
+        model_packages_to_load = ["text_encoder", "text_encoder_2", "unet", "vae_decoder"]
+
+    else:
+        coreml_pipe_kwargs = {
+            "tokenizer": pytorch_pipe.tokenizer,
+            "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
+            "feature_extractor": pytorch_pipe.feature_extractor,
+        }
+        model_packages_to_load = ["text_encoder", "unet", "vae_decoder"]
+
+    coreml_pipe_kwargs["force_zeros_for_empty_prompt"] = force_zeros_for_empty_prompt
+
+    if getattr(pytorch_pipe, "safety_checker", None) is not None:
+        model_packages_to_load.append("safety_checker")
+    else:
+        logger.warning(
+            f"Original diffusers pipeline for {model_version} does not have a safety_checker, "
+            "Core ML pipeline will mirror this behavior.")
+        coreml_pipe_kwargs["safety_checker"] = None
+
+    if delete_original_pipe:
+        del pytorch_pipe
+        gc.collect()
+        logger.info("Removed PyTorch pipe to reduce peak memory consumption")
+
+    if controlnet_models:
+        model_packages_to_load.remove("unet")
+        coreml_pipe_kwargs["unet"] = _load_mlpackage(
+            submodule_name="control-unet",
+            mlpackages_dir=mlpackages_dir,
+            model_version=model_version,
+            compute_unit=compute_unit,
+        )
+        coreml_pipe_kwargs["controlnet"] = [_load_mlpackage_controlnet(
+            mlpackages_dir,
+            model_version,
+            compute_unit,
+        ) for model_version in controlnet_models]
+    else:
+        coreml_pipe_kwargs["controlnet"] = None
+
+    # Load Core ML models
+    logger.info(f"Loading Core ML models in memory from {mlpackages_dir}")
+    coreml_pipe_kwargs.update({
+        model_name: _load_mlpackage(
+            submodule_name=model_name,
+            mlpackages_dir=mlpackages_dir,
+            model_version=model_version,
+            compute_unit=compute_unit,
+            sources=sources,
+        )
+        for model_name in model_packages_to_load
+    })
+    logger.info("Done.")
+
+    logger.info("Initializing Core ML pipe for image generation")
+    coreml_pipe = CoreMLStableDiffusionPipeline(**coreml_pipe_kwargs)
+    logger.info("Done.")
+
+    return coreml_pipe
+
+
+def get_image_path(args, **override_kwargs):
+    """ mkdir output folder and encode metadata in the filename
+    """
+    out_folder = os.path.join(args.o, "_".join(args.prompt.replace("/", "_").rsplit(" ")))
+    os.makedirs(out_folder, exist_ok=True)
+
+    out_fname = f"randomSeed_{override_kwargs.get('seed', None) or args.seed}"
+    out_fname += f"_computeUnit_{override_kwargs.get('compute_unit', None) or args.compute_unit}"
+    out_fname += f"_modelVersion_{override_kwargs.get('model_version', None) or args.model_version.replace('/', '_')}"
+
+    if args.scheduler is not None:
+        out_fname += f"_customScheduler_{override_kwargs.get('scheduler', None) or args.scheduler}"
+        out_fname += f"_numInferenceSteps{override_kwargs.get('num_inference_steps', None) or args.num_inference_steps}"
+
+    return os.path.join(out_folder, out_fname + ".png")
+
+
+def prepare_controlnet_cond(image_path, height, width):
+    image = Image.open(image_path).convert("RGB")
+    image = image.resize((height, width), resample=Image.LANCZOS)
+    image = np.array(image).transpose(2, 0, 1) / 255.0
+    return image
+
+
+def main(args):
+    logger.info(f"Setting random seed to {args.seed}")
+    np.random.seed(args.seed)
+
+    logger.info("Initializing PyTorch pipe for reference configuration")
+
+    SDP = StableDiffusionXLPipeline if 'xl' in args.model_version else StableDiffusionPipeline
+
+    pytorch_pipe = SDP.from_pretrained(
+        args.model_version,
+        use_auth_token=True,
+    )
+
+    # Get Scheduler
+    user_specified_scheduler = None
+    if args.scheduler is not None:
+        user_specified_scheduler = SCHEDULER_MAP[
+            args.scheduler].from_config(pytorch_pipe.scheduler.config)
+
+    # Get Force Zeros Config if it exists
+    force_zeros_for_empty_prompt: bool = False
+    if 'xl' in args.model_version and 'force_zeros_for_empty_prompt' in pytorch_pipe.config:
+        force_zeros_for_empty_prompt = pytorch_pipe.config['force_zeros_for_empty_prompt']
+
+    coreml_pipe = get_coreml_pipe(
+        pytorch_pipe=pytorch_pipe,
+        mlpackages_dir=args.i,
+        model_version=args.model_version,
+        compute_unit=args.compute_unit,
+        scheduler_override=user_specified_scheduler,
+        controlnet_models=args.controlnet,
+        force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
+        sources=args.model_sources,
+    )
+
+    if args.controlnet:
+        controlnet_cond = []
+        for i, _ in enumerate(args.controlnet):
+            image_path = args.controlnet_inputs[i]
+            image = prepare_controlnet_cond(image_path, coreml_pipe.height, coreml_pipe.width)
+            controlnet_cond.append(image)
+    else:
+        controlnet_cond = None
+
+    logger.info("Beginning image generation.")
+    image = coreml_pipe(
+        prompt=args.prompt,
+        height=coreml_pipe.height,
+        width=coreml_pipe.width,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        controlnet_cond=controlnet_cond,
+        negative_prompt=args.negative_prompt,
+        unet_batch_one=args.unet_batch_one,
+    )
+
+    out_path = get_image_path(args)
+    logger.info(f"Saving generated image to {out_path}")
+    image["images"][0].save(out_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--prompt",
+        required=True,
+        help="The text prompt to be used for text-to-image generation.")
+    parser.add_argument(
+        "-i",
+        required=True,
+        help=("Path to input directory with the .mlpackage files generated by "
+              "python_coreml_stable_diffusion.torch2coreml"))
+    parser.add_argument("-o", required=True)
+    parser.add_argument("--seed",
+                        "-s",
+                        default=93,
+                        type=int,
+                        help="Random seed to be able to reproduce results")
+    parser.add_argument(
+        "--model-version",
+        default="CompVis/stable-diffusion-v1-4",
+        help=
+        ("The pre-trained model checkpoint and configuration to restore. "
+         "For available versions: https://huggingface.co/models?search=stable-diffusion"
+         ))
+    parser.add_argument(
+        "--compute-unit",
+        choices=get_available_compute_units(),
+        default="ALL",
+        help=("The compute units to be used when executing Core ML models. "
+              f"Options: {get_available_compute_units()}"))
+    parser.add_argument(
+        "--scheduler",
+        choices=tuple(SCHEDULER_MAP.keys()),
+        default=None,
+        help=("The scheduler to use for running the reverse diffusion process. "
+              "If not specified, the default scheduler from the diffusers pipeline is utilized"))
+    parser.add_argument(
+        "--num-inference-steps",
+        default=50,
+        type=int,
+        help="The number of iterations the unet model will be executed throughout the reverse diffusion process")
+    parser.add_argument(
+        "--guidance-scale",
+        default=7.5,
+        type=float,
+        help="Controls the influence of the text prompt on sampling process (0=random images)")
+    parser.add_argument(
+        "--controlnet",
+        nargs="*",
+        type=str,
+        help=("Enables ControlNet and use control-unet instead of unet for additional inputs. "
+              "For Multi-Controlnet, provide the model names separated by spaces."))
+    parser.add_argument(
+        "--controlnet-inputs",
+        nargs="*",
+        type=str,
+        help=("Image paths for ControlNet inputs. "
+              "Please enter images corresponding to each controlnet provided at --controlnet option in same order."))
+    parser.add_argument(
+        "--negative-prompt",
+        default=None,
+        help="The negative text prompt to be used for text-to-image generation.")
+    parser.add_argument(
+        "--unet-batch-one",
+        action="store_true",
+        help="Do not batch unet predictions for the prompt and negative prompt.")
+    parser.add_argument('--model-sources',
+                        default=None,
+                        choices=['packages', 'compiled'],
+                        help='Force build from `packages` or `compiled`')
+
+    args = parser.parse_args()
+    main(args)