# # 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)