diffusers/community-pipelines-mirror

1#2# Copyright 2024 The HuggingFace Inc. team.3# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.4# SPDX-License-Identifier: Apache-2.05#6# Licensed under the Apache License, Version 2.0 (the "License");7# you may not use this file except in compliance with the License.8# You may obtain a copy of the License at9#10# http://www.apache.org/licenses/LICENSE-2.011#12# Unless required by applicable law or agreed to in writing, software13# distributed under the License is distributed on an "AS IS" BASIS,14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15# See the License for the specific language governing permissions and16# limitations under the License.17 18import gc19import os20from collections import OrderedDict21from copy import copy22from typing import List, Optional, Union23 24import numpy as np25import onnx26import onnx_graphsurgeon as gs27import PIL.Image28import tensorrt as trt29import torch30from huggingface_hub import snapshot_download31from huggingface_hub.utils import validate_hf_hub_args32from onnx import shape_inference33from polygraphy import cuda34from polygraphy.backend.common import bytes_from_path35from polygraphy.backend.onnx.loader import fold_constants36from polygraphy.backend.trt import (37    CreateConfig,38    Profile,39    engine_from_bytes,40    engine_from_network,41    network_from_onnx_path,42    save_engine,43)44from polygraphy.backend.trt import util as trt_util45from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection46 47from diffusers.models import AutoencoderKL, UNet2DConditionModel48from diffusers.pipelines.stable_diffusion import (49    StableDiffusionInpaintPipeline,50    StableDiffusionPipelineOutput,51    StableDiffusionSafetyChecker,52)53from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint import prepare_mask_and_masked_image54from diffusers.schedulers import DDIMScheduler55from diffusers.utils import logging56 57 58"""59Installation instructions60python3 -m pip install --upgrade transformers diffusers>=0.16.061python3 -m pip install --upgrade tensorrt>=8.6.162python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com63python3 -m pip install onnxruntime64"""65 66TRT_LOGGER = trt.Logger(trt.Logger.ERROR)67logger = logging.get_logger(__name__)  # pylint: disable=invalid-name68 69# Map of numpy dtype -> torch dtype70numpy_to_torch_dtype_dict = {71    np.uint8: torch.uint8,72    np.int8: torch.int8,73    np.int16: torch.int16,74    np.int32: torch.int32,75    np.int64: torch.int64,76    np.float16: torch.float16,77    np.float32: torch.float32,78    np.float64: torch.float64,79    np.complex64: torch.complex64,80    np.complex128: torch.complex128,81}82if np.version.full_version >= "1.24.0":83    numpy_to_torch_dtype_dict[np.bool_] = torch.bool84else:85    numpy_to_torch_dtype_dict[np.bool] = torch.bool86 87# Map of torch dtype -> numpy dtype88torch_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()}89 90 91def device_view(t):92    return cuda.DeviceView(ptr=t.data_ptr(), shape=t.shape, dtype=torch_to_numpy_dtype_dict[t.dtype])93 94 95def preprocess_image(image):96    """97    image: torch.Tensor98    """99    w, h = image.size100    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32101    image = image.resize((w, h))102    image = np.array(image).astype(np.float32) / 255.0103    image = image[None].transpose(0, 3, 1, 2)104    image = torch.from_numpy(image).contiguous()105    return 2.0 * image - 1.0106 107 108class Engine:109    def __init__(self, engine_path):110        self.engine_path = engine_path111        self.engine = None112        self.context = None113        self.buffers = OrderedDict()114        self.tensors = OrderedDict()115 116    def __del__(self):117        [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray)]118        del self.engine119        del self.context120        del self.buffers121        del self.tensors122 123    def build(124        self,125        onnx_path,126        fp16,127        input_profile=None,128        enable_preview=False,129        enable_all_tactics=False,130        timing_cache=None,131        workspace_size=0,132    ):133        logger.warning(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")134        p = Profile()135        if input_profile:136            for name, dims in input_profile.items():137                assert len(dims) == 3138                p.add(name, min=dims[0], opt=dims[1], max=dims[2])139 140        config_kwargs = {}141 142        config_kwargs["preview_features"] = [trt.PreviewFeature.DISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805]143        if enable_preview:144            # Faster dynamic shapes made optional since it increases engine build time.145            config_kwargs["preview_features"].append(trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805)146        if workspace_size > 0:147            config_kwargs["memory_pool_limits"] = {trt.MemoryPoolType.WORKSPACE: workspace_size}148        if not enable_all_tactics:149            config_kwargs["tactic_sources"] = []150 151        engine = engine_from_network(152            network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]),153            config=CreateConfig(fp16=fp16, profiles=[p], load_timing_cache=timing_cache, **config_kwargs),154            save_timing_cache=timing_cache,155        )156        save_engine(engine, path=self.engine_path)157 158    def load(self):159        logger.warning(f"Loading TensorRT engine: {self.engine_path}")160        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))161 162    def activate(self):163        self.context = self.engine.create_execution_context()164 165    def allocate_buffers(self, shape_dict=None, device="cuda"):166        for idx in range(trt_util.get_bindings_per_profile(self.engine)):167            binding = self.engine[idx]168            if shape_dict and binding in shape_dict:169                shape = shape_dict[binding]170            else:171                shape = self.engine.get_binding_shape(binding)172            dtype = trt.nptype(self.engine.get_binding_dtype(binding))173            if self.engine.binding_is_input(binding):174                self.context.set_binding_shape(idx, shape)175            tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device)176            self.tensors[binding] = tensor177            self.buffers[binding] = cuda.DeviceView(ptr=tensor.data_ptr(), shape=shape, dtype=dtype)178 179    def infer(self, feed_dict, stream):180        start_binding, end_binding = trt_util.get_active_profile_bindings(self.context)181        # shallow copy of ordered dict182        device_buffers = copy(self.buffers)183        for name, buf in feed_dict.items():184            assert isinstance(buf, cuda.DeviceView)185            device_buffers[name] = buf186        bindings = [0] * start_binding + [buf.ptr for buf in device_buffers.values()]187        noerror = self.context.execute_async_v2(bindings=bindings, stream_handle=stream.ptr)188        if not noerror:189            raise ValueError("ERROR: inference failed.")190 191        return self.tensors192 193 194class Optimizer:195    def __init__(self, onnx_graph):196        self.graph = gs.import_onnx(onnx_graph)197 198    def cleanup(self, return_onnx=False):199        self.graph.cleanup().toposort()200        if return_onnx:201            return gs.export_onnx(self.graph)202 203    def select_outputs(self, keep, names=None):204        self.graph.outputs = [self.graph.outputs[o] for o in keep]205        if names:206            for i, name in enumerate(names):207                self.graph.outputs[i].name = name208 209    def fold_constants(self, return_onnx=False):210        onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)211        self.graph = gs.import_onnx(onnx_graph)212        if return_onnx:213            return onnx_graph214 215    def infer_shapes(self, return_onnx=False):216        onnx_graph = gs.export_onnx(self.graph)217        if onnx_graph.ByteSize() > 2147483648:218            raise TypeError("ERROR: model size exceeds supported 2GB limit")219        else:220            onnx_graph = shape_inference.infer_shapes(onnx_graph)221 222        self.graph = gs.import_onnx(onnx_graph)223        if return_onnx:224            return onnx_graph225 226 227class BaseModel:228    def __init__(self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77):229        self.model = model230        self.name = "SD Model"231        self.fp16 = fp16232        self.device = device233 234        self.min_batch = 1235        self.max_batch = max_batch_size236        self.min_image_shape = 256  # min image resolution: 256x256237        self.max_image_shape = 1024  # max image resolution: 1024x1024238        self.min_latent_shape = self.min_image_shape // 8239        self.max_latent_shape = self.max_image_shape // 8240 241        self.embedding_dim = embedding_dim242        self.text_maxlen = text_maxlen243 244    def get_model(self):245        return self.model246 247    def get_input_names(self):248        pass249 250    def get_output_names(self):251        pass252 253    def get_dynamic_axes(self):254        return None255 256    def get_sample_input(self, batch_size, image_height, image_width):257        pass258 259    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):260        return None261 262    def get_shape_dict(self, batch_size, image_height, image_width):263        return None264 265    def optimize(self, onnx_graph):266        opt = Optimizer(onnx_graph)267        opt.cleanup()268        opt.fold_constants()269        opt.infer_shapes()270        onnx_opt_graph = opt.cleanup(return_onnx=True)271        return onnx_opt_graph272 273    def check_dims(self, batch_size, image_height, image_width):274        assert batch_size >= self.min_batch and batch_size <= self.max_batch275        assert image_height % 8 == 0 or image_width % 8 == 0276        latent_height = image_height // 8277        latent_width = image_width // 8278        assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape279        assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape280        return (latent_height, latent_width)281 282    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):283        min_batch = batch_size if static_batch else self.min_batch284        max_batch = batch_size if static_batch else self.max_batch285        latent_height = image_height // 8286        latent_width = image_width // 8287        min_image_height = image_height if static_shape else self.min_image_shape288        max_image_height = image_height if static_shape else self.max_image_shape289        min_image_width = image_width if static_shape else self.min_image_shape290        max_image_width = image_width if static_shape else self.max_image_shape291        min_latent_height = latent_height if static_shape else self.min_latent_shape292        max_latent_height = latent_height if static_shape else self.max_latent_shape293        min_latent_width = latent_width if static_shape else self.min_latent_shape294        max_latent_width = latent_width if static_shape else self.max_latent_shape295        return (296            min_batch,297            max_batch,298            min_image_height,299            max_image_height,300            min_image_width,301            max_image_width,302            min_latent_height,303            max_latent_height,304            min_latent_width,305            max_latent_width,306        )307 308 309def getOnnxPath(model_name, onnx_dir, opt=True):310    return os.path.join(onnx_dir, model_name + (".opt" if opt else "") + ".onnx")311 312 313def getEnginePath(model_name, engine_dir):314    return os.path.join(engine_dir, model_name + ".plan")315 316 317def build_engines(318    models: dict,319    engine_dir,320    onnx_dir,321    onnx_opset,322    opt_image_height,323    opt_image_width,324    opt_batch_size=1,325    force_engine_rebuild=False,326    static_batch=False,327    static_shape=True,328    enable_preview=False,329    enable_all_tactics=False,330    timing_cache=None,331    max_workspace_size=0,332):333    built_engines = {}334    if not os.path.isdir(onnx_dir):335        os.makedirs(onnx_dir)336    if not os.path.isdir(engine_dir):337        os.makedirs(engine_dir)338 339    # Export models to ONNX340    for model_name, model_obj in models.items():341        engine_path = getEnginePath(model_name, engine_dir)342        if force_engine_rebuild or not os.path.exists(engine_path):343            logger.warning("Building Engines...")344            logger.warning("Engine build can take a while to complete")345            onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)346            onnx_opt_path = getOnnxPath(model_name, onnx_dir)347            if force_engine_rebuild or not os.path.exists(onnx_opt_path):348                if force_engine_rebuild or not os.path.exists(onnx_path):349                    logger.warning(f"Exporting model: {onnx_path}")350                    model = model_obj.get_model()351                    with torch.inference_mode(), torch.autocast("cuda"):352                        inputs = model_obj.get_sample_input(opt_batch_size, opt_image_height, opt_image_width)353                        torch.onnx.export(354                            model,355                            inputs,356                            onnx_path,357                            export_params=True,358                            opset_version=onnx_opset,359                            do_constant_folding=True,360                            input_names=model_obj.get_input_names(),361                            output_names=model_obj.get_output_names(),362                            dynamic_axes=model_obj.get_dynamic_axes(),363                        )364                    del model365                    torch.cuda.empty_cache()366                    gc.collect()367                else:368                    logger.warning(f"Found cached model: {onnx_path}")369 370                # Optimize onnx371                if force_engine_rebuild or not os.path.exists(onnx_opt_path):372                    logger.warning(f"Generating optimizing model: {onnx_opt_path}")373                    onnx_opt_graph = model_obj.optimize(onnx.load(onnx_path))374                    onnx.save(onnx_opt_graph, onnx_opt_path)375                else:376                    logger.warning(f"Found cached optimized model: {onnx_opt_path} ")377 378    # Build TensorRT engines379    for model_name, model_obj in models.items():380        engine_path = getEnginePath(model_name, engine_dir)381        engine = Engine(engine_path)382        onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)383        onnx_opt_path = getOnnxPath(model_name, onnx_dir)384 385        if force_engine_rebuild or not os.path.exists(engine.engine_path):386            engine.build(387                onnx_opt_path,388                fp16=True,389                input_profile=model_obj.get_input_profile(390                    opt_batch_size,391                    opt_image_height,392                    opt_image_width,393                    static_batch=static_batch,394                    static_shape=static_shape,395                ),396                enable_preview=enable_preview,397                timing_cache=timing_cache,398                workspace_size=max_workspace_size,399            )400        built_engines[model_name] = engine401 402    # Load and activate TensorRT engines403    for model_name, model_obj in models.items():404        engine = built_engines[model_name]405        engine.load()406        engine.activate()407 408    return built_engines409 410 411def runEngine(engine, feed_dict, stream):412    return engine.infer(feed_dict, stream)413 414 415class CLIP(BaseModel):416    def __init__(self, model, device, max_batch_size, embedding_dim):417        super(CLIP, self).__init__(418            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim419        )420        self.name = "CLIP"421 422    def get_input_names(self):423        return ["input_ids"]424 425    def get_output_names(self):426        return ["text_embeddings", "pooler_output"]427 428    def get_dynamic_axes(self):429        return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}430 431    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):432        self.check_dims(batch_size, image_height, image_width)433        min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(434            batch_size, image_height, image_width, static_batch, static_shape435        )436        return {437            "input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]438        }439 440    def get_shape_dict(self, batch_size, image_height, image_width):441        self.check_dims(batch_size, image_height, image_width)442        return {443            "input_ids": (batch_size, self.text_maxlen),444            "text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),445        }446 447    def get_sample_input(self, batch_size, image_height, image_width):448        self.check_dims(batch_size, image_height, image_width)449        return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)450 451    def optimize(self, onnx_graph):452        opt = Optimizer(onnx_graph)453        opt.select_outputs([0])  # delete graph output#1454        opt.cleanup()455        opt.fold_constants()456        opt.infer_shapes()457        opt.select_outputs([0], names=["text_embeddings"])  # rename network output458        opt_onnx_graph = opt.cleanup(return_onnx=True)459        return opt_onnx_graph460 461 462def make_CLIP(model, device, max_batch_size, embedding_dim, inpaint=False):463    return CLIP(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)464 465 466class UNet(BaseModel):467    def __init__(468        self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77, unet_dim=4469    ):470        super(UNet, self).__init__(471            model=model,472            fp16=fp16,473            device=device,474            max_batch_size=max_batch_size,475            embedding_dim=embedding_dim,476            text_maxlen=text_maxlen,477        )478        self.unet_dim = unet_dim479        self.name = "UNet"480 481    def get_input_names(self):482        return ["sample", "timestep", "encoder_hidden_states"]483 484    def get_output_names(self):485        return ["latent"]486 487    def get_dynamic_axes(self):488        return {489            "sample": {0: "2B", 2: "H", 3: "W"},490            "encoder_hidden_states": {0: "2B"},491            "latent": {0: "2B", 2: "H", 3: "W"},492        }493 494    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):495        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)496        (497            min_batch,498            max_batch,499            _,500            _,501            _,502            _,503            min_latent_height,504            max_latent_height,505            min_latent_width,506            max_latent_width,507        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)508        return {509            "sample": [510                (2 * min_batch, self.unet_dim, min_latent_height, min_latent_width),511                (2 * batch_size, self.unet_dim, latent_height, latent_width),512                (2 * max_batch, self.unet_dim, max_latent_height, max_latent_width),513            ],514            "encoder_hidden_states": [515                (2 * min_batch, self.text_maxlen, self.embedding_dim),516                (2 * batch_size, self.text_maxlen, self.embedding_dim),517                (2 * max_batch, self.text_maxlen, self.embedding_dim),518            ],519        }520 521    def get_shape_dict(self, batch_size, image_height, image_width):522        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)523        return {524            "sample": (2 * batch_size, self.unet_dim, latent_height, latent_width),525            "encoder_hidden_states": (2 * batch_size, self.text_maxlen, self.embedding_dim),526            "latent": (2 * batch_size, 4, latent_height, latent_width),527        }528 529    def get_sample_input(self, batch_size, image_height, image_width):530        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)531        dtype = torch.float16 if self.fp16 else torch.float32532        return (533            torch.randn(534                2 * batch_size, self.unet_dim, latent_height, latent_width, dtype=torch.float32, device=self.device535            ),536            torch.tensor([1.0], dtype=torch.float32, device=self.device),537            torch.randn(2 * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),538        )539 540 541def make_UNet(model, device, max_batch_size, embedding_dim, inpaint=False, unet_dim=4):542    return UNet(543        model,544        fp16=True,545        device=device,546        max_batch_size=max_batch_size,547        embedding_dim=embedding_dim,548        unet_dim=unet_dim,549    )550 551 552class VAE(BaseModel):553    def __init__(self, model, device, max_batch_size, embedding_dim):554        super(VAE, self).__init__(555            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim556        )557        self.name = "VAE decoder"558 559    def get_input_names(self):560        return ["latent"]561 562    def get_output_names(self):563        return ["images"]564 565    def get_dynamic_axes(self):566        return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}567 568    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):569        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)570        (571            min_batch,572            max_batch,573            _,574            _,575            _,576            _,577            min_latent_height,578            max_latent_height,579            min_latent_width,580            max_latent_width,581        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)582        return {583            "latent": [584                (min_batch, 4, min_latent_height, min_latent_width),585                (batch_size, 4, latent_height, latent_width),586                (max_batch, 4, max_latent_height, max_latent_width),587            ]588        }589 590    def get_shape_dict(self, batch_size, image_height, image_width):591        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)592        return {593            "latent": (batch_size, 4, latent_height, latent_width),594            "images": (batch_size, 3, image_height, image_width),595        }596 597    def get_sample_input(self, batch_size, image_height, image_width):598        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)599        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)600 601 602def make_VAE(model, device, max_batch_size, embedding_dim, inpaint=False):603    return VAE(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)604 605 606class TorchVAEEncoder(torch.nn.Module):607    def __init__(self, model):608        super().__init__()609        self.vae_encoder = model610 611    def forward(self, x):612        return self.vae_encoder.encode(x).latent_dist.sample()613 614 615class VAEEncoder(BaseModel):616    def __init__(self, model, device, max_batch_size, embedding_dim):617        super(VAEEncoder, self).__init__(618            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim619        )620        self.name = "VAE encoder"621 622    def get_model(self):623        vae_encoder = TorchVAEEncoder(self.model)624        return vae_encoder625 626    def get_input_names(self):627        return ["images"]628 629    def get_output_names(self):630        return ["latent"]631 632    def get_dynamic_axes(self):633        return {"images": {0: "B", 2: "8H", 3: "8W"}, "latent": {0: "B", 2: "H", 3: "W"}}634 635    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):636        assert batch_size >= self.min_batch and batch_size <= self.max_batch637        min_batch = batch_size if static_batch else self.min_batch638        max_batch = batch_size if static_batch else self.max_batch639        self.check_dims(batch_size, image_height, image_width)640        (641            min_batch,642            max_batch,643            min_image_height,644            max_image_height,645            min_image_width,646            max_image_width,647            _,648            _,649            _,650            _,651        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)652 653        return {654            "images": [655                (min_batch, 3, min_image_height, min_image_width),656                (batch_size, 3, image_height, image_width),657                (max_batch, 3, max_image_height, max_image_width),658            ]659        }660 661    def get_shape_dict(self, batch_size, image_height, image_width):662        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)663        return {664            "images": (batch_size, 3, image_height, image_width),665            "latent": (batch_size, 4, latent_height, latent_width),666        }667 668    def get_sample_input(self, batch_size, image_height, image_width):669        self.check_dims(batch_size, image_height, image_width)670        return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device)671 672 673def make_VAEEncoder(model, device, max_batch_size, embedding_dim, inpaint=False):674    return VAEEncoder(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)675 676 677class TensorRTStableDiffusionInpaintPipeline(StableDiffusionInpaintPipeline):678    r"""679    Pipeline for inpainting using TensorRT accelerated Stable Diffusion.680 681    This model inherits from [`StableDiffusionInpaintPipeline`]. Check the superclass documentation for the generic methods the682    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)683 684    Args:685        vae ([`AutoencoderKL`]):686            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.687        text_encoder ([`CLIPTextModel`]):688            Frozen text-encoder. Stable Diffusion uses the text portion of689            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically690            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.691        tokenizer (`CLIPTokenizer`):692            Tokenizer of class693            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).694        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.695        scheduler ([`SchedulerMixin`]):696            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of697            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].698        safety_checker ([`StableDiffusionSafetyChecker`]):699            Classification module that estimates whether generated images could be considered offensive or harmful.700            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.701        feature_extractor ([`CLIPFeatureExtractor`]):702            Model that extracts features from generated images to be used as inputs for the `safety_checker`.703    """704 705    def __init__(706        self,707        vae: AutoencoderKL,708        text_encoder: CLIPTextModel,709        tokenizer: CLIPTokenizer,710        unet: UNet2DConditionModel,711        scheduler: DDIMScheduler,712        safety_checker: StableDiffusionSafetyChecker,713        feature_extractor: CLIPFeatureExtractor,714        image_encoder: CLIPVisionModelWithProjection = None,715        requires_safety_checker: bool = True,716        stages=["clip", "unet", "vae", "vae_encoder"],717        image_height: int = 512,718        image_width: int = 512,719        max_batch_size: int = 16,720        # ONNX export parameters721        onnx_opset: int = 17,722        onnx_dir: str = "onnx",723        # TensorRT engine build parameters724        engine_dir: str = "engine",725        build_preview_features: bool = True,726        force_engine_rebuild: bool = False,727        timing_cache: str = "timing_cache",728    ):729        super().__init__(730            vae,731            text_encoder,732            tokenizer,733            unet,734            scheduler,735            safety_checker=safety_checker,736            feature_extractor=feature_extractor,737            image_encoder=image_encoder,738            requires_safety_checker=requires_safety_checker,739        )740 741        self.vae.forward = self.vae.decode742 743        self.stages = stages744        self.image_height, self.image_width = image_height, image_width745        self.inpaint = True746        self.onnx_opset = onnx_opset747        self.onnx_dir = onnx_dir748        self.engine_dir = engine_dir749        self.force_engine_rebuild = force_engine_rebuild750        self.timing_cache = timing_cache751        self.build_static_batch = False752        self.build_dynamic_shape = False753        self.build_preview_features = build_preview_features754 755        self.max_batch_size = max_batch_size756        # TODO: Restrict batch size to 4 for larger image dimensions as a WAR for TensorRT limitation.757        if self.build_dynamic_shape or self.image_height > 512 or self.image_width > 512:758            self.max_batch_size = 4759 760        self.stream = None  # loaded in loadResources()761        self.models = {}  # loaded in __loadModels()762        self.engine = {}  # loaded in build_engines()763 764    def __loadModels(self):765        # Load pipeline models766        self.embedding_dim = self.text_encoder.config.hidden_size767        models_args = {768            "device": self.torch_device,769            "max_batch_size": self.max_batch_size,770            "embedding_dim": self.embedding_dim,771            "inpaint": self.inpaint,772        }773        if "clip" in self.stages:774            self.models["clip"] = make_CLIP(self.text_encoder, **models_args)775        if "unet" in self.stages:776            self.models["unet"] = make_UNet(self.unet, **models_args, unet_dim=self.unet.config.in_channels)777        if "vae" in self.stages:778            self.models["vae"] = make_VAE(self.vae, **models_args)779        if "vae_encoder" in self.stages:780            self.models["vae_encoder"] = make_VAEEncoder(self.vae, **models_args)781 782    @classmethod783    @validate_hf_hub_args784    def set_cached_folder(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):785        cache_dir = kwargs.pop("cache_dir", None)786        resume_download = kwargs.pop("resume_download", False)787        proxies = kwargs.pop("proxies", None)788        local_files_only = kwargs.pop("local_files_only", False)789        token = kwargs.pop("token", None)790        revision = kwargs.pop("revision", None)791 792        cls.cached_folder = (793            pretrained_model_name_or_path794            if os.path.isdir(pretrained_model_name_or_path)795            else snapshot_download(796                pretrained_model_name_or_path,797                cache_dir=cache_dir,798                resume_download=resume_download,799                proxies=proxies,800                local_files_only=local_files_only,801                token=token,802                revision=revision,803            )804        )805 806    def to(self, torch_device: Optional[Union[str, torch.device]] = None, silence_dtype_warnings: bool = False):807        super().to(torch_device, silence_dtype_warnings=silence_dtype_warnings)808 809        self.onnx_dir = os.path.join(self.cached_folder, self.onnx_dir)810        self.engine_dir = os.path.join(self.cached_folder, self.engine_dir)811        self.timing_cache = os.path.join(self.cached_folder, self.timing_cache)812 813        # set device814        self.torch_device = self._execution_device815        logger.warning(f"Running inference on device: {self.torch_device}")816 817        # load models818        self.__loadModels()819 820        # build engines821        self.engine = build_engines(822            self.models,823            self.engine_dir,824            self.onnx_dir,825            self.onnx_opset,826            opt_image_height=self.image_height,827            opt_image_width=self.image_width,828            force_engine_rebuild=self.force_engine_rebuild,829            static_batch=self.build_static_batch,830            static_shape=not self.build_dynamic_shape,831            enable_preview=self.build_preview_features,832            timing_cache=self.timing_cache,833        )834 835        return self836 837    def __initialize_timesteps(self, num_inference_steps, strength):838        self.scheduler.set_timesteps(num_inference_steps)839        offset = self.scheduler.config.steps_offset if hasattr(self.scheduler, "steps_offset") else 0840        init_timestep = int(num_inference_steps * strength) + offset841        init_timestep = min(init_timestep, num_inference_steps)842        t_start = max(num_inference_steps - init_timestep + offset, 0)843        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :].to(self.torch_device)844        return timesteps, num_inference_steps - t_start845 846    def __preprocess_images(self, batch_size, images=()):847        init_images = []848        for image in images:849            image = image.to(self.torch_device).float()850            image = image.repeat(batch_size, 1, 1, 1)851            init_images.append(image)852        return tuple(init_images)853 854    def __encode_image(self, init_image):855        init_latents = runEngine(self.engine["vae_encoder"], {"images": device_view(init_image)}, self.stream)[856            "latent"857        ]858        init_latents = 0.18215 * init_latents859        return init_latents860 861    def __encode_prompt(self, prompt, negative_prompt):862        r"""863        Encodes the prompt into text encoder hidden states.864 865        Args:866             prompt (`str` or `List[str]`, *optional*):867                prompt to be encoded868            negative_prompt (`str` or `List[str]`, *optional*):869                The prompt or prompts not to guide the image generation. If not defined, one has to pass870                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.871                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).872        """873        # Tokenize prompt874        text_input_ids = (875            self.tokenizer(876                prompt,877                padding="max_length",878                max_length=self.tokenizer.model_max_length,879                truncation=True,880                return_tensors="pt",881            )882            .input_ids.type(torch.int32)883            .to(self.torch_device)884        )885 886        text_input_ids_inp = device_view(text_input_ids)887        # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt888        text_embeddings = runEngine(self.engine["clip"], {"input_ids": text_input_ids_inp}, self.stream)[889            "text_embeddings"890        ].clone()891 892        # Tokenize negative prompt893        uncond_input_ids = (894            self.tokenizer(895                negative_prompt,896                padding="max_length",897                max_length=self.tokenizer.model_max_length,898                truncation=True,899                return_tensors="pt",900            )901            .input_ids.type(torch.int32)902            .to(self.torch_device)903        )904        uncond_input_ids_inp = device_view(uncond_input_ids)905        uncond_embeddings = runEngine(self.engine["clip"], {"input_ids": uncond_input_ids_inp}, self.stream)[906            "text_embeddings"907        ]908 909        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance910        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)911 912        return text_embeddings913 914    def __denoise_latent(915        self, latents, text_embeddings, timesteps=None, step_offset=0, mask=None, masked_image_latents=None916    ):917        if not isinstance(timesteps, torch.Tensor):918            timesteps = self.scheduler.timesteps919        for step_index, timestep in enumerate(timesteps):920            # Expand the latents if we are doing classifier free guidance921            latent_model_input = torch.cat([latents] * 2)922            latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)923            if isinstance(mask, torch.Tensor):924                latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)925 926            # Predict the noise residual927            timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep928 929            sample_inp = device_view(latent_model_input)930            timestep_inp = device_view(timestep_float)931            embeddings_inp = device_view(text_embeddings)932            noise_pred = runEngine(933                self.engine["unet"],934                {"sample": sample_inp, "timestep": timestep_inp, "encoder_hidden_states": embeddings_inp},935                self.stream,936            )["latent"]937 938            # Perform guidance939            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)940            noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)941 942            latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample943 944        latents = 1.0 / 0.18215 * latents945        return latents946 947    def __decode_latent(self, latents):948        images = runEngine(self.engine["vae"], {"latent": device_view(latents)}, self.stream)["images"]949        images = (images / 2 + 0.5).clamp(0, 1)950        return images.cpu().permute(0, 2, 3, 1).float().numpy()951 952    def __loadResources(self, image_height, image_width, batch_size):953        self.stream = cuda.Stream()954 955        # Allocate buffers for TensorRT engine bindings956        for model_name, obj in self.models.items():957            self.engine[model_name].allocate_buffers(958                shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.torch_device959            )960 961    @torch.no_grad()962    def __call__(963        self,964        prompt: Union[str, List[str]] = None,965        image: Union[torch.Tensor, PIL.Image.Image] = None,966        mask_image: Union[torch.Tensor, PIL.Image.Image] = None,967        strength: float = 1.0,968        num_inference_steps: int = 50,969        guidance_scale: float = 7.5,970        negative_prompt: Optional[Union[str, List[str]]] = None,971        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,972    ):973        r"""974        Function invoked when calling the pipeline for generation.975 976        Args:977            prompt (`str` or `List[str]`, *optional*):978                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.979                instead.980            image (`PIL.Image.Image`):981                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will982                be masked out with `mask_image` and repainted according to `prompt`.983            mask_image (`PIL.Image.Image`):984                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be985                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted986                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)987                instead of 3, so the expected shape would be `(B, H, W, 1)`.988            strength (`float`, *optional*, defaults to 0.8):989                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`990                will be used as a starting point, adding more noise to it the larger the `strength`. The number of991                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will992                be maximum and the denoising process will run for the full number of iterations specified in993                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.994            num_inference_steps (`int`, *optional*, defaults to 50):995                The number of denoising steps. More denoising steps usually lead to a higher quality image at the996                expense of slower inference.997            guidance_scale (`float`, *optional*, defaults to 7.5):998                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).999                `guidance_scale` is defined as `w` of equation 2. of [Imagen1000                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1001                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1002                usually at the expense of lower image quality.1003            negative_prompt (`str` or `List[str]`, *optional*):1004                The prompt or prompts not to guide the image generation. If not defined, one has to pass1005                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.1006                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).1007            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):1008                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)1009                to make generation deterministic.1010 1011        """1012        self.generator = generator1013        self.denoising_steps = num_inference_steps1014        self._guidance_scale = guidance_scale1015 1016        # Pre-compute latent input scales and linear multistep coefficients1017        self.scheduler.set_timesteps(self.denoising_steps, device=self.torch_device)1018 1019        # Define call parameters1020        if prompt is not None and isinstance(prompt, str):1021            batch_size = 11022            prompt = [prompt]1023        elif prompt is not None and isinstance(prompt, list):1024            batch_size = len(prompt)1025        else:1026            raise ValueError(f"Expected prompt to be of type list or str but got {type(prompt)}")1027 1028        if negative_prompt is None:1029            negative_prompt = [""] * batch_size1030 1031        if negative_prompt is not None and isinstance(negative_prompt, str):1032            negative_prompt = [negative_prompt]1033 1034        assert len(prompt) == len(negative_prompt)1035 1036        if batch_size > self.max_batch_size:1037            raise ValueError(1038                f"Batch size {len(prompt)} is larger than allowed {self.max_batch_size}. If dynamic shape is used, then maximum batch size is 4"1039            )1040 1041        # Validate image dimensions1042        mask_width, mask_height = mask_image.size1043        if mask_height != self.image_height or mask_width != self.image_width:1044            raise ValueError(1045                f"Input image height and width {self.image_height} and {self.image_width} are not equal to "1046                f"the respective dimensions of the mask image {mask_height} and {mask_width}"1047            )1048 1049        # load resources1050        self.__loadResources(self.image_height, self.image_width, batch_size)1051 1052        with torch.inference_mode(), torch.autocast("cuda"), trt.Runtime(TRT_LOGGER):1053            # Spatial dimensions of latent tensor1054            latent_height = self.image_height // 81055            latent_width = self.image_width // 81056 1057            # Pre-process input images1058            mask, masked_image, init_image = self.__preprocess_images(1059                batch_size,1060                prepare_mask_and_masked_image(1061                    image,1062                    mask_image,1063                    self.image_height,1064                    self.image_width,1065                    return_image=True,1066                ),1067            )1068 1069            mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width))1070            mask = torch.cat([mask] * 2)1071 1072            # Initialize timesteps1073            timesteps, t_start = self.__initialize_timesteps(self.denoising_steps, strength)1074 1075            # at which timestep to set the initial noise (n.b. 50% if strength is 0.5)1076            latent_timestep = timesteps[:1].repeat(batch_size)1077            # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise1078            is_strength_max = strength == 1.01079 1080            # Pre-initialize latents1081            num_channels_latents = self.vae.config.latent_channels1082            latents_outputs = self.prepare_latents(1083                batch_size,1084                num_channels_latents,1085                self.image_height,1086                self.image_width,1087                torch.float32,1088                self.torch_device,1089                generator,1090                image=init_image,1091                timestep=latent_timestep,1092                is_strength_max=is_strength_max,1093            )1094 1095            latents = latents_outputs[0]1096 1097            # VAE encode masked image1098            masked_latents = self.__encode_image(masked_image)1099            masked_latents = torch.cat([masked_latents] * 2)1100 1101            # CLIP text encoder1102            text_embeddings = self.__encode_prompt(prompt, negative_prompt)1103 1104            # UNet denoiser1105            latents = self.__denoise_latent(1106                latents,1107                text_embeddings,1108                timesteps=timesteps,1109                step_offset=t_start,1110                mask=mask,1111                masked_image_latents=masked_latents,1112            )1113 1114            # VAE decode latent1115            images = self.__decode_latent(latents)1116 1117        images = self.numpy_to_pil(images)1118        return StableDiffusionPipelineOutput(images=images, nsfw_content_detected=None)1119