wehub-resource-sync
1102 lines
52 KiB
Python
1102 lines
52 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
import gc
|
|
import inspect
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import sys
|
|
import time
|
|
from hashlib import md5
|
|
from typing import List, Optional
|
|
|
|
import modelopt.torch.opt as mto
|
|
import modelopt.torch.quantization as mtq
|
|
import numpy as np
|
|
import nvtx
|
|
import tensorrt as trt
|
|
import torch
|
|
from cuda.bindings import runtime as cudart
|
|
from diffusers import (
|
|
DDIMScheduler,
|
|
DDPMScheduler,
|
|
DDPMWuerstchenScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
EulerDiscreteScheduler,
|
|
LCMScheduler,
|
|
LMSDiscreteScheduler,
|
|
PNDMScheduler,
|
|
UniPCMultistepScheduler,
|
|
)
|
|
|
|
import demo_diffusion.engine as engine_module
|
|
import demo_diffusion.image as image_module
|
|
from demo_diffusion.model import (
|
|
CLIPModel,
|
|
CLIPWithProjModel,
|
|
SDLoraLoader,
|
|
UNet2DConditionControlNetModel,
|
|
UNetModel,
|
|
UNetXLModel,
|
|
UNetXLModelControlNet,
|
|
VAEEncoderModel,
|
|
VAEModel,
|
|
get_clip_embedding_dim,
|
|
make_scheduler,
|
|
make_tokenizer,
|
|
merge_loras,
|
|
unload_torch_model,
|
|
)
|
|
from demo_diffusion.pipeline.calibrate import load_calib_prompts
|
|
from demo_diffusion.pipeline.type import PIPELINE_TYPE
|
|
from demo_diffusion.utils_modelopt import (
|
|
SD_FP8_FP16_DEFAULT_CONFIG,
|
|
SD_FP8_FP32_DEFAULT_CONFIG,
|
|
check_lora,
|
|
filter_func,
|
|
generate_fp8_scales,
|
|
get_int8_config,
|
|
quantize_lvl,
|
|
set_fmha,
|
|
)
|
|
|
|
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
|
|
|
|
|
|
class StableDiffusionPipeline:
|
|
SCHEDULER_DEFAULTS = {
|
|
"1.4": "PNDM",
|
|
"dreamshaper-7": "PNDM",
|
|
"xl-1.0" : "Euler",
|
|
"xl-turbo": "EulerA",
|
|
"svd-xt-1.1": "Euler",
|
|
"cascade": "DDPMWuerstchen"
|
|
}
|
|
"""
|
|
Application showcasing the acceleration of Stable Diffusion pipelines using NVidia TensorRT.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
version='1.4',
|
|
pipeline_type=PIPELINE_TYPE.TXT2IMG,
|
|
max_batch_size=16,
|
|
denoising_steps=30,
|
|
scheduler=None,
|
|
guidance_scale=7.5,
|
|
device='cuda',
|
|
output_dir='.',
|
|
hf_token=None,
|
|
verbose=False,
|
|
nvtx_profile=False,
|
|
use_cuda_graph=False,
|
|
vae_scaling_factor=0.18215,
|
|
framework_model_dir='pytorch_model',
|
|
controlnets=None,
|
|
lora_scale: float = 1.0,
|
|
lora_weight: Optional[List[float]] = None,
|
|
lora_path: Optional[List[str]] = None,
|
|
return_latents=False,
|
|
torch_inference='',
|
|
):
|
|
"""
|
|
Initializes the Diffusion pipeline.
|
|
|
|
Args:
|
|
version (str):
|
|
The version of the pipeline. Should be one of [1.4, SDXL]
|
|
pipeline_type (PIPELINE_TYPE):
|
|
Type of current pipeline.
|
|
max_batch_size (int):
|
|
Maximum batch size for dynamic batch engine.
|
|
denoising_steps (int):
|
|
The number of denoising steps.
|
|
More denoising steps usually lead to a higher quality image at the expense of slower inference.
|
|
scheduler (str):
|
|
The scheduler to guide the denoising process. Must be one of [DDIM, DPM, EulerA, Euler, LCM, LMSD, PNDM].
|
|
guidance_scale (float):
|
|
Guidance scale is enabled by setting as > 1.
|
|
Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality.
|
|
device (str):
|
|
PyTorch device to run inference. Default: 'cuda'
|
|
output_dir (str):
|
|
Output directory for log files and image artifacts
|
|
hf_token (str):
|
|
HuggingFace User Access Token to use for downloading Stable Diffusion model checkpoints.
|
|
verbose (bool):
|
|
Enable verbose logging.
|
|
nvtx_profile (bool):
|
|
Insert NVTX profiling markers.
|
|
use_cuda_graph (bool):
|
|
Use CUDA graph to capture engine execution and then launch inference
|
|
vae_scaling_factor (float):
|
|
VAE scaling factor
|
|
framework_model_dir (str):
|
|
cache directory for framework checkpoints
|
|
controlnets (str):
|
|
Which ControlNet/ControlNets to use.
|
|
return_latents (bool):
|
|
Skip decoding the image and return latents instead.
|
|
torch_inference (str):
|
|
Run inference with PyTorch (using specified compilation mode) instead of TensorRT.
|
|
"""
|
|
|
|
self.denoising_steps = denoising_steps
|
|
self.guidance_scale = guidance_scale
|
|
self.do_classifier_free_guidance = (guidance_scale > 1.0)
|
|
self.vae_scaling_factor = vae_scaling_factor
|
|
|
|
self.max_batch_size = max_batch_size
|
|
|
|
self.framework_model_dir = framework_model_dir
|
|
self.output_dir = output_dir
|
|
for directory in [self.framework_model_dir, self.output_dir]:
|
|
if not os.path.exists(directory):
|
|
print(f"[I] Create directory: {directory}")
|
|
pathlib.Path(directory).mkdir(parents=True)
|
|
|
|
self.hf_token = hf_token
|
|
self.device = device
|
|
self.verbose = verbose
|
|
self.nvtx_profile = nvtx_profile
|
|
|
|
self.version = version
|
|
self.controlnets = controlnets
|
|
|
|
# Pipeline type
|
|
self.pipeline_type = pipeline_type
|
|
if self.pipeline_type.is_txt2img():
|
|
self.stages = ['clip','unet','vae']
|
|
elif self.pipeline_type.is_img2img():
|
|
self.stages = ['vae_encoder', 'clip','unet','vae']
|
|
elif self.pipeline_type.is_sd_xl_base():
|
|
self.stages = ['clip', 'clip2', 'unetxl']
|
|
if not return_latents:
|
|
self.stages.append('vae')
|
|
elif self.pipeline_type.is_sd_xl_refiner():
|
|
self.stages = ['clip2', 'unetxl', 'vae']
|
|
elif self.pipeline_type.is_img2vid():
|
|
self.stages = ['clip-vis', 'clip-imgfe', 'unet-temp', 'vae-temp']
|
|
elif self.pipeline_type.is_cascade_prior():
|
|
self.stages = ['clip', 'unet']
|
|
elif self.pipeline_type.is_cascade_decoder():
|
|
self.stages = ['clip', 'unet', 'vqgan']
|
|
else:
|
|
raise ValueError(f"Unsupported pipeline {self.pipeline_type.name}.")
|
|
self.return_latents = return_latents
|
|
|
|
if not scheduler:
|
|
scheduler = 'UniPC' if self.pipeline_type.is_controlnet() else self.SCHEDULER_DEFAULTS.get(version, 'DDIM')
|
|
print(f"[I] Autoselected scheduler: {scheduler}")
|
|
|
|
scheduler_class_map = {
|
|
"DDIM" : DDIMScheduler,
|
|
"DDPM" : DDPMScheduler,
|
|
"EulerA" : EulerAncestralDiscreteScheduler,
|
|
"Euler" : EulerDiscreteScheduler,
|
|
"LCM" : LCMScheduler,
|
|
"LMSD" : LMSDiscreteScheduler,
|
|
"PNDM" : PNDMScheduler,
|
|
"UniPC" : UniPCMultistepScheduler,
|
|
"DDPMWuerstchen" : DDPMWuerstchenScheduler,
|
|
}
|
|
try:
|
|
scheduler_class = scheduler_class_map[scheduler]
|
|
except KeyError:
|
|
raise ValueError(f"Unsupported scheduler {scheduler}. Should be one of {list(scheduler_class.keys())}.")
|
|
self.scheduler = make_scheduler(scheduler_class, version, pipeline_type, hf_token, framework_model_dir)
|
|
|
|
self.config = {}
|
|
if self.pipeline_type.is_sd_xl():
|
|
self.config['clip_hidden_states'] = True
|
|
self.torch_inference = torch_inference
|
|
if self.torch_inference:
|
|
torch._inductor.config.conv_1x1_as_mm = True
|
|
torch._inductor.config.coordinate_descent_tuning = True
|
|
torch._inductor.config.epilogue_fusion = False
|
|
torch._inductor.config.coordinate_descent_check_all_directions = True
|
|
self.use_cuda_graph = use_cuda_graph
|
|
|
|
# initialized in loadEngines()
|
|
self.models = {}
|
|
self.torch_models = {}
|
|
self.engine = {}
|
|
self.shared_device_memory = None
|
|
|
|
# initialize lora loader and scales
|
|
self.lora_loader = None
|
|
self.lora_weights = dict()
|
|
if lora_path:
|
|
self.lora_loader = SDLoraLoader(lora_path, lora_weight, lora_scale)
|
|
assert len(lora_path) == len(lora_weight)
|
|
for i, path in enumerate(lora_path):
|
|
self.lora_weights[path] = lora_weight[i]
|
|
|
|
# initialized in loadResources()
|
|
self.events = {}
|
|
self.generator = None
|
|
self.markers = {}
|
|
self.seed = None
|
|
self.stream = None
|
|
self.tokenizer = None
|
|
|
|
def loadResources(self, image_height, image_width, batch_size, seed):
|
|
# Initialize noise generator
|
|
if seed:
|
|
self.seed = seed
|
|
self.generator = torch.Generator(device="cuda").manual_seed(seed)
|
|
|
|
# Create CUDA events and stream
|
|
for stage in self.stages:
|
|
self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
|
|
self.stream = cudart.cudaStreamCreate()[1]
|
|
|
|
# Allocate TensorRT I/O buffers
|
|
if not self.torch_inference:
|
|
for model_name, obj in self.models.items():
|
|
self.engine[model_name].allocate_buffers(shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.device)
|
|
|
|
def teardown(self):
|
|
for e in self.events.values():
|
|
cudart.cudaEventDestroy(e[0])
|
|
cudart.cudaEventDestroy(e[1])
|
|
|
|
for engine in self.engine.values():
|
|
del engine
|
|
|
|
if self.shared_device_memory:
|
|
cudart.cudaFree(self.shared_device_memory)
|
|
|
|
for torch_model in self.torch_models.values():
|
|
del torch_model
|
|
|
|
cudart.cudaStreamDestroy(self.stream)
|
|
del self.stream
|
|
|
|
def cachedModelName(self, model_name):
|
|
return model_name
|
|
|
|
def getOnnxPath(self, model_name, onnx_dir, opt=True, suffix=''):
|
|
onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+suffix+('.opt' if opt else ''))
|
|
os.makedirs(onnx_model_dir, exist_ok=True)
|
|
return os.path.join(onnx_model_dir, 'model.onnx')
|
|
|
|
def getEnginePath(self, model_name, engine_dir, enable_refit=False, suffix=''):
|
|
return os.path.join(engine_dir, self.cachedModelName(model_name)+suffix+('.refit' if enable_refit else '')+'.trt'+trt.__version__+'.plan')
|
|
|
|
def getWeightsMapPath(self, model_name, onnx_dir):
|
|
onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+'.opt')
|
|
os.makedirs(onnx_model_dir, exist_ok=True)
|
|
return os.path.join(onnx_model_dir, 'weights_map.json')
|
|
|
|
def getRefitNodesPath(self, model_name, onnx_dir, suffix=''):
|
|
onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+'.opt')
|
|
os.makedirs(onnx_model_dir, exist_ok=True)
|
|
return os.path.join(onnx_model_dir, 'refit'+suffix+'.json')
|
|
|
|
def getStateDictPath(self, model_name, onnx_dir, suffix=''):
|
|
onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+suffix)
|
|
os.makedirs(onnx_model_dir, exist_ok=True)
|
|
return os.path.join(onnx_model_dir, 'state_dict.pt')
|
|
|
|
def initializeModels(self, framework_model_dir, int8, fp8):
|
|
# Load text tokenizer(s)
|
|
if not self.pipeline_type.is_sd_xl_refiner():
|
|
self.tokenizer = make_tokenizer(self.version, self.pipeline_type, self.hf_token, framework_model_dir)
|
|
if self.pipeline_type.is_sd_xl():
|
|
self.tokenizer2 = make_tokenizer(self.version, self.pipeline_type, self.hf_token, framework_model_dir, subfolder='tokenizer_2')
|
|
|
|
# Load pipeline models
|
|
models_args = {'version': self.version, 'pipeline': self.pipeline_type, 'device': self.device,
|
|
'hf_token': self.hf_token, 'verbose': self.verbose, 'framework_model_dir': framework_model_dir,
|
|
'max_batch_size': self.max_batch_size}
|
|
|
|
if 'clip' in self.stages:
|
|
subfolder = 'text_encoder'
|
|
self.models['clip'] = CLIPModel(**models_args, fp16=True, embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type), output_hidden_states=self.config.get('clip_hidden_states', False), subfolder=subfolder)
|
|
|
|
if 'clip2' in self.stages:
|
|
subfolder = 'text_encoder_2'
|
|
self.models['clip2'] = CLIPWithProjModel(**models_args, fp16=True, output_hidden_states=self.config.get('clip_hidden_states', False), subfolder=subfolder)
|
|
|
|
if 'unet' in self.stages:
|
|
self.models['unet'] = UNetModel(**models_args, fp16=True, int8=int8, fp8=fp8, controlnets=self.controlnets, do_classifier_free_guidance=self.do_classifier_free_guidance)
|
|
|
|
if 'unetxl' in self.stages:
|
|
if not self.controlnets:
|
|
self.models["unetxl"] = UNetXLModel(
|
|
**models_args,
|
|
fp16=True,
|
|
int8=int8,
|
|
fp8=fp8,
|
|
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
|
)
|
|
else:
|
|
self.models["unetxl"] = UNetXLModelControlNet(
|
|
**models_args,
|
|
fp16=True,
|
|
int8=int8,
|
|
fp8=fp8,
|
|
controlnets=self.controlnets,
|
|
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
|
)
|
|
|
|
vae_fp16 = not self.pipeline_type.is_sd_xl()
|
|
|
|
if 'vae' in self.stages:
|
|
self.models['vae'] = VAEModel(**models_args, fp16=vae_fp16, tf32=True)
|
|
|
|
if 'vae_encoder' in self.stages:
|
|
self.models['vae_encoder'] = VAEEncoderModel(**models_args, fp16=vae_fp16)
|
|
|
|
def loadEngines(
|
|
self,
|
|
engine_dir,
|
|
framework_model_dir,
|
|
onnx_dir,
|
|
onnx_opset,
|
|
opt_batch_size,
|
|
opt_image_height,
|
|
opt_image_width,
|
|
optimization_level=3,
|
|
static_batch=False,
|
|
static_shape=True,
|
|
enable_refit=False,
|
|
enable_all_tactics=False,
|
|
timing_cache=None,
|
|
int8=False,
|
|
fp8=False,
|
|
quantization_level=2.5,
|
|
quantization_percentile=1.0,
|
|
quantization_alpha=0.8,
|
|
calibration_size=32,
|
|
calib_batch_size=2,
|
|
**_kwargs,
|
|
):
|
|
"""
|
|
Build and load engines for TensorRT accelerated inference.
|
|
Export ONNX models first, if applicable.
|
|
|
|
Args:
|
|
engine_dir (str):
|
|
Directory to store the TensorRT engines.
|
|
framework_model_dir (str):
|
|
Directory to store the framework model ckpt.
|
|
onnx_dir (str):
|
|
Directory to store the ONNX models.
|
|
onnx_opset (int):
|
|
ONNX opset version to export the models.
|
|
opt_batch_size (int):
|
|
Batch size to optimize for during engine building.
|
|
opt_image_height (int):
|
|
Image height to optimize for during engine building. Must be a multiple of 8.
|
|
opt_image_width (int):
|
|
Image width to optimize for during engine building. Must be a multiple of 8.
|
|
optimization_level (int):
|
|
Optimization level to build the TensorRT engine with.
|
|
static_batch (bool):
|
|
Build engine only for specified opt_batch_size.
|
|
static_shape (bool):
|
|
Build engine only for specified opt_image_height & opt_image_width. Default = True.
|
|
enable_refit (bool):
|
|
Build engines with refit option enabled.
|
|
enable_all_tactics (bool):
|
|
Enable all tactic sources during TensorRT engine builds.
|
|
timing_cache (str):
|
|
Path to the timing cache to speed up TensorRT build.
|
|
int8 (bool):
|
|
Whether to quantize to int8 format or not (SDXL, SD15 and SD21 only).
|
|
fp8 (bool):
|
|
Whether to quantize to fp8 format or not (SDXL, SD15 and SD21 only).
|
|
quantization_level (float):
|
|
Controls which layers to quantize. 1: CNN, 2: CNN+FFN, 2.5: CNN+FFN+QKV, 3: CNN+FC
|
|
quantization_percentile (float):
|
|
Control quantization scaling factors (amax) collecting range, where the minimum amax in
|
|
range(n_steps * percentile) will be collected. Recommendation: 1.0
|
|
quantization_alpha (float):
|
|
The alpha parameter for SmoothQuant quantization used for linear layers.
|
|
Recommendation: 0.8 for SDXL
|
|
calibration_size (int):
|
|
The number of steps to use for calibrating the model for quantization.
|
|
Recommendation: 32, 64, 128 for SDXL
|
|
calib_batch_size (int):
|
|
The batch size to use for calibration. Defaults to 2.
|
|
"""
|
|
# Create directories if missing
|
|
for directory in [engine_dir, onnx_dir]:
|
|
if not os.path.exists(directory):
|
|
print(f"[I] Create directory: {directory}")
|
|
pathlib.Path(directory).mkdir(parents=True)
|
|
|
|
# Initialize models
|
|
self.initializeModels(framework_model_dir, int8, fp8)
|
|
|
|
# Configure pipeline models to load
|
|
model_names = self.models.keys()
|
|
lora_suffix = '-'+'-'.join([str(md5(path.encode('utf-8')).hexdigest())+'-'+('%.2f' % self.lora_weights[path])+'-'+('%.2f' % self.lora_loader.scale) for path in sorted(self.lora_loader.paths)]) if self.lora_loader else ''
|
|
# Enable refit and LoRA merging only for UNet & UNetXL for now
|
|
do_engine_refit = dict(zip(model_names, [not self.pipeline_type.is_sd_xl_refiner() and enable_refit and model_name.startswith('unet') for model_name in model_names]))
|
|
do_lora_merge = dict(zip(model_names, [not enable_refit and self.lora_loader and model_name.startswith('unet') for model_name in model_names]))
|
|
# Torch fallback for VAE if specified
|
|
torch_fallback = dict(zip(model_names, [self.torch_inference for model_name in model_names]))
|
|
model_suffix = dict(zip(model_names, [lora_suffix if do_lora_merge[model_name] else '' for model_name in model_names]))
|
|
use_int8 = dict.fromkeys(model_names, False)
|
|
use_fp8 = dict.fromkeys(model_names, False)
|
|
if int8:
|
|
assert self.pipeline_type.is_sd_xl_base() or self.version == "1.4", "int8 quantization only supported for SDXL and SD1.4 pipeline"
|
|
model_name = 'unetxl' if self.pipeline_type.is_sd_xl() else 'unet'
|
|
use_int8[model_name] = True
|
|
model_suffix[model_name] += f"-int8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}.p{quantization_percentile}.a{quantization_alpha}"
|
|
elif fp8:
|
|
assert self.pipeline_type.is_sd_xl() or self.version == "1.4", "fp8 quantization only supported for SDXL and SD1.4 pipeline"
|
|
model_name = 'unetxl' if self.pipeline_type.is_sd_xl() else 'unet'
|
|
use_fp8[model_name] = True
|
|
model_suffix[model_name] += f"-fp8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}.p{quantization_percentile}.a{quantization_alpha}"
|
|
onnx_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, opt=False, suffix=model_suffix[model_name]) for model_name in model_names]))
|
|
onnx_opt_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, suffix=model_suffix[model_name]) for model_name in model_names]))
|
|
engine_path = dict(zip(model_names, [self.getEnginePath(model_name, engine_dir, do_engine_refit[model_name], suffix=model_suffix[model_name]) for model_name in model_names]))
|
|
weights_map_path = dict(zip(model_names, [(self.getWeightsMapPath(model_name, onnx_dir) if do_engine_refit[model_name] else None) for model_name in model_names]))
|
|
|
|
for model_name, obj in self.models.items():
|
|
if torch_fallback[model_name]:
|
|
continue
|
|
# Export models to ONNX and save weights name mapping
|
|
do_export_onnx = not os.path.exists(engine_path[model_name]) and not os.path.exists(onnx_opt_path[model_name])
|
|
do_export_weights_map = weights_map_path[model_name] and not os.path.exists(weights_map_path[model_name])
|
|
if do_export_onnx or do_export_weights_map:
|
|
# Non-quantized ONNX export
|
|
if not use_int8[model_name] and not use_fp8[model_name]:
|
|
obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width, enable_lora_merge=do_lora_merge[model_name], static_shape=static_shape, lora_loader=self.lora_loader)
|
|
else:
|
|
pipeline = obj.get_pipeline()
|
|
if self.pipeline_type.is_controlnet():
|
|
model = UNet2DConditionControlNetModel(pipeline.unet, pipeline.controlnet.nets)
|
|
else:
|
|
model = pipeline.unet
|
|
if use_fp8[model_name] and quantization_level == 4.0:
|
|
set_fmha(model)
|
|
|
|
state_dict_path = self.getStateDictPath(model_name, onnx_dir, suffix=model_suffix[model_name])
|
|
if not os.path.exists(state_dict_path):
|
|
print(f"[I] Calibrated weights not found, generating {state_dict_path}")
|
|
root_dir = os.path.dirname(os.path.abspath(sys.modules["__main__"].__file__))
|
|
calibration_file = os.path.join(root_dir, "calibration_data", "calibration-prompts.txt")
|
|
calibration_prompts = load_calib_prompts(calib_batch_size, calibration_file)
|
|
if self.pipeline_type.is_controlnet():
|
|
calibration_image_canny = image_module.download_image(
|
|
"https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0/resolve/main/out_bird.png"
|
|
)
|
|
# "out_bird.png" has 5 images combined in a row. We pick the first image which is the input image.
|
|
calibration_image_canny = calibration_image_canny.crop(
|
|
(0, 0, calibration_image_canny.width / 5, calibration_image_canny.height)
|
|
)
|
|
calibration_images = [calibration_image_canny]
|
|
|
|
# TODO check size > calibration_size
|
|
def do_calibrate(pipeline, calibration_prompts, **kwargs):
|
|
for i_th, prompts in enumerate(calibration_prompts):
|
|
if i_th >= kwargs["calib_size"]:
|
|
return
|
|
pipeline_call_kwargs = {
|
|
"prompt": prompts,
|
|
"num_inference_steps": kwargs["n_steps"],
|
|
"negative_prompt": [
|
|
"normal quality, low quality, worst quality, low res, blurry, nsfw, nude"
|
|
]
|
|
* len(prompts),
|
|
}
|
|
if self.pipeline_type.is_controlnet():
|
|
pipeline_call_kwargs["image"] = calibration_images
|
|
pipeline(**pipeline_call_kwargs).images
|
|
|
|
def forward_loop(model):
|
|
if self.pipeline_type.is_controlnet():
|
|
pipeline.unet = model.unet
|
|
pipeline.controlnet.nets = model.controlnets
|
|
else:
|
|
pipeline.unet = model
|
|
do_calibrate(
|
|
pipeline=pipeline,
|
|
calibration_prompts=calibration_prompts,
|
|
calib_size=calibration_size // calib_batch_size,
|
|
n_steps=self.denoising_steps,
|
|
)
|
|
|
|
print(f"[I] Performing calibration for {calibration_size} steps.")
|
|
if use_int8[model_name]:
|
|
quant_config = get_int8_config(
|
|
model,
|
|
quantization_level,
|
|
quantization_alpha,
|
|
quantization_percentile,
|
|
self.denoising_steps
|
|
)
|
|
elif use_fp8[model_name]:
|
|
quant_config = SD_FP8_FP16_DEFAULT_CONFIG
|
|
|
|
# Handle LoRA
|
|
if do_lora_merge[model_name]:
|
|
assert self.lora_loader is not None
|
|
model = merge_loras(model, self.lora_loader)
|
|
check_lora(model)
|
|
mtq.quantize(model, quant_config, forward_loop)
|
|
mto.save(model, state_dict_path)
|
|
else:
|
|
mto.restore(model, state_dict_path)
|
|
|
|
print(f"[I] Generating quantized ONNX model: {onnx_opt_path[model_name]}")
|
|
if not os.path.exists(onnx_path[model_name]):
|
|
quantize_lvl(self.version, model, quantization_level)
|
|
mtq.disable_quantizer(model, filter_func)
|
|
if use_fp8[model_name]:
|
|
generate_fp8_scales(model)
|
|
else:
|
|
model = None
|
|
obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width, custom_model=model, static_shape=static_shape)
|
|
|
|
# FIXME do_export_weights_map needs ONNX graph
|
|
if do_export_weights_map:
|
|
print(f"[I] Saving weights map: {weights_map_path[model_name]}")
|
|
obj.export_weights_map(onnx_opt_path[model_name], weights_map_path[model_name])
|
|
|
|
# Release temp GPU memory during onnx export to avoid OOM.
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Build TensorRT engines
|
|
for model_name, obj in self.models.items():
|
|
if torch_fallback[model_name]:
|
|
continue
|
|
engine = engine_module.Engine(engine_path[model_name])
|
|
if not os.path.exists(engine_path[model_name]):
|
|
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
|
|
# TF32 can be enabled for all precisions (including INT8/FP8)
|
|
tf32amp = obj.tf32
|
|
precision_constraints = 'none'
|
|
engine.build(onnx_opt_path[model_name],
|
|
tf32=tf32amp,
|
|
input_profile=obj.get_input_profile(
|
|
opt_batch_size, opt_image_height, opt_image_width,
|
|
static_batch=static_batch, static_shape=static_shape
|
|
),
|
|
enable_refit=do_engine_refit[model_name],
|
|
enable_all_tactics=enable_all_tactics,
|
|
timing_cache=timing_cache,
|
|
update_output_names=update_output_names,
|
|
verbose=self.verbose,
|
|
builder_optimization_level=optimization_level,
|
|
precision_constraints=precision_constraints,
|
|
)
|
|
self.engine[model_name] = engine
|
|
|
|
# Load TensorRT engines
|
|
for model_name, obj in self.models.items():
|
|
if torch_fallback[model_name]:
|
|
continue
|
|
self.engine[model_name].load()
|
|
if do_engine_refit[model_name] and self.lora_loader:
|
|
assert weights_map_path[model_name]
|
|
with open(weights_map_path[model_name], 'r') as fp_wts:
|
|
print(f"[I] Loading weights map: {weights_map_path[model_name]} ")
|
|
[weights_name_mapping, weights_shape_mapping] = json.load(fp_wts)
|
|
refit_weights_path = self.getRefitNodesPath(model_name, engine_dir, suffix=lora_suffix)
|
|
if not os.path.exists(refit_weights_path):
|
|
print(f"[I] Saving refit weights: {refit_weights_path}")
|
|
model = merge_loras(obj.get_model(), self.lora_loader)
|
|
refit_weights, updated_weight_names = engine_module.get_refit_weights(
|
|
model.state_dict(), onnx_opt_path[model_name], weights_name_mapping, weights_shape_mapping
|
|
)
|
|
torch.save((refit_weights, updated_weight_names), refit_weights_path)
|
|
unload_torch_model(model)
|
|
else:
|
|
print(f"[I] Loading refit weights: {refit_weights_path}")
|
|
refit_weights, updated_weight_names = torch.load(refit_weights_path)
|
|
self.engine[model_name].refit(refit_weights, updated_weight_names)
|
|
|
|
# Load torch models
|
|
for model_name, obj in self.models.items():
|
|
if torch_fallback[model_name]:
|
|
self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference)
|
|
|
|
# Release temp GPU memory during onnx export to avoid OOM.
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
def calculateMaxDeviceMemory(self):
|
|
max_device_memory = 0
|
|
for model_name, engine in self.engine.items():
|
|
max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
|
|
return max_device_memory
|
|
|
|
def activateEngines(self, shared_device_memory=None):
|
|
if shared_device_memory is None:
|
|
max_device_memory = self.calculateMaxDeviceMemory()
|
|
_, shared_device_memory = cudart.cudaMalloc(max_device_memory)
|
|
self.shared_device_memory = shared_device_memory
|
|
# Load and activate TensorRT engines
|
|
for engine in self.engine.values():
|
|
engine.activate(device_memory=self.shared_device_memory)
|
|
|
|
def runEngine(self, model_name, feed_dict):
|
|
engine = self.engine[model_name]
|
|
return engine.infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph)
|
|
|
|
def initialize_latents(self, batch_size, unet_channels, latent_height, latent_width, latents_dtype=torch.float32):
|
|
latents_dtype = latents_dtype # text_embeddings.dtype
|
|
latents_shape = (batch_size, unet_channels, latent_height, latent_width)
|
|
latents = torch.randn(latents_shape, device=self.device, dtype=latents_dtype, generator=self.generator)
|
|
# Scale the initial noise by the standard deviation required by the scheduler
|
|
latents = latents * self.scheduler.init_noise_sigma
|
|
return latents
|
|
|
|
def profile_start(self, name, color='blue'):
|
|
if self.nvtx_profile:
|
|
self.markers[name] = nvtx.start_range(message=name, color=color)
|
|
if name in self.events:
|
|
cudart.cudaEventRecord(self.events[name][0], 0)
|
|
|
|
def profile_stop(self, name):
|
|
if name in self.events:
|
|
cudart.cudaEventRecord(self.events[name][1], 0)
|
|
if self.nvtx_profile:
|
|
nvtx.end_range(self.markers[name])
|
|
|
|
def preprocess_images(self, batch_size, images=()):
|
|
if not images:
|
|
return ()
|
|
self.profile_start('preprocess', color='pink')
|
|
input_images=[]
|
|
for image in images:
|
|
image = image.to(self.device).float()
|
|
if image.shape[0] != batch_size:
|
|
image = image.repeat(batch_size, 1, 1, 1)
|
|
input_images.append(image)
|
|
self.profile_stop('preprocess')
|
|
return tuple(input_images)
|
|
|
|
def preprocess_controlnet_images(self, batch_size, images=None):
|
|
'''
|
|
images: List of PIL.Image.Image
|
|
'''
|
|
if images is None:
|
|
return None
|
|
self.profile_start('preprocess', color='pink')
|
|
images = [(np.array(i.convert("RGB")).astype(np.float32) / 255.0)[..., None].transpose(3, 2, 0, 1).repeat(batch_size, axis=0) for i in images]
|
|
# do_classifier_free_guidance
|
|
images = [torch.cat([torch.from_numpy(i).to(self.device).float()] * 2) for i in images]
|
|
images = torch.cat([image[None, ...] for image in images], dim=0)
|
|
self.profile_stop('preprocess')
|
|
return images
|
|
|
|
def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False):
|
|
self.profile_start(encoder, color='green')
|
|
|
|
tokenizer = self.tokenizer2 if encoder == 'clip2' else self.tokenizer
|
|
|
|
def tokenize(prompt, output_hidden_states):
|
|
text_input_ids = tokenizer(
|
|
prompt,
|
|
padding="max_length",
|
|
max_length=tokenizer.model_max_length,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
).input_ids.type(torch.int32).to(self.device)
|
|
|
|
text_hidden_states = None
|
|
if self.torch_inference:
|
|
outputs = self.torch_models[encoder](text_input_ids, output_hidden_states=output_hidden_states)
|
|
text_embeddings = outputs[0].clone()
|
|
if output_hidden_states:
|
|
text_hidden_states = outputs['hidden_states'][-2].clone()
|
|
else:
|
|
# NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt
|
|
outputs = self.runEngine(encoder, {'input_ids': text_input_ids})
|
|
text_embeddings = outputs['text_embeddings'].clone()
|
|
if output_hidden_states:
|
|
text_hidden_states = outputs['hidden_states'].clone()
|
|
return text_embeddings, text_hidden_states
|
|
|
|
# Tokenize prompt
|
|
text_embeddings, text_hidden_states = tokenize(prompt, output_hidden_states)
|
|
|
|
if self.do_classifier_free_guidance:
|
|
# Tokenize negative prompt
|
|
uncond_embeddings, uncond_hidden_states = tokenize(negative_prompt, output_hidden_states)
|
|
|
|
# Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance
|
|
text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)
|
|
|
|
if pooled_outputs:
|
|
pooled_output = text_embeddings
|
|
|
|
if output_hidden_states:
|
|
text_embeddings = torch.cat([uncond_hidden_states, text_hidden_states]).to(dtype=torch.float16) if self.do_classifier_free_guidance else text_hidden_states
|
|
|
|
self.profile_stop(encoder)
|
|
if pooled_outputs:
|
|
return text_embeddings, pooled_output
|
|
return text_embeddings
|
|
|
|
# from diffusers (get_timesteps)
|
|
def get_timesteps(self, num_inference_steps, strength, denoising_start=None):
|
|
# get the original timestep using init_timestep
|
|
if denoising_start is None:
|
|
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
|
|
t_start = max(num_inference_steps - init_timestep, 0)
|
|
else:
|
|
t_start = 0
|
|
|
|
timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
|
|
|
|
# Strength is irrelevant if we directly request a timestep to start at;
|
|
# that is, strength is determined by the denoising_start instead.
|
|
if denoising_start is not None:
|
|
discrete_timestep_cutoff = int(
|
|
round(
|
|
self.scheduler.config.num_train_timesteps
|
|
- (denoising_start * self.scheduler.config.num_train_timesteps)
|
|
)
|
|
)
|
|
|
|
num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
|
|
if self.scheduler.order == 2 and num_inference_steps % 2 == 0:
|
|
# if the scheduler is a 2nd order scheduler we might have to do +1
|
|
# because `num_inference_steps` might be even given that every timestep
|
|
# (except the highest one) is duplicated. If `num_inference_steps` is even it would
|
|
# mean that we cut the timesteps in the middle of the denoising step
|
|
# (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
|
|
# we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
|
|
num_inference_steps = num_inference_steps + 1
|
|
|
|
# because t_n+1 >= t_n, we slice the timesteps starting from the end
|
|
timesteps = timesteps[-num_inference_steps:]
|
|
return timesteps, num_inference_steps
|
|
|
|
return timesteps, num_inference_steps - t_start
|
|
|
|
def denoise_latent(self,
|
|
latents,
|
|
text_embeddings,
|
|
denoiser='unet',
|
|
timesteps=None,
|
|
step_offset=0,
|
|
mask=None,
|
|
masked_image_latents=None,
|
|
image_guidance=1.5,
|
|
controlnet_imgs=None,
|
|
controlnet_scales=None,
|
|
text_embeds=None,
|
|
time_ids=None):
|
|
|
|
assert image_guidance > 1.0, "Image guidance has to be > 1.0"
|
|
|
|
controlnet_imgs = self.preprocess_controlnet_images(latents.shape[0], controlnet_imgs)
|
|
|
|
do_autocast = self.torch_inference != '' and self.models[denoiser].fp16
|
|
with torch.autocast('cuda', enabled=do_autocast):
|
|
self.profile_start(denoiser, color='blue')
|
|
for step_index, timestep in enumerate(timesteps):
|
|
# Expand the latents if we are doing classifier free guidance
|
|
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
|
|
latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)
|
|
if isinstance(mask, torch.Tensor):
|
|
latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)
|
|
|
|
# Predict the noise residual
|
|
if self.torch_inference:
|
|
params = {"sample": latent_model_input, "timestep": timestep, "encoder_hidden_states": text_embeddings}
|
|
if controlnet_imgs is not None:
|
|
params.update({"images": controlnet_imgs, "controlnet_scales": controlnet_scales})
|
|
added_cond_kwargs = {}
|
|
if text_embeds != None:
|
|
added_cond_kwargs.update({'text_embeds': text_embeds})
|
|
if time_ids != None:
|
|
added_cond_kwargs.update({'time_ids': time_ids})
|
|
if text_embeds != None or time_ids != None:
|
|
params.update({'added_cond_kwargs': added_cond_kwargs})
|
|
noise_pred = self.torch_models[denoiser](**params)["sample"]
|
|
else:
|
|
timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep
|
|
|
|
params = {"sample": latent_model_input, "timestep": timestep_float, "encoder_hidden_states": text_embeddings}
|
|
if controlnet_imgs is not None:
|
|
params.update({"images": controlnet_imgs, "controlnet_scales": controlnet_scales})
|
|
if text_embeds != None:
|
|
params.update({'text_embeds': text_embeds})
|
|
if time_ids != None:
|
|
params.update({'time_ids': time_ids})
|
|
noise_pred = self.runEngine(denoiser, params)['latent']
|
|
|
|
# Perform guidance
|
|
if self.do_classifier_free_guidance:
|
|
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
|
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
|
|
|
|
# from diffusers (prepare_extra_step_kwargs)
|
|
extra_step_kwargs = {}
|
|
if "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()):
|
|
# TODO: configurable eta
|
|
eta = 0.0
|
|
extra_step_kwargs["eta"] = eta
|
|
if "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()):
|
|
extra_step_kwargs["generator"] = self.generator
|
|
|
|
latents = self.scheduler.step(noise_pred, timestep, latents, **extra_step_kwargs, return_dict=False)[0]
|
|
|
|
latents = 1. / self.vae_scaling_factor * latents
|
|
latents = latents.to(dtype=torch.float32)
|
|
|
|
self.profile_stop(denoiser)
|
|
return latents
|
|
|
|
def encode_image(self, input_image):
|
|
self.profile_start('vae_encoder', color='red')
|
|
cast_to = torch.float16 if self.models['vae_encoder'].fp16 else torch.bfloat16 if self.models['vae_encoder'].bf16 else torch.float32
|
|
input_image = input_image.to(dtype=cast_to)
|
|
if self.torch_inference:
|
|
image_latents = self.torch_models['vae_encoder'](input_image)
|
|
else:
|
|
image_latents = self.runEngine('vae_encoder', {'images': input_image})['latent']
|
|
image_latents = self.vae_scaling_factor * image_latents
|
|
self.profile_stop('vae_encoder')
|
|
return image_latents
|
|
|
|
def decode_latent(self, latents):
|
|
self.profile_start('vae', color='red')
|
|
cast_to = torch.float16 if self.models['vae'].fp16 else torch.bfloat16 if self.models['vae'].bf16 else torch.float32
|
|
latents = latents.to(dtype=cast_to)
|
|
|
|
if self.torch_inference:
|
|
images = self.torch_models['vae'](latents, return_dict=False)[0]
|
|
else:
|
|
images = self.runEngine('vae', {'latent': latents})['images']
|
|
self.profile_stop('vae')
|
|
return images
|
|
|
|
def print_summary(self, denoising_steps, walltime_ms, batch_size):
|
|
print('|-----------------|--------------|')
|
|
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
|
|
print('|-----------------|--------------|')
|
|
for stage in self.stages:
|
|
stage_name = stage
|
|
if "unet" in stage:
|
|
if self.pipeline_type.is_controlnet():
|
|
stage_name += '+cnet'
|
|
stage_name += ' x ' + str(denoising_steps)
|
|
print(
|
|
"| {:^15} | {:>9.2f} ms |".format(
|
|
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
|
|
)
|
|
)
|
|
print('|-----------------|--------------|')
|
|
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
|
|
print('|-----------------|--------------|')
|
|
print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms))
|
|
|
|
def save_image(self, images, pipeline, prompt, seed):
|
|
# Save image
|
|
image_name_prefix = pipeline+''.join(set(['-'+prompt[i].replace(' ','_')[:10] for i in range(len(prompt))]))+'-'+str(seed)+'-'
|
|
image_name_suffix = 'torch' if self.torch_inference else 'trt'
|
|
image_module.save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
|
|
|
|
def infer(
|
|
self,
|
|
prompt,
|
|
negative_prompt,
|
|
image_height,
|
|
image_width,
|
|
input_image=None,
|
|
image_strength=0.75,
|
|
controlnet_scales=None,
|
|
aesthetic_score=6.0,
|
|
negative_aesthetic_score=2.5,
|
|
warmup=False,
|
|
verbose=False,
|
|
save_image=True,
|
|
):
|
|
"""
|
|
Run the diffusion pipeline.
|
|
|
|
Args:
|
|
prompt (str):
|
|
The text prompt to guide image generation.
|
|
negative_prompt (str):
|
|
The prompt not to guide the image generation.
|
|
image_height (int):
|
|
Height (in pixels) of the image to be generated. Must be a multiple of 8.
|
|
image_width (int):
|
|
Width (in pixels) of the image to be generated. Must be a multiple of 8.
|
|
input_image (image):
|
|
Input image used to initialize the latents.
|
|
image_strength (float):
|
|
Strength of transformation applied to input_image. Must be between 0 and 1.
|
|
controlnet_scales (torch.Tensor)
|
|
A tensor which containes ControlNet scales, essential for multi ControlNet.
|
|
Must be equal to number of Controlnets.
|
|
warmup (bool):
|
|
Indicate if this is a warmup run.
|
|
verbose (bool):
|
|
Verbose in logging
|
|
save_image (bool):
|
|
Save the generated image (if applicable)
|
|
"""
|
|
assert len(prompt) == len(negative_prompt)
|
|
batch_size = len(prompt)
|
|
|
|
# Spatial dimensions of latent tensor
|
|
latent_height = image_height // 8
|
|
latent_width = image_width // 8
|
|
|
|
if self.generator and self.seed:
|
|
self.generator.manual_seed(self.seed)
|
|
|
|
num_inference_steps = self.denoising_steps
|
|
|
|
with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
|
|
torch.cuda.synchronize()
|
|
e2e_tic = time.perf_counter()
|
|
|
|
# TODO: support custom timesteps
|
|
timesteps = None
|
|
if timesteps is not None:
|
|
if "timesteps" not in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()):
|
|
raise ValueError(
|
|
f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom"
|
|
f" timestep schedules. Please check whether you are using the correct scheduler."
|
|
)
|
|
self.scheduler.set_timesteps(timesteps=timesteps, device=self.device)
|
|
assert self.denoising_steps == len(self.scheduler.timesteps)
|
|
else:
|
|
self.scheduler.set_timesteps(self.denoising_steps, device=self.device)
|
|
timesteps = self.scheduler.timesteps.to(self.device)
|
|
|
|
denoise_kwargs = {}
|
|
if not (self.pipeline_type.is_img2img() or self.pipeline_type.is_sd_xl_refiner()):
|
|
# Initialize latents
|
|
latents = self.initialize_latents(batch_size=batch_size,
|
|
unet_channels=4,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width)
|
|
if self.pipeline_type.is_controlnet():
|
|
denoise_kwargs.update({'controlnet_imgs': input_image, 'controlnet_scales': controlnet_scales})
|
|
|
|
# Pre-process and VAE encode input image
|
|
if self.pipeline_type.is_img2img() or self.pipeline_type.is_sd_xl_refiner():
|
|
assert input_image != None
|
|
# Initialize timesteps and pre-process input image
|
|
timesteps, num_inference_steps = self.get_timesteps(self.denoising_steps, image_strength)
|
|
denoise_kwargs.update({'timesteps': timesteps})
|
|
if self.pipeline_type.is_img2img() or self.pipeline_type.is_sd_xl_refiner():
|
|
latent_timestep = timesteps[:1].repeat(batch_size)
|
|
input_image = self.preprocess_images(batch_size, (input_image,))[0]
|
|
# Encode if not a latent
|
|
image_latents = input_image if input_image.shape[1] == 4 else self.encode_image(input_image)
|
|
# Add noise to latents using timesteps
|
|
noise = torch.randn(image_latents.shape, generator=self.generator, device=self.device, dtype=torch.float32)
|
|
latents = self.scheduler.add_noise(image_latents, noise, latent_timestep)
|
|
|
|
# CLIP text encoder(s)
|
|
if self.pipeline_type.is_sd_xl():
|
|
text_embeddings2, pooled_embeddings2 = self.encode_prompt(prompt, negative_prompt,
|
|
encoder='clip2', pooled_outputs=True, output_hidden_states=True)
|
|
|
|
# Merge text embeddings
|
|
if self.pipeline_type.is_sd_xl_base():
|
|
text_embeddings = self.encode_prompt(prompt, negative_prompt, output_hidden_states=True)
|
|
text_embeddings = torch.cat([text_embeddings, text_embeddings2], dim=-1)
|
|
else:
|
|
text_embeddings = text_embeddings2
|
|
|
|
# Time embeddings
|
|
def _get_add_time_ids(original_size, crops_coords_top_left, target_size, dtype, aesthetic_score=None, negative_aesthetic_score=None):
|
|
if self.pipeline_type.is_sd_xl_refiner(): #self.requires_aesthetics_score:
|
|
add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
|
|
if self.do_classifier_free_guidance:
|
|
add_neg_time_ids = list(original_size + crops_coords_top_left + (negative_aesthetic_score,))
|
|
else:
|
|
add_time_ids = list(original_size + crops_coords_top_left + target_size)
|
|
if self.do_classifier_free_guidance:
|
|
add_neg_time_ids = list(original_size + crops_coords_top_left + target_size)
|
|
add_time_ids = torch.tensor([add_time_ids], dtype=dtype, device=self.device)
|
|
if self.do_classifier_free_guidance:
|
|
add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype, device=self.device)
|
|
add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0)
|
|
return add_time_ids
|
|
|
|
original_size = (image_height, image_width)
|
|
crops_coords_top_left = (0, 0)
|
|
target_size = (image_height, image_width)
|
|
if self.pipeline_type.is_sd_xl_refiner():
|
|
add_time_ids = _get_add_time_ids(
|
|
original_size, crops_coords_top_left, target_size, dtype=text_embeddings.dtype, aesthetic_score=aesthetic_score, negative_aesthetic_score=negative_aesthetic_score
|
|
)
|
|
else:
|
|
add_time_ids = _get_add_time_ids(
|
|
original_size, crops_coords_top_left, target_size, dtype=text_embeddings.dtype
|
|
)
|
|
add_time_ids = add_time_ids.repeat(batch_size, 1)
|
|
denoise_kwargs.update({'text_embeds': pooled_embeddings2, 'time_ids': add_time_ids})
|
|
else:
|
|
text_embeddings = self.encode_prompt(prompt, negative_prompt)
|
|
|
|
# UNet denoiser + (optional) ControlNet(s)
|
|
denoiser = 'unetxl' if self.pipeline_type.is_sd_xl() else 'unet'
|
|
latents = self.denoise_latent(latents, text_embeddings, denoiser=denoiser, **denoise_kwargs)
|
|
|
|
# VAE decode latent (if applicable)
|
|
if self.return_latents:
|
|
latents = latents * self.vae_scaling_factor
|
|
else:
|
|
images = self.decode_latent(latents)
|
|
|
|
torch.cuda.synchronize()
|
|
e2e_toc = time.perf_counter()
|
|
|
|
walltime_ms = (e2e_toc - e2e_tic) * 1000.
|
|
if not warmup:
|
|
self.print_summary(num_inference_steps, walltime_ms, batch_size)
|
|
if not self.return_latents and save_image:
|
|
# post-process images
|
|
images = ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy()
|
|
self.save_image(images, self.pipeline_type.name.lower(), prompt, self.seed)
|
|
|
|
return (latents, walltime_ms) if self.return_latents else (images, walltime_ms)
|
|
|
|
def run(self, prompt, negative_prompt, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph, **kwargs):
|
|
# Process prompt
|
|
if not isinstance(prompt, list):
|
|
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
|
|
prompt = prompt * batch_size
|
|
|
|
if not isinstance(negative_prompt, list):
|
|
raise ValueError(f"`--negative-prompt` must be of type `str` list, but is {type(negative_prompt)}")
|
|
if len(negative_prompt) == 1:
|
|
negative_prompt = negative_prompt * batch_size
|
|
|
|
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
|
|
if num_warmup_runs > 0:
|
|
print("[I] Warming up ..")
|
|
for _ in range(num_warmup_runs):
|
|
self.infer(prompt, negative_prompt, height, width, warmup=True, **kwargs)
|
|
|
|
for _ in range(batch_count):
|
|
print("[I] Running StableDiffusion pipeline")
|
|
if self.nvtx_profile:
|
|
cudart.cudaProfilerStart()
|
|
self.infer(prompt, negative_prompt, height, width, warmup=False, **kwargs)
|
|
if self.nvtx_profile:
|
|
cudart.cudaProfilerStop()
|