diffusers/community-pipelines-mirror

1#2# Copyright 2023 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 onnx import shape_inference32from polygraphy import cuda33from polygraphy.backend.common import bytes_from_path34from polygraphy.backend.onnx.loader import fold_constants35from polygraphy.backend.trt import (36    CreateConfig,37    Profile,38    engine_from_bytes,39    engine_from_network,40    network_from_onnx_path,41    save_engine,42)43from polygraphy.backend.trt import util as trt_util44from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer45 46from diffusers.models import AutoencoderKL, UNet2DConditionModel47from diffusers.pipelines.stable_diffusion import (48    StableDiffusionInpaintPipeline,49    StableDiffusionPipelineOutput,50    StableDiffusionSafetyChecker,51)52from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint import prepare_mask_and_masked_image53from diffusers.schedulers import DDIMScheduler54from diffusers.utils import DIFFUSERS_CACHE, logging55 56 57"""58Installation instructions59python3 -m pip install --upgrade transformers diffusers>=0.16.060python3 -m pip install --upgrade tensorrt>=8.6.161python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com62python3 -m pip install onnxruntime63"""64 65TRT_LOGGER = trt.Logger(trt.Logger.ERROR)66logger = logging.get_logger(__name__)  # pylint: disable=invalid-name67 68# Map of numpy dtype -> torch dtype69numpy_to_torch_dtype_dict = {70    np.uint8: torch.uint8,71    np.int8: torch.int8,72    np.int16: torch.int16,73    np.int32: torch.int32,74    np.int64: torch.int64,75    np.float16: torch.float16,76    np.float32: torch.float32,77    np.float64: torch.float64,78    np.complex64: torch.complex64,79    np.complex128: torch.complex128,80}81if np.version.full_version >= "1.24.0":82    numpy_to_torch_dtype_dict[np.bool_] = torch.bool83else:84    numpy_to_torch_dtype_dict[np.bool] = torch.bool85 86# Map of torch dtype -> numpy dtype87torch_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()}88 89 90def device_view(t):91    return cuda.DeviceView(ptr=t.data_ptr(), shape=t.shape, dtype=torch_to_numpy_dtype_dict[t.dtype])92 93 94def preprocess_image(image):95    """96    image: torch.Tensor97    """98    w, h = image.size99    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32100    image = image.resize((w, h))101    image = np.array(image).astype(np.float32) / 255.0102    image = image[None].transpose(0, 3, 1, 2)103    image = torch.from_numpy(image).contiguous()104    return 2.0 * image - 1.0105 106 107class Engine:108    def __init__(self, engine_path):109        self.engine_path = engine_path110        self.engine = None111        self.context = None112        self.buffers = OrderedDict()113        self.tensors = OrderedDict()114 115    def __del__(self):116        [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray)]117        del self.engine118        del self.context119        del self.buffers120        del self.tensors121 122    def build(123        self,124        onnx_path,125        fp16,126        input_profile=None,127        enable_preview=False,128        enable_all_tactics=False,129        timing_cache=None,130        workspace_size=0,131    ):132        logger.warning(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")133        p = Profile()134        if input_profile:135            for name, dims in input_profile.items():136                assert len(dims) == 3137                p.add(name, min=dims[0], opt=dims[1], max=dims[2])138 139        config_kwargs = {}140 141        config_kwargs["preview_features"] = [trt.PreviewFeature.DISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805]142        if enable_preview:143            # Faster dynamic shapes made optional since it increases engine build time.144            config_kwargs["preview_features"].append(trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805)145        if workspace_size > 0:146            config_kwargs["memory_pool_limits"] = {trt.MemoryPoolType.WORKSPACE: workspace_size}147        if not enable_all_tactics:148            config_kwargs["tactic_sources"] = []149 150        engine = engine_from_network(151            network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]),152            config=CreateConfig(fp16=fp16, profiles=[p], load_timing_cache=timing_cache, **config_kwargs),153            save_timing_cache=timing_cache,154        )155        save_engine(engine, path=self.engine_path)156 157    def load(self):158        logger.warning(f"Loading TensorRT engine: {self.engine_path}")159        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))160 161    def activate(self):162        self.context = self.engine.create_execution_context()163 164    def allocate_buffers(self, shape_dict=None, device="cuda"):165        for idx in range(trt_util.get_bindings_per_profile(self.engine)):166            binding = self.engine[idx]167            if shape_dict and binding in shape_dict:168                shape = shape_dict[binding]169            else:170                shape = self.engine.get_binding_shape(binding)171            dtype = trt.nptype(self.engine.get_binding_dtype(binding))172            if self.engine.binding_is_input(binding):173                self.context.set_binding_shape(idx, shape)174            tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device)175            self.tensors[binding] = tensor176            self.buffers[binding] = cuda.DeviceView(ptr=tensor.data_ptr(), shape=shape, dtype=dtype)177 178    def infer(self, feed_dict, stream):179        start_binding, end_binding = trt_util.get_active_profile_bindings(self.context)180        # shallow copy of ordered dict181        device_buffers = copy(self.buffers)182        for name, buf in feed_dict.items():183            assert isinstance(buf, cuda.DeviceView)184            device_buffers[name] = buf185        bindings = [0] * start_binding + [buf.ptr for buf in device_buffers.values()]186        noerror = self.context.execute_async_v2(bindings=bindings, stream_handle=stream.ptr)187        if not noerror:188            raise ValueError("ERROR: inference failed.")189 190        return self.tensors191 192 193class Optimizer:194    def __init__(self, onnx_graph):195        self.graph = gs.import_onnx(onnx_graph)196 197    def cleanup(self, return_onnx=False):198        self.graph.cleanup().toposort()199        if return_onnx:200            return gs.export_onnx(self.graph)201 202    def select_outputs(self, keep, names=None):203        self.graph.outputs = [self.graph.outputs[o] for o in keep]204        if names:205            for i, name in enumerate(names):206                self.graph.outputs[i].name = name207 208    def fold_constants(self, return_onnx=False):209        onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)210        self.graph = gs.import_onnx(onnx_graph)211        if return_onnx:212            return onnx_graph213 214    def infer_shapes(self, return_onnx=False):215        onnx_graph = gs.export_onnx(self.graph)216        if onnx_graph.ByteSize() > 2147483648:217            raise TypeError("ERROR: model size exceeds supported 2GB limit")218        else:219            onnx_graph = shape_inference.infer_shapes(onnx_graph)220 221        self.graph = gs.import_onnx(onnx_graph)222        if return_onnx:223            return onnx_graph224 225 226class BaseModel:227    def __init__(self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77):228        self.model = model229        self.name = "SD Model"230        self.fp16 = fp16231        self.device = device232 233        self.min_batch = 1234        self.max_batch = max_batch_size235        self.min_image_shape = 256  # min image resolution: 256x256236        self.max_image_shape = 1024  # max image resolution: 1024x1024237        self.min_latent_shape = self.min_image_shape // 8238        self.max_latent_shape = self.max_image_shape // 8239 240        self.embedding_dim = embedding_dim241        self.text_maxlen = text_maxlen242 243    def get_model(self):244        return self.model245 246    def get_input_names(self):247        pass248 249    def get_output_names(self):250        pass251 252    def get_dynamic_axes(self):253        return None254 255    def get_sample_input(self, batch_size, image_height, image_width):256        pass257 258    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):259        return None260 261    def get_shape_dict(self, batch_size, image_height, image_width):262        return None263 264    def optimize(self, onnx_graph):265        opt = Optimizer(onnx_graph)266        opt.cleanup()267        opt.fold_constants()268        opt.infer_shapes()269        onnx_opt_graph = opt.cleanup(return_onnx=True)270        return onnx_opt_graph271 272    def check_dims(self, batch_size, image_height, image_width):273        assert batch_size >= self.min_batch and batch_size <= self.max_batch274        assert image_height % 8 == 0 or image_width % 8 == 0275        latent_height = image_height // 8276        latent_width = image_width // 8277        assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape278        assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape279        return (latent_height, latent_width)280 281    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):282        min_batch = batch_size if static_batch else self.min_batch283        max_batch = batch_size if static_batch else self.max_batch284        latent_height = image_height // 8285        latent_width = image_width // 8286        min_image_height = image_height if static_shape else self.min_image_shape287        max_image_height = image_height if static_shape else self.max_image_shape288        min_image_width = image_width if static_shape else self.min_image_shape289        max_image_width = image_width if static_shape else self.max_image_shape290        min_latent_height = latent_height if static_shape else self.min_latent_shape291        max_latent_height = latent_height if static_shape else self.max_latent_shape292        min_latent_width = latent_width if static_shape else self.min_latent_shape293        max_latent_width = latent_width if static_shape else self.max_latent_shape294        return (295            min_batch,296            max_batch,297            min_image_height,298            max_image_height,299            min_image_width,300            max_image_width,301            min_latent_height,302            max_latent_height,303            min_latent_width,304            max_latent_width,305        )306 307 308def getOnnxPath(model_name, onnx_dir, opt=True):309    return os.path.join(onnx_dir, model_name + (".opt" if opt else "") + ".onnx")310 311 312def getEnginePath(model_name, engine_dir):313    return os.path.join(engine_dir, model_name + ".plan")314 315 316def build_engines(317    models: dict,318    engine_dir,319    onnx_dir,320    onnx_opset,321    opt_image_height,322    opt_image_width,323    opt_batch_size=1,324    force_engine_rebuild=False,325    static_batch=False,326    static_shape=True,327    enable_preview=False,328    enable_all_tactics=False,329    timing_cache=None,330    max_workspace_size=0,331):332    built_engines = {}333    if not os.path.isdir(onnx_dir):334        os.makedirs(onnx_dir)335    if not os.path.isdir(engine_dir):336        os.makedirs(engine_dir)337 338    # Export models to ONNX339    for model_name, model_obj in models.items():340        engine_path = getEnginePath(model_name, engine_dir)341        if force_engine_rebuild or not os.path.exists(engine_path):342            logger.warning("Building Engines...")343            logger.warning("Engine build can take a while to complete")344            onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)345            onnx_opt_path = getOnnxPath(model_name, onnx_dir)346            if force_engine_rebuild or not os.path.exists(onnx_opt_path):347                if force_engine_rebuild or not os.path.exists(onnx_path):348                    logger.warning(f"Exporting model: {onnx_path}")349                    model = model_obj.get_model()350                    with torch.inference_mode(), torch.autocast("cuda"):351                        inputs = model_obj.get_sample_input(opt_batch_size, opt_image_height, opt_image_width)352                        torch.onnx.export(353                            model,354                            inputs,355                            onnx_path,356                            export_params=True,357                            opset_version=onnx_opset,358                            do_constant_folding=True,359                            input_names=model_obj.get_input_names(),360                            output_names=model_obj.get_output_names(),361                            dynamic_axes=model_obj.get_dynamic_axes(),362                        )363                    del model364                    torch.cuda.empty_cache()365                    gc.collect()366                else:367                    logger.warning(f"Found cached model: {onnx_path}")368 369                # Optimize onnx370                if force_engine_rebuild or not os.path.exists(onnx_opt_path):371                    logger.warning(f"Generating optimizing model: {onnx_opt_path}")372                    onnx_opt_graph = model_obj.optimize(onnx.load(onnx_path))373                    onnx.save(onnx_opt_graph, onnx_opt_path)374                else:375                    logger.warning(f"Found cached optimized model: {onnx_opt_path} ")376 377    # Build TensorRT engines378    for model_name, model_obj in models.items():379        engine_path = getEnginePath(model_name, engine_dir)380        engine = Engine(engine_path)381        onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)382        onnx_opt_path = getOnnxPath(model_name, onnx_dir)383 384        if force_engine_rebuild or not os.path.exists(engine.engine_path):385            engine.build(386                onnx_opt_path,387                fp16=True,388                input_profile=model_obj.get_input_profile(389                    opt_batch_size,390                    opt_image_height,391                    opt_image_width,392                    static_batch=static_batch,393                    static_shape=static_shape,394                ),395                enable_preview=enable_preview,396                timing_cache=timing_cache,397                workspace_size=max_workspace_size,398            )399        built_engines[model_name] = engine400 401    # Load and activate TensorRT engines402    for model_name, model_obj in models.items():403        engine = built_engines[model_name]404        engine.load()405        engine.activate()406 407    return built_engines408 409 410def runEngine(engine, feed_dict, stream):411    return engine.infer(feed_dict, stream)412 413 414class CLIP(BaseModel):415    def __init__(self, model, device, max_batch_size, embedding_dim):416        super(CLIP, self).__init__(417            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim418        )419        self.name = "CLIP"420 421    def get_input_names(self):422        return ["input_ids"]423 424    def get_output_names(self):425        return ["text_embeddings", "pooler_output"]426 427    def get_dynamic_axes(self):428        return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}429 430    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):431        self.check_dims(batch_size, image_height, image_width)432        min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(433            batch_size, image_height, image_width, static_batch, static_shape434        )435        return {436            "input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]437        }438 439    def get_shape_dict(self, batch_size, image_height, image_width):440        self.check_dims(batch_size, image_height, image_width)441        return {442            "input_ids": (batch_size, self.text_maxlen),443            "text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),444        }445 446    def get_sample_input(self, batch_size, image_height, image_width):447        self.check_dims(batch_size, image_height, image_width)448        return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)449 450    def optimize(self, onnx_graph):451        opt = Optimizer(onnx_graph)452        opt.select_outputs([0])  # delete graph output#1453        opt.cleanup()454        opt.fold_constants()455        opt.infer_shapes()456        opt.select_outputs([0], names=["text_embeddings"])  # rename network output457        opt_onnx_graph = opt.cleanup(return_onnx=True)458        return opt_onnx_graph459 460 461def make_CLIP(model, device, max_batch_size, embedding_dim, inpaint=False):462    return CLIP(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)463 464 465class UNet(BaseModel):466    def __init__(467        self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77, unet_dim=4468    ):469        super(UNet, self).__init__(470            model=model,471            fp16=fp16,472            device=device,473            max_batch_size=max_batch_size,474            embedding_dim=embedding_dim,475            text_maxlen=text_maxlen,476        )477        self.unet_dim = unet_dim478        self.name = "UNet"479 480    def get_input_names(self):481        return ["sample", "timestep", "encoder_hidden_states"]482 483    def get_output_names(self):484        return ["latent"]485 486    def get_dynamic_axes(self):487        return {488            "sample": {0: "2B", 2: "H", 3: "W"},489            "encoder_hidden_states": {0: "2B"},490            "latent": {0: "2B", 2: "H", 3: "W"},491        }492 493    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):494        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)495        (496            min_batch,497            max_batch,498            _,499            _,500            _,501            _,502            min_latent_height,503            max_latent_height,504            min_latent_width,505            max_latent_width,506        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)507        return {508            "sample": [509                (2 * min_batch, self.unet_dim, min_latent_height, min_latent_width),510                (2 * batch_size, self.unet_dim, latent_height, latent_width),511                (2 * max_batch, self.unet_dim, max_latent_height, max_latent_width),512            ],513            "encoder_hidden_states": [514                (2 * min_batch, self.text_maxlen, self.embedding_dim),515                (2 * batch_size, self.text_maxlen, self.embedding_dim),516                (2 * max_batch, self.text_maxlen, self.embedding_dim),517            ],518        }519 520    def get_shape_dict(self, batch_size, image_height, image_width):521        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)522        return {523            "sample": (2 * batch_size, self.unet_dim, latent_height, latent_width),524            "encoder_hidden_states": (2 * batch_size, self.text_maxlen, self.embedding_dim),525            "latent": (2 * batch_size, 4, latent_height, latent_width),526        }527 528    def get_sample_input(self, batch_size, image_height, image_width):529        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)530        dtype = torch.float16 if self.fp16 else torch.float32531        return (532            torch.randn(533                2 * batch_size, self.unet_dim, latent_height, latent_width, dtype=torch.float32, device=self.device534            ),535            torch.tensor([1.0], dtype=torch.float32, device=self.device),536            torch.randn(2 * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),537        )538 539 540def make_UNet(model, device, max_batch_size, embedding_dim, inpaint=False, unet_dim=4):541    return UNet(542        model,543        fp16=True,544        device=device,545        max_batch_size=max_batch_size,546        embedding_dim=embedding_dim,547        unet_dim=unet_dim,548    )549 550 551class VAE(BaseModel):552    def __init__(self, model, device, max_batch_size, embedding_dim):553        super(VAE, self).__init__(554            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim555        )556        self.name = "VAE decoder"557 558    def get_input_names(self):559        return ["latent"]560 561    def get_output_names(self):562        return ["images"]563 564    def get_dynamic_axes(self):565        return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}566 567    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):568        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)569        (570            min_batch,571            max_batch,572            _,573            _,574            _,575            _,576            min_latent_height,577            max_latent_height,578            min_latent_width,579            max_latent_width,580        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)581        return {582            "latent": [583                (min_batch, 4, min_latent_height, min_latent_width),584                (batch_size, 4, latent_height, latent_width),585                (max_batch, 4, max_latent_height, max_latent_width),586            ]587        }588 589    def get_shape_dict(self, batch_size, image_height, image_width):590        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)591        return {592            "latent": (batch_size, 4, latent_height, latent_width),593            "images": (batch_size, 3, image_height, image_width),594        }595 596    def get_sample_input(self, batch_size, image_height, image_width):597        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)598        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)599 600 601def make_VAE(model, device, max_batch_size, embedding_dim, inpaint=False):602    return VAE(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)603 604 605class TorchVAEEncoder(torch.nn.Module):606    def __init__(self, model):607        super().__init__()608        self.vae_encoder = model609 610    def forward(self, x):611        return self.vae_encoder.encode(x).latent_dist.sample()612 613 614class VAEEncoder(BaseModel):615    def __init__(self, model, device, max_batch_size, embedding_dim):616        super(VAEEncoder, self).__init__(617            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim618        )619        self.name = "VAE encoder"620 621    def get_model(self):622        vae_encoder = TorchVAEEncoder(self.model)623        return vae_encoder624 625    def get_input_names(self):626        return ["images"]627 628    def get_output_names(self):629        return ["latent"]630 631    def get_dynamic_axes(self):632        return {"images": {0: "B", 2: "8H", 3: "8W"}, "latent": {0: "B", 2: "H", 3: "W"}}633 634    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):635        assert batch_size >= self.min_batch and batch_size <= self.max_batch636        min_batch = batch_size if static_batch else self.min_batch637        max_batch = batch_size if static_batch else self.max_batch638        self.check_dims(batch_size, image_height, image_width)639        (640            min_batch,641            max_batch,642            min_image_height,643            max_image_height,644            min_image_width,645            max_image_width,646            _,647            _,648            _,649            _,650        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)651 652        return {653            "images": [654                (min_batch, 3, min_image_height, min_image_width),655                (batch_size, 3, image_height, image_width),656                (max_batch, 3, max_image_height, max_image_width),657            ]658        }659 660    def get_shape_dict(self, batch_size, image_height, image_width):661        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)662        return {663            "images": (batch_size, 3, image_height, image_width),664            "latent": (batch_size, 4, latent_height, latent_width),665        }666 667    def get_sample_input(self, batch_size, image_height, image_width):668        self.check_dims(batch_size, image_height, image_width)669        return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device)670 671 672def make_VAEEncoder(model, device, max_batch_size, embedding_dim, inpaint=False):673    return VAEEncoder(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)674 675 676class TensorRTStableDiffusionInpaintPipeline(StableDiffusionInpaintPipeline):677    r"""678    Pipeline for inpainting using TensorRT accelerated Stable Diffusion.679 680    This model inherits from [`StableDiffusionInpaintPipeline`]. Check the superclass documentation for the generic methods the681    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)682 683    Args:684        vae ([`AutoencoderKL`]):685            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.686        text_encoder ([`CLIPTextModel`]):687            Frozen text-encoder. Stable Diffusion uses the text portion of688            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically689            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.690        tokenizer (`CLIPTokenizer`):691            Tokenizer of class692            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).693        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.694        scheduler ([`SchedulerMixin`]):695            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of696            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].697        safety_checker ([`StableDiffusionSafetyChecker`]):698            Classification module that estimates whether generated images could be considered offensive or harmful.699            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.700        feature_extractor ([`CLIPFeatureExtractor`]):701            Model that extracts features from generated images to be used as inputs for the `safety_checker`.702    """703 704    def __init__(705        self,706        vae: AutoencoderKL,707        text_encoder: CLIPTextModel,708        tokenizer: CLIPTokenizer,709        unet: UNet2DConditionModel,710        scheduler: DDIMScheduler,711        safety_checker: StableDiffusionSafetyChecker,712        feature_extractor: CLIPFeatureExtractor,713        requires_safety_checker: bool = True,714        stages=["clip", "unet", "vae", "vae_encoder"],715        image_height: int = 512,716        image_width: int = 512,717        max_batch_size: int = 16,718        # ONNX export parameters719        onnx_opset: int = 17,720        onnx_dir: str = "onnx",721        # TensorRT engine build parameters722        engine_dir: str = "engine",723        build_preview_features: bool = True,724        force_engine_rebuild: bool = False,725        timing_cache: str = "timing_cache",726    ):727        super().__init__(728            vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker729        )730 731        self.vae.forward = self.vae.decode732 733        self.stages = stages734        self.image_height, self.image_width = image_height, image_width735        self.inpaint = True736        self.onnx_opset = onnx_opset737        self.onnx_dir = onnx_dir738        self.engine_dir = engine_dir739        self.force_engine_rebuild = force_engine_rebuild740        self.timing_cache = timing_cache741        self.build_static_batch = False742        self.build_dynamic_shape = False743        self.build_preview_features = build_preview_features744 745        self.max_batch_size = max_batch_size746        # TODO: Restrict batch size to 4 for larger image dimensions as a WAR for TensorRT limitation.747        if self.build_dynamic_shape or self.image_height > 512 or self.image_width > 512:748            self.max_batch_size = 4749 750        self.stream = None  # loaded in loadResources()751        self.models = {}  # loaded in __loadModels()752        self.engine = {}  # loaded in build_engines()753 754    def __loadModels(self):755        # Load pipeline models756        self.embedding_dim = self.text_encoder.config.hidden_size757        models_args = {758            "device": self.torch_device,759            "max_batch_size": self.max_batch_size,760            "embedding_dim": self.embedding_dim,761            "inpaint": self.inpaint,762        }763        if "clip" in self.stages:764            self.models["clip"] = make_CLIP(self.text_encoder, **models_args)765        if "unet" in self.stages:766            self.models["unet"] = make_UNet(self.unet, **models_args, unet_dim=self.unet.config.in_channels)767        if "vae" in self.stages:768            self.models["vae"] = make_VAE(self.vae, **models_args)769        if "vae_encoder" in self.stages:770            self.models["vae_encoder"] = make_VAEEncoder(self.vae, **models_args)771 772    @classmethod773    def set_cached_folder(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):774        cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)775        resume_download = kwargs.pop("resume_download", False)776        proxies = kwargs.pop("proxies", None)777        local_files_only = kwargs.pop("local_files_only", False)778        use_auth_token = kwargs.pop("use_auth_token", None)779        revision = kwargs.pop("revision", None)780 781        cls.cached_folder = (782            pretrained_model_name_or_path783            if os.path.isdir(pretrained_model_name_or_path)784            else snapshot_download(785                pretrained_model_name_or_path,786                cache_dir=cache_dir,787                resume_download=resume_download,788                proxies=proxies,789                local_files_only=local_files_only,790                use_auth_token=use_auth_token,791                revision=revision,792            )793        )794 795    def to(self, torch_device: Optional[Union[str, torch.device]] = None, silence_dtype_warnings: bool = False):796        super().to(torch_device, silence_dtype_warnings=silence_dtype_warnings)797 798        self.onnx_dir = os.path.join(self.cached_folder, self.onnx_dir)799        self.engine_dir = os.path.join(self.cached_folder, self.engine_dir)800        self.timing_cache = os.path.join(self.cached_folder, self.timing_cache)801 802        # set device803        self.torch_device = self._execution_device804        logger.warning(f"Running inference on device: {self.torch_device}")805 806        # load models807        self.__loadModels()808 809        # build engines810        self.engine = build_engines(811            self.models,812            self.engine_dir,813            self.onnx_dir,814            self.onnx_opset,815            opt_image_height=self.image_height,816            opt_image_width=self.image_width,817            force_engine_rebuild=self.force_engine_rebuild,818            static_batch=self.build_static_batch,819            static_shape=not self.build_dynamic_shape,820            enable_preview=self.build_preview_features,821            timing_cache=self.timing_cache,822        )823 824        return self825 826    def __initialize_timesteps(self, num_inference_steps, strength):827        self.scheduler.set_timesteps(num_inference_steps)828        offset = self.scheduler.config.steps_offset if hasattr(self.scheduler, "steps_offset") else 0829        init_timestep = int(num_inference_steps * strength) + offset830        init_timestep = min(init_timestep, num_inference_steps)831        t_start = max(num_inference_steps - init_timestep + offset, 0)832        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :].to(self.torch_device)833        return timesteps, num_inference_steps - t_start834 835    def __preprocess_images(self, batch_size, images=()):836        init_images = []837        for image in images:838            image = image.to(self.torch_device).float()839            image = image.repeat(batch_size, 1, 1, 1)840            init_images.append(image)841        return tuple(init_images)842 843    def __encode_image(self, init_image):844        init_latents = runEngine(self.engine["vae_encoder"], {"images": device_view(init_image)}, self.stream)[845            "latent"846        ]847        init_latents = 0.18215 * init_latents848        return init_latents849 850    def __encode_prompt(self, prompt, negative_prompt):851        r"""852        Encodes the prompt into text encoder hidden states.853 854        Args:855             prompt (`str` or `List[str]`, *optional*):856                prompt to be encoded857            negative_prompt (`str` or `List[str]`, *optional*):858                The prompt or prompts not to guide the image generation. If not defined, one has to pass859                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.860                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).861        """862        # Tokenize prompt863        text_input_ids = (864            self.tokenizer(865                prompt,866                padding="max_length",867                max_length=self.tokenizer.model_max_length,868                truncation=True,869                return_tensors="pt",870            )871            .input_ids.type(torch.int32)872            .to(self.torch_device)873        )874 875        text_input_ids_inp = device_view(text_input_ids)876        # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt877        text_embeddings = runEngine(self.engine["clip"], {"input_ids": text_input_ids_inp}, self.stream)[878            "text_embeddings"879        ].clone()880 881        # Tokenize negative prompt882        uncond_input_ids = (883            self.tokenizer(884                negative_prompt,885                padding="max_length",886                max_length=self.tokenizer.model_max_length,887                truncation=True,888                return_tensors="pt",889            )890            .input_ids.type(torch.int32)891            .to(self.torch_device)892        )893        uncond_input_ids_inp = device_view(uncond_input_ids)894        uncond_embeddings = runEngine(self.engine["clip"], {"input_ids": uncond_input_ids_inp}, self.stream)[895            "text_embeddings"896        ]897 898        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance899        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)900 901        return text_embeddings902 903    def __denoise_latent(904        self, latents, text_embeddings, timesteps=None, step_offset=0, mask=None, masked_image_latents=None905    ):906        if not isinstance(timesteps, torch.Tensor):907            timesteps = self.scheduler.timesteps908        for step_index, timestep in enumerate(timesteps):909            # Expand the latents if we are doing classifier free guidance910            latent_model_input = torch.cat([latents] * 2)911            latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)912            if isinstance(mask, torch.Tensor):913                latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)914 915            # Predict the noise residual916            timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep917 918            sample_inp = device_view(latent_model_input)919            timestep_inp = device_view(timestep_float)920            embeddings_inp = device_view(text_embeddings)921            noise_pred = runEngine(922                self.engine["unet"],923                {"sample": sample_inp, "timestep": timestep_inp, "encoder_hidden_states": embeddings_inp},924                self.stream,925            )["latent"]926 927            # Perform guidance928            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)929            noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)930 931            latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample932 933        latents = 1.0 / 0.18215 * latents934        return latents935 936    def __decode_latent(self, latents):937        images = runEngine(self.engine["vae"], {"latent": device_view(latents)}, self.stream)["images"]938        images = (images / 2 + 0.5).clamp(0, 1)939        return images.cpu().permute(0, 2, 3, 1).float().numpy()940 941    def __loadResources(self, image_height, image_width, batch_size):942        self.stream = cuda.Stream()943 944        # Allocate buffers for TensorRT engine bindings945        for model_name, obj in self.models.items():946            self.engine[model_name].allocate_buffers(947                shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.torch_device948            )949 950    @torch.no_grad()951    def __call__(952        self,953        prompt: Union[str, List[str]] = None,954        image: Union[torch.FloatTensor, PIL.Image.Image] = None,955        mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None,956        strength: float = 1.0,957        num_inference_steps: int = 50,958        guidance_scale: float = 7.5,959        negative_prompt: Optional[Union[str, List[str]]] = None,960        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,961    ):962        r"""963        Function invoked when calling the pipeline for generation.964 965        Args:966            prompt (`str` or `List[str]`, *optional*):967                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.968                instead.969            image (`PIL.Image.Image`):970                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will971                be masked out with `mask_image` and repainted according to `prompt`.972            mask_image (`PIL.Image.Image`):973                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be974                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted975                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)976                instead of 3, so the expected shape would be `(B, H, W, 1)`.977            strength (`float`, *optional*, defaults to 0.8):978                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`979                will be used as a starting point, adding more noise to it the larger the `strength`. The number of980                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will981                be maximum and the denoising process will run for the full number of iterations specified in982                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.983            num_inference_steps (`int`, *optional*, defaults to 50):984                The number of denoising steps. More denoising steps usually lead to a higher quality image at the985                expense of slower inference.986            guidance_scale (`float`, *optional*, defaults to 7.5):987                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).988                `guidance_scale` is defined as `w` of equation 2. of [Imagen989                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >990                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,991                usually at the expense of lower image quality.992            negative_prompt (`str` or `List[str]`, *optional*):993                The prompt or prompts not to guide the image generation. If not defined, one has to pass994                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.995                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).996            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):997                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)998                to make generation deterministic.999 1000        """1001        self.generator = generator1002        self.denoising_steps = num_inference_steps1003        self.guidance_scale = guidance_scale1004 1005        # Pre-compute latent input scales and linear multistep coefficients1006        self.scheduler.set_timesteps(self.denoising_steps, device=self.torch_device)1007 1008        # Define call parameters1009        if prompt is not None and isinstance(prompt, str):1010            batch_size = 11011            prompt = [prompt]1012        elif prompt is not None and isinstance(prompt, list):1013            batch_size = len(prompt)1014        else:1015            raise ValueError(f"Expected prompt to be of type list or str but got {type(prompt)}")1016 1017        if negative_prompt is None:1018            negative_prompt = [""] * batch_size1019 1020        if negative_prompt is not None and isinstance(negative_prompt, str):1021            negative_prompt = [negative_prompt]1022 1023        assert len(prompt) == len(negative_prompt)1024 1025        if batch_size > self.max_batch_size:1026            raise ValueError(1027                f"Batch size {len(prompt)} is larger than allowed {self.max_batch_size}. If dynamic shape is used, then maximum batch size is 4"1028            )1029 1030        # Validate image dimensions1031        mask_width, mask_height = mask_image.size1032        if mask_height != self.image_height or mask_width != self.image_width:1033            raise ValueError(1034                f"Input image height and width {self.image_height} and {self.image_width} are not equal to "1035                f"the respective dimensions of the mask image {mask_height} and {mask_width}"1036            )1037 1038        # load resources1039        self.__loadResources(self.image_height, self.image_width, batch_size)1040 1041        with torch.inference_mode(), torch.autocast("cuda"), trt.Runtime(TRT_LOGGER):1042            # Spatial dimensions of latent tensor1043            latent_height = self.image_height // 81044            latent_width = self.image_width // 81045 1046            # Pre-process input images1047            mask, masked_image, init_image = self.__preprocess_images(1048                batch_size,1049                prepare_mask_and_masked_image(1050                    image,1051                    mask_image,1052                    self.image_height,1053                    self.image_width,1054                    return_image=True,1055                ),1056            )1057 1058            mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width))1059            mask = torch.cat([mask] * 2)1060 1061            # Initialize timesteps1062            timesteps, t_start = self.__initialize_timesteps(self.denoising_steps, strength)1063 1064            # at which timestep to set the initial noise (n.b. 50% if strength is 0.5)1065            latent_timestep = timesteps[:1].repeat(batch_size)1066            # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise1067            is_strength_max = strength == 1.01068 1069            # Pre-initialize latents1070            num_channels_latents = self.vae.config.latent_channels1071            latents_outputs = self.prepare_latents(1072                batch_size,1073                num_channels_latents,1074                self.image_height,1075                self.image_width,1076                torch.float32,1077                self.torch_device,1078                generator,1079                image=init_image,1080                timestep=latent_timestep,1081                is_strength_max=is_strength_max,1082            )1083 1084            latents = latents_outputs[0]1085 1086            # VAE encode masked image1087            masked_latents = self.__encode_image(masked_image)1088            masked_latents = torch.cat([masked_latents] * 2)1089 1090            # CLIP text encoder1091            text_embeddings = self.__encode_prompt(prompt, negative_prompt)1092 1093            # UNet denoiser1094            latents = self.__denoise_latent(1095                latents,1096                text_embeddings,1097                timesteps=timesteps,1098                step_offset=t_start,1099                mask=mask,1100                masked_image_latents=masked_latents,1101            )1102 1103            # VAE decode latent1104            images = self.__decode_latent(latents)1105 1106        images = self.numpy_to_pil(images)1107        return StableDiffusionPipelineOutput(images=images, nsfw_content_detected=None)1108