wehub/apple--ml-stable-diffusion
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:46 +08:00
commit 8199cf3c39
90 changed files with 62775 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python_coreml_stable_diffusion/activation_quantization.py
+503
View File
@@ -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)