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    StableDiffusionImg2ImgPipeline,49    StableDiffusionPipelineOutput,50    StableDiffusionSafetyChecker,51)52from diffusers.schedulers import DDIMScheduler53from diffusers.utils import DIFFUSERS_CACHE, logging54 55 56"""57Installation instructions58python3 -m pip install --upgrade transformers diffusers>=0.16.059python3 -m pip install --upgrade tensorrt>=8.6.160python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com61python3 -m pip install onnxruntime62"""63 64TRT_LOGGER = trt.Logger(trt.Logger.ERROR)65logger = logging.get_logger(__name__)  # pylint: disable=invalid-name66 67# Map of numpy dtype -> torch dtype68numpy_to_torch_dtype_dict = {69    np.uint8: torch.uint8,70    np.int8: torch.int8,71    np.int16: torch.int16,72    np.int32: torch.int32,73    np.int64: torch.int64,74    np.float16: torch.float16,75    np.float32: torch.float32,76    np.float64: torch.float64,77    np.complex64: torch.complex64,78    np.complex128: torch.complex128,79}80if np.version.full_version >= "1.24.0":81    numpy_to_torch_dtype_dict[np.bool_] = torch.bool82else:83    numpy_to_torch_dtype_dict[np.bool] = torch.bool84 85# Map of torch dtype -> numpy dtype86torch_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()}87 88 89def device_view(t):90    return cuda.DeviceView(ptr=t.data_ptr(), shape=t.shape, dtype=torch_to_numpy_dtype_dict[t.dtype])91 92 93def preprocess_image(image):94    """95    image: torch.Tensor96    """97    w, h = image.size98    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 3299    image = image.resize((w, h))100    image = np.array(image).astype(np.float32) / 255.0101    image = image[None].transpose(0, 3, 1, 2)102    image = torch.from_numpy(image).contiguous()103    return 2.0 * image - 1.0104 105 106class Engine:107    def __init__(self, engine_path):108        self.engine_path = engine_path109        self.engine = None110        self.context = None111        self.buffers = OrderedDict()112        self.tensors = OrderedDict()113 114    def __del__(self):115        [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray)]116        del self.engine117        del self.context118        del self.buffers119        del self.tensors120 121    def build(122        self,123        onnx_path,124        fp16,125        input_profile=None,126        enable_preview=False,127        enable_all_tactics=False,128        timing_cache=None,129        workspace_size=0,130    ):131        logger.warning(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")132        p = Profile()133        if input_profile:134            for name, dims in input_profile.items():135                assert len(dims) == 3136                p.add(name, min=dims[0], opt=dims[1], max=dims[2])137 138        config_kwargs = {}139 140        config_kwargs["preview_features"] = [trt.PreviewFeature.DISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805]141        if enable_preview:142            # Faster dynamic shapes made optional since it increases engine build time.143            config_kwargs["preview_features"].append(trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805)144        if workspace_size > 0:145            config_kwargs["memory_pool_limits"] = {trt.MemoryPoolType.WORKSPACE: workspace_size}146        if not enable_all_tactics:147            config_kwargs["tactic_sources"] = []148 149        engine = engine_from_network(150            network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]),151            config=CreateConfig(fp16=fp16, profiles=[p], load_timing_cache=timing_cache, **config_kwargs),152            save_timing_cache=timing_cache,153        )154        save_engine(engine, path=self.engine_path)155 156    def load(self):157        logger.warning(f"Loading TensorRT engine: {self.engine_path}")158        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))159 160    def activate(self):161        self.context = self.engine.create_execution_context()162 163    def allocate_buffers(self, shape_dict=None, device="cuda"):164        for idx in range(trt_util.get_bindings_per_profile(self.engine)):165            binding = self.engine[idx]166            if shape_dict and binding in shape_dict:167                shape = shape_dict[binding]168            else:169                shape = self.engine.get_binding_shape(binding)170            dtype = trt.nptype(self.engine.get_binding_dtype(binding))171            if self.engine.binding_is_input(binding):172                self.context.set_binding_shape(idx, shape)173            tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device)174            self.tensors[binding] = tensor175            self.buffers[binding] = cuda.DeviceView(ptr=tensor.data_ptr(), shape=shape, dtype=dtype)176 177    def infer(self, feed_dict, stream):178        start_binding, end_binding = trt_util.get_active_profile_bindings(self.context)179        # shallow copy of ordered dict180        device_buffers = copy(self.buffers)181        for name, buf in feed_dict.items():182            assert isinstance(buf, cuda.DeviceView)183            device_buffers[name] = buf184        bindings = [0] * start_binding + [buf.ptr for buf in device_buffers.values()]185        noerror = self.context.execute_async_v2(bindings=bindings, stream_handle=stream.ptr)186        if not noerror:187            raise ValueError("ERROR: inference failed.")188 189        return self.tensors190 191 192class Optimizer:193    def __init__(self, onnx_graph):194        self.graph = gs.import_onnx(onnx_graph)195 196    def cleanup(self, return_onnx=False):197        self.graph.cleanup().toposort()198        if return_onnx:199            return gs.export_onnx(self.graph)200 201    def select_outputs(self, keep, names=None):202        self.graph.outputs = [self.graph.outputs[o] for o in keep]203        if names:204            for i, name in enumerate(names):205                self.graph.outputs[i].name = name206 207    def fold_constants(self, return_onnx=False):208        onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)209        self.graph = gs.import_onnx(onnx_graph)210        if return_onnx:211            return onnx_graph212 213    def infer_shapes(self, return_onnx=False):214        onnx_graph = gs.export_onnx(self.graph)215        if onnx_graph.ByteSize() > 2147483648:216            raise TypeError("ERROR: model size exceeds supported 2GB limit")217        else:218            onnx_graph = shape_inference.infer_shapes(onnx_graph)219 220        self.graph = gs.import_onnx(onnx_graph)221        if return_onnx:222            return onnx_graph223 224 225class BaseModel:226    def __init__(self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77):227        self.model = model228        self.name = "SD Model"229        self.fp16 = fp16230        self.device = device231 232        self.min_batch = 1233        self.max_batch = max_batch_size234        self.min_image_shape = 256  # min image resolution: 256x256235        self.max_image_shape = 1024  # max image resolution: 1024x1024236        self.min_latent_shape = self.min_image_shape // 8237        self.max_latent_shape = self.max_image_shape // 8238 239        self.embedding_dim = embedding_dim240        self.text_maxlen = text_maxlen241 242    def get_model(self):243        return self.model244 245    def get_input_names(self):246        pass247 248    def get_output_names(self):249        pass250 251    def get_dynamic_axes(self):252        return None253 254    def get_sample_input(self, batch_size, image_height, image_width):255        pass256 257    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):258        return None259 260    def get_shape_dict(self, batch_size, image_height, image_width):261        return None262 263    def optimize(self, onnx_graph):264        opt = Optimizer(onnx_graph)265        opt.cleanup()266        opt.fold_constants()267        opt.infer_shapes()268        onnx_opt_graph = opt.cleanup(return_onnx=True)269        return onnx_opt_graph270 271    def check_dims(self, batch_size, image_height, image_width):272        assert batch_size >= self.min_batch and batch_size <= self.max_batch273        assert image_height % 8 == 0 or image_width % 8 == 0274        latent_height = image_height // 8275        latent_width = image_width // 8276        assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape277        assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape278        return (latent_height, latent_width)279 280    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):281        min_batch = batch_size if static_batch else self.min_batch282        max_batch = batch_size if static_batch else self.max_batch283        latent_height = image_height // 8284        latent_width = image_width // 8285        min_image_height = image_height if static_shape else self.min_image_shape286        max_image_height = image_height if static_shape else self.max_image_shape287        min_image_width = image_width if static_shape else self.min_image_shape288        max_image_width = image_width if static_shape else self.max_image_shape289        min_latent_height = latent_height if static_shape else self.min_latent_shape290        max_latent_height = latent_height if static_shape else self.max_latent_shape291        min_latent_width = latent_width if static_shape else self.min_latent_shape292        max_latent_width = latent_width if static_shape else self.max_latent_shape293        return (294            min_batch,295            max_batch,296            min_image_height,297            max_image_height,298            min_image_width,299            max_image_width,300            min_latent_height,301            max_latent_height,302            min_latent_width,303            max_latent_width,304        )305 306 307def getOnnxPath(model_name, onnx_dir, opt=True):308    return os.path.join(onnx_dir, model_name + (".opt" if opt else "") + ".onnx")309 310 311def getEnginePath(model_name, engine_dir):312    return os.path.join(engine_dir, model_name + ".plan")313 314 315def build_engines(316    models: dict,317    engine_dir,318    onnx_dir,319    onnx_opset,320    opt_image_height,321    opt_image_width,322    opt_batch_size=1,323    force_engine_rebuild=False,324    static_batch=False,325    static_shape=True,326    enable_preview=False,327    enable_all_tactics=False,328    timing_cache=None,329    max_workspace_size=0,330):331    built_engines = {}332    if not os.path.isdir(onnx_dir):333        os.makedirs(onnx_dir)334    if not os.path.isdir(engine_dir):335        os.makedirs(engine_dir)336 337    # Export models to ONNX338    for model_name, model_obj in models.items():339        engine_path = getEnginePath(model_name, engine_dir)340        if force_engine_rebuild or not os.path.exists(engine_path):341            logger.warning("Building Engines...")342            logger.warning("Engine build can take a while to complete")343            onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)344            onnx_opt_path = getOnnxPath(model_name, onnx_dir)345            if force_engine_rebuild or not os.path.exists(onnx_opt_path):346                if force_engine_rebuild or not os.path.exists(onnx_path):347                    logger.warning(f"Exporting model: {onnx_path}")348                    model = model_obj.get_model()349                    with torch.inference_mode(), torch.autocast("cuda"):350                        inputs = model_obj.get_sample_input(opt_batch_size, opt_image_height, opt_image_width)351                        torch.onnx.export(352                            model,353                            inputs,354                            onnx_path,355                            export_params=True,356                            opset_version=onnx_opset,357                            do_constant_folding=True,358                            input_names=model_obj.get_input_names(),359                            output_names=model_obj.get_output_names(),360                            dynamic_axes=model_obj.get_dynamic_axes(),361                        )362                    del model363                    torch.cuda.empty_cache()364                    gc.collect()365                else:366                    logger.warning(f"Found cached model: {onnx_path}")367 368                # Optimize onnx369                if force_engine_rebuild or not os.path.exists(onnx_opt_path):370                    logger.warning(f"Generating optimizing model: {onnx_opt_path}")371                    onnx_opt_graph = model_obj.optimize(onnx.load(onnx_path))372                    onnx.save(onnx_opt_graph, onnx_opt_path)373                else:374                    logger.warning(f"Found cached optimized model: {onnx_opt_path} ")375 376    # Build TensorRT engines377    for model_name, model_obj in models.items():378        engine_path = getEnginePath(model_name, engine_dir)379        engine = Engine(engine_path)380        onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)381        onnx_opt_path = getOnnxPath(model_name, onnx_dir)382 383        if force_engine_rebuild or not os.path.exists(engine.engine_path):384            engine.build(385                onnx_opt_path,386                fp16=True,387                input_profile=model_obj.get_input_profile(388                    opt_batch_size,389                    opt_image_height,390                    opt_image_width,391                    static_batch=static_batch,392                    static_shape=static_shape,393                ),394                enable_preview=enable_preview,395                timing_cache=timing_cache,396                workspace_size=max_workspace_size,397            )398        built_engines[model_name] = engine399 400    # Load and activate TensorRT engines401    for model_name, model_obj in models.items():402        engine = built_engines[model_name]403        engine.load()404        engine.activate()405 406    return built_engines407 408 409def runEngine(engine, feed_dict, stream):410    return engine.infer(feed_dict, stream)411 412 413class CLIP(BaseModel):414    def __init__(self, model, device, max_batch_size, embedding_dim):415        super(CLIP, self).__init__(416            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim417        )418        self.name = "CLIP"419 420    def get_input_names(self):421        return ["input_ids"]422 423    def get_output_names(self):424        return ["text_embeddings", "pooler_output"]425 426    def get_dynamic_axes(self):427        return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}428 429    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):430        self.check_dims(batch_size, image_height, image_width)431        min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(432            batch_size, image_height, image_width, static_batch, static_shape433        )434        return {435            "input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]436        }437 438    def get_shape_dict(self, batch_size, image_height, image_width):439        self.check_dims(batch_size, image_height, image_width)440        return {441            "input_ids": (batch_size, self.text_maxlen),442            "text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),443        }444 445    def get_sample_input(self, batch_size, image_height, image_width):446        self.check_dims(batch_size, image_height, image_width)447        return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)448 449    def optimize(self, onnx_graph):450        opt = Optimizer(onnx_graph)451        opt.select_outputs([0])  # delete graph output#1452        opt.cleanup()453        opt.fold_constants()454        opt.infer_shapes()455        opt.select_outputs([0], names=["text_embeddings"])  # rename network output456        opt_onnx_graph = opt.cleanup(return_onnx=True)457        return opt_onnx_graph458 459 460def make_CLIP(model, device, max_batch_size, embedding_dim, inpaint=False):461    return CLIP(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)462 463 464class UNet(BaseModel):465    def __init__(466        self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77, unet_dim=4467    ):468        super(UNet, self).__init__(469            model=model,470            fp16=fp16,471            device=device,472            max_batch_size=max_batch_size,473            embedding_dim=embedding_dim,474            text_maxlen=text_maxlen,475        )476        self.unet_dim = unet_dim477        self.name = "UNet"478 479    def get_input_names(self):480        return ["sample", "timestep", "encoder_hidden_states"]481 482    def get_output_names(self):483        return ["latent"]484 485    def get_dynamic_axes(self):486        return {487            "sample": {0: "2B", 2: "H", 3: "W"},488            "encoder_hidden_states": {0: "2B"},489            "latent": {0: "2B", 2: "H", 3: "W"},490        }491 492    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):493        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)494        (495            min_batch,496            max_batch,497            _,498            _,499            _,500            _,501            min_latent_height,502            max_latent_height,503            min_latent_width,504            max_latent_width,505        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)506        return {507            "sample": [508                (2 * min_batch, self.unet_dim, min_latent_height, min_latent_width),509                (2 * batch_size, self.unet_dim, latent_height, latent_width),510                (2 * max_batch, self.unet_dim, max_latent_height, max_latent_width),511            ],512            "encoder_hidden_states": [513                (2 * min_batch, self.text_maxlen, self.embedding_dim),514                (2 * batch_size, self.text_maxlen, self.embedding_dim),515                (2 * max_batch, self.text_maxlen, self.embedding_dim),516            ],517        }518 519    def get_shape_dict(self, batch_size, image_height, image_width):520        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)521        return {522            "sample": (2 * batch_size, self.unet_dim, latent_height, latent_width),523            "encoder_hidden_states": (2 * batch_size, self.text_maxlen, self.embedding_dim),524            "latent": (2 * batch_size, 4, latent_height, latent_width),525        }526 527    def get_sample_input(self, batch_size, image_height, image_width):528        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)529        dtype = torch.float16 if self.fp16 else torch.float32530        return (531            torch.randn(532                2 * batch_size, self.unet_dim, latent_height, latent_width, dtype=torch.float32, device=self.device533            ),534            torch.tensor([1.0], dtype=torch.float32, device=self.device),535            torch.randn(2 * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),536        )537 538 539def make_UNet(model, device, max_batch_size, embedding_dim, inpaint=False):540    return UNet(541        model,542        fp16=True,543        device=device,544        max_batch_size=max_batch_size,545        embedding_dim=embedding_dim,546        unet_dim=(9 if inpaint else 4),547    )548 549 550class VAE(BaseModel):551    def __init__(self, model, device, max_batch_size, embedding_dim):552        super(VAE, self).__init__(553            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim554        )555        self.name = "VAE decoder"556 557    def get_input_names(self):558        return ["latent"]559 560    def get_output_names(self):561        return ["images"]562 563    def get_dynamic_axes(self):564        return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}565 566    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):567        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)568        (569            min_batch,570            max_batch,571            _,572            _,573            _,574            _,575            min_latent_height,576            max_latent_height,577            min_latent_width,578            max_latent_width,579        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)580        return {581            "latent": [582                (min_batch, 4, min_latent_height, min_latent_width),583                (batch_size, 4, latent_height, latent_width),584                (max_batch, 4, max_latent_height, max_latent_width),585            ]586        }587 588    def get_shape_dict(self, batch_size, image_height, image_width):589        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)590        return {591            "latent": (batch_size, 4, latent_height, latent_width),592            "images": (batch_size, 3, image_height, image_width),593        }594 595    def get_sample_input(self, batch_size, image_height, image_width):596        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)597        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)598 599 600def make_VAE(model, device, max_batch_size, embedding_dim, inpaint=False):601    return VAE(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)602 603 604class TorchVAEEncoder(torch.nn.Module):605    def __init__(self, model):606        super().__init__()607        self.vae_encoder = model608 609    def forward(self, x):610        return self.vae_encoder.encode(x).latent_dist.sample()611 612 613class VAEEncoder(BaseModel):614    def __init__(self, model, device, max_batch_size, embedding_dim):615        super(VAEEncoder, self).__init__(616            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim617        )618        self.name = "VAE encoder"619 620    def get_model(self):621        vae_encoder = TorchVAEEncoder(self.model)622        return vae_encoder623 624    def get_input_names(self):625        return ["images"]626 627    def get_output_names(self):628        return ["latent"]629 630    def get_dynamic_axes(self):631        return {"images": {0: "B", 2: "8H", 3: "8W"}, "latent": {0: "B", 2: "H", 3: "W"}}632 633    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):634        assert batch_size >= self.min_batch and batch_size <= self.max_batch635        min_batch = batch_size if static_batch else self.min_batch636        max_batch = batch_size if static_batch else self.max_batch637        self.check_dims(batch_size, image_height, image_width)638        (639            min_batch,640            max_batch,641            min_image_height,642            max_image_height,643            min_image_width,644            max_image_width,645            _,646            _,647            _,648            _,649        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)650 651        return {652            "images": [653                (min_batch, 3, min_image_height, min_image_width),654                (batch_size, 3, image_height, image_width),655                (max_batch, 3, max_image_height, max_image_width),656            ]657        }658 659    def get_shape_dict(self, batch_size, image_height, image_width):660        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)661        return {662            "images": (batch_size, 3, image_height, image_width),663            "latent": (batch_size, 4, latent_height, latent_width),664        }665 666    def get_sample_input(self, batch_size, image_height, image_width):667        self.check_dims(batch_size, image_height, image_width)668        return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device)669 670 671def make_VAEEncoder(model, device, max_batch_size, embedding_dim, inpaint=False):672    return VAEEncoder(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)673 674 675class TensorRTStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):676    r"""677    Pipeline for image-to-image generation using TensorRT accelerated Stable Diffusion.678 679    This model inherits from [`StableDiffusionImg2ImgPipeline`]. Check the superclass documentation for the generic methods the680    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)681 682    Args:683        vae ([`AutoencoderKL`]):684            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.685        text_encoder ([`CLIPTextModel`]):686            Frozen text-encoder. Stable Diffusion uses the text portion of687            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically688            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.689        tokenizer (`CLIPTokenizer`):690            Tokenizer of class691            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).692        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.693        scheduler ([`SchedulerMixin`]):694            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of695            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].696        safety_checker ([`StableDiffusionSafetyChecker`]):697            Classification module that estimates whether generated images could be considered offensive or harmful.698            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.699        feature_extractor ([`CLIPFeatureExtractor`]):700            Model that extracts features from generated images to be used as inputs for the `safety_checker`.701    """702 703    def __init__(704        self,705        vae: AutoencoderKL,706        text_encoder: CLIPTextModel,707        tokenizer: CLIPTokenizer,708        unet: UNet2DConditionModel,709        scheduler: DDIMScheduler,710        safety_checker: StableDiffusionSafetyChecker,711        feature_extractor: CLIPFeatureExtractor,712        requires_safety_checker: bool = True,713        stages=["clip", "unet", "vae", "vae_encoder"],714        image_height: int = 512,715        image_width: int = 512,716        max_batch_size: int = 16,717        # ONNX export parameters718        onnx_opset: int = 17,719        onnx_dir: str = "onnx",720        # TensorRT engine build parameters721        engine_dir: str = "engine",722        build_preview_features: bool = True,723        force_engine_rebuild: bool = False,724        timing_cache: str = "timing_cache",725    ):726        super().__init__(727            vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker728        )729 730        self.vae.forward = self.vae.decode731 732        self.stages = stages733        self.image_height, self.image_width = image_height, image_width734        self.inpaint = False735        self.onnx_opset = onnx_opset736        self.onnx_dir = onnx_dir737        self.engine_dir = engine_dir738        self.force_engine_rebuild = force_engine_rebuild739        self.timing_cache = timing_cache740        self.build_static_batch = False741        self.build_dynamic_shape = False742        self.build_preview_features = build_preview_features743 744        self.max_batch_size = max_batch_size745        # TODO: Restrict batch size to 4 for larger image dimensions as a WAR for TensorRT limitation.746        if self.build_dynamic_shape or self.image_height > 512 or self.image_width > 512:747            self.max_batch_size = 4748 749        self.stream = None  # loaded in loadResources()750        self.models = {}  # loaded in __loadModels()751        self.engine = {}  # loaded in build_engines()752 753    def __loadModels(self):754        # Load pipeline models755        self.embedding_dim = self.text_encoder.config.hidden_size756        models_args = {757            "device": self.torch_device,758            "max_batch_size": self.max_batch_size,759            "embedding_dim": self.embedding_dim,760            "inpaint": self.inpaint,761        }762        if "clip" in self.stages:763            self.models["clip"] = make_CLIP(self.text_encoder, **models_args)764        if "unet" in self.stages:765            self.models["unet"] = make_UNet(self.unet, **models_args)766        if "vae" in self.stages:767            self.models["vae"] = make_VAE(self.vae, **models_args)768        if "vae_encoder" in self.stages:769            self.models["vae_encoder"] = make_VAEEncoder(self.vae, **models_args)770 771    @classmethod772    def set_cached_folder(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):773        cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)774        resume_download = kwargs.pop("resume_download", False)775        proxies = kwargs.pop("proxies", None)776        local_files_only = kwargs.pop("local_files_only", False)777        use_auth_token = kwargs.pop("use_auth_token", None)778        revision = kwargs.pop("revision", None)779 780        cls.cached_folder = (781            pretrained_model_name_or_path782            if os.path.isdir(pretrained_model_name_or_path)783            else snapshot_download(784                pretrained_model_name_or_path,785                cache_dir=cache_dir,786                resume_download=resume_download,787                proxies=proxies,788                local_files_only=local_files_only,789                use_auth_token=use_auth_token,790                revision=revision,791            )792        )793 794    def to(self, torch_device: Optional[Union[str, torch.device]] = None, silence_dtype_warnings: bool = False):795        super().to(torch_device, silence_dtype_warnings=silence_dtype_warnings)796 797        self.onnx_dir = os.path.join(self.cached_folder, self.onnx_dir)798        self.engine_dir = os.path.join(self.cached_folder, self.engine_dir)799        self.timing_cache = os.path.join(self.cached_folder, self.timing_cache)800 801        # set device802        self.torch_device = self._execution_device803        logger.warning(f"Running inference on device: {self.torch_device}")804 805        # load models806        self.__loadModels()807 808        # build engines809        self.engine = build_engines(810            self.models,811            self.engine_dir,812            self.onnx_dir,813            self.onnx_opset,814            opt_image_height=self.image_height,815            opt_image_width=self.image_width,816            force_engine_rebuild=self.force_engine_rebuild,817            static_batch=self.build_static_batch,818            static_shape=not self.build_dynamic_shape,819            enable_preview=self.build_preview_features,820            timing_cache=self.timing_cache,821        )822 823        return self824 825    def __initialize_timesteps(self, timesteps, strength):826        self.scheduler.set_timesteps(timesteps)827        offset = self.scheduler.steps_offset if hasattr(self.scheduler, "steps_offset") else 0828        init_timestep = int(timesteps * strength) + offset829        init_timestep = min(init_timestep, timesteps)830        t_start = max(timesteps - init_timestep + offset, 0)831        timesteps = self.scheduler.timesteps[t_start:].to(self.torch_device)832        return timesteps, t_start833 834    def __preprocess_images(self, batch_size, images=()):835        init_images = []836        for image in images:837            image = image.to(self.torch_device).float()838            image = image.repeat(batch_size, 1, 1, 1)839            init_images.append(image)840        return tuple(init_images)841 842    def __encode_image(self, init_image):843        init_latents = runEngine(self.engine["vae_encoder"], {"images": device_view(init_image)}, self.stream)[844            "latent"845        ]846        init_latents = 0.18215 * init_latents847        return init_latents848 849    def __encode_prompt(self, prompt, negative_prompt):850        r"""851        Encodes the prompt into text encoder hidden states.852 853        Args:854             prompt (`str` or `List[str]`, *optional*):855                prompt to be encoded856            negative_prompt (`str` or `List[str]`, *optional*):857                The prompt or prompts not to guide the image generation. If not defined, one has to pass858                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.859                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).860        """861        # Tokenize prompt862        text_input_ids = (863            self.tokenizer(864                prompt,865                padding="max_length",866                max_length=self.tokenizer.model_max_length,867                truncation=True,868                return_tensors="pt",869            )870            .input_ids.type(torch.int32)871            .to(self.torch_device)872        )873 874        text_input_ids_inp = device_view(text_input_ids)875        # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt876        text_embeddings = runEngine(self.engine["clip"], {"input_ids": text_input_ids_inp}, self.stream)[877            "text_embeddings"878        ].clone()879 880        # Tokenize negative prompt881        uncond_input_ids = (882            self.tokenizer(883                negative_prompt,884                padding="max_length",885                max_length=self.tokenizer.model_max_length,886                truncation=True,887                return_tensors="pt",888            )889            .input_ids.type(torch.int32)890            .to(self.torch_device)891        )892        uncond_input_ids_inp = device_view(uncond_input_ids)893        uncond_embeddings = runEngine(self.engine["clip"], {"input_ids": uncond_input_ids_inp}, self.stream)[894            "text_embeddings"895        ]896 897        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance898        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)899 900        return text_embeddings901 902    def __denoise_latent(903        self, latents, text_embeddings, timesteps=None, step_offset=0, mask=None, masked_image_latents=None904    ):905        if not isinstance(timesteps, torch.Tensor):906            timesteps = self.scheduler.timesteps907        for step_index, timestep in enumerate(timesteps):908            # Expand the latents if we are doing classifier free guidance909            latent_model_input = torch.cat([latents] * 2)910            latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)911            if isinstance(mask, torch.Tensor):912                latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)913 914            # Predict the noise residual915            timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep916 917            sample_inp = device_view(latent_model_input)918            timestep_inp = device_view(timestep_float)919            embeddings_inp = device_view(text_embeddings)920            noise_pred = runEngine(921                self.engine["unet"],922                {"sample": sample_inp, "timestep": timestep_inp, "encoder_hidden_states": embeddings_inp},923                self.stream,924            )["latent"]925 926            # Perform guidance927            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)928            noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)929 930            latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample931 932        latents = 1.0 / 0.18215 * latents933        return latents934 935    def __decode_latent(self, latents):936        images = runEngine(self.engine["vae"], {"latent": device_view(latents)}, self.stream)["images"]937        images = (images / 2 + 0.5).clamp(0, 1)938        return images.cpu().permute(0, 2, 3, 1).float().numpy()939 940    def __loadResources(self, image_height, image_width, batch_size):941        self.stream = cuda.Stream()942 943        # Allocate buffers for TensorRT engine bindings944        for model_name, obj in self.models.items():945            self.engine[model_name].allocate_buffers(946                shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.torch_device947            )948 949    @torch.no_grad()950    def __call__(951        self,952        prompt: Union[str, List[str]] = None,953        image: Union[torch.FloatTensor, PIL.Image.Image] = None,954        strength: float = 0.8,955        num_inference_steps: int = 50,956        guidance_scale: float = 7.5,957        negative_prompt: Optional[Union[str, List[str]]] = None,958        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,959    ):960        r"""961        Function invoked when calling the pipeline for generation.962 963        Args:964            prompt (`str` or `List[str]`, *optional*):965                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.966                instead.967            image (`PIL.Image.Image`):968                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will969                be masked out with `mask_image` and repainted according to `prompt`.970            strength (`float`, *optional*, defaults to 0.8):971                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`972                will be used as a starting point, adding more noise to it the larger the `strength`. The number of973                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will974                be maximum and the denoising process will run for the full number of iterations specified in975                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.976            num_inference_steps (`int`, *optional*, defaults to 50):977                The number of denoising steps. More denoising steps usually lead to a higher quality image at the978                expense of slower inference.979            guidance_scale (`float`, *optional*, defaults to 7.5):980                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).981                `guidance_scale` is defined as `w` of equation 2. of [Imagen982                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >983                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,984                usually at the expense of lower image quality.985            negative_prompt (`str` or `List[str]`, *optional*):986                The prompt or prompts not to guide the image generation. If not defined, one has to pass987                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.988                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).989            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):990                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)991                to make generation deterministic.992 993        """994        self.generator = generator995        self.denoising_steps = num_inference_steps996        self.guidance_scale = guidance_scale997 998        # Pre-compute latent input scales and linear multistep coefficients999        self.scheduler.set_timesteps(self.denoising_steps, device=self.torch_device)1000 1001        # Define call parameters1002        if prompt is not None and isinstance(prompt, str):1003            batch_size = 11004            prompt = [prompt]1005        elif prompt is not None and isinstance(prompt, list):1006            batch_size = len(prompt)1007        else:1008            raise ValueError(f"Expected prompt to be of type list or str but got {type(prompt)}")1009 1010        if negative_prompt is None:1011            negative_prompt = [""] * batch_size1012 1013        if negative_prompt is not None and isinstance(negative_prompt, str):1014            negative_prompt = [negative_prompt]1015 1016        assert len(prompt) == len(negative_prompt)1017 1018        if batch_size > self.max_batch_size:1019            raise ValueError(1020                f"Batch size {len(prompt)} is larger than allowed {self.max_batch_size}. If dynamic shape is used, then maximum batch size is 4"1021            )1022 1023        # load resources1024        self.__loadResources(self.image_height, self.image_width, batch_size)1025 1026        with torch.inference_mode(), torch.autocast("cuda"), trt.Runtime(TRT_LOGGER):1027            # Initialize timesteps1028            timesteps, t_start = self.__initialize_timesteps(self.denoising_steps, strength)1029            latent_timestep = timesteps[:1].repeat(batch_size)1030 1031            # Pre-process input image1032            if isinstance(image, PIL.Image.Image):1033                image = preprocess_image(image)1034            init_image = self.__preprocess_images(batch_size, (image,))[0]1035 1036            # VAE encode init image1037            init_latents = self.__encode_image(init_image)1038 1039            # Add noise to latents using timesteps1040            noise = torch.randn(1041                init_latents.shape, generator=self.generator, device=self.torch_device, dtype=torch.float321042            )1043            latents = self.scheduler.add_noise(init_latents, noise, latent_timestep)1044 1045            # CLIP text encoder1046            text_embeddings = self.__encode_prompt(prompt, negative_prompt)1047 1048            # UNet denoiser1049            latents = self.__denoise_latent(latents, text_embeddings, timesteps=timesteps, step_offset=t_start)1050 1051            # VAE decode latent1052            images = self.__decode_latent(latents)1053 1054        images = self.numpy_to_pil(images)1055        return StableDiffusionPipelineOutput(images=images, nsfw_content_detected=None)1056