diffusers/community-pipelines-mirror

1#2# Copyright 2024 The HuggingFace Inc. team.3# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.4# SPDX-License-Identifier: Apache-2.05#6# Licensed under the Apache License, Version 2.0 (the "License");7# you may not use this file except in compliance with the License.8# You may obtain a copy of the License at9#10# http://www.apache.org/licenses/LICENSE-2.011#12# Unless required by applicable law or agreed to in writing, software13# distributed under the License is distributed on an "AS IS" BASIS,14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15# See the License for the specific language governing permissions and16# limitations under the License.17 18import gc19import os20from collections import OrderedDict21from typing import List, Optional, Tuple, Union22 23import numpy as np24import onnx25import onnx_graphsurgeon as gs26import PIL.Image27import tensorrt as trt28import torch29from cuda import cudart30from huggingface_hub import snapshot_download31from huggingface_hub.utils import validate_hf_hub_args32from onnx import shape_inference33from packaging import version34from polygraphy import cuda35from polygraphy.backend.common import bytes_from_path36from polygraphy.backend.onnx.loader import fold_constants37from polygraphy.backend.trt import (38    CreateConfig,39    Profile,40    engine_from_bytes,41    engine_from_network,42    network_from_onnx_path,43    save_engine,44)45from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection46 47from diffusers import DiffusionPipeline48from diffusers.configuration_utils import FrozenDict, deprecate49from diffusers.image_processor import VaeImageProcessor50from diffusers.models import AutoencoderKL, UNet2DConditionModel51from diffusers.pipelines.stable_diffusion import (52    StableDiffusionPipelineOutput,53    StableDiffusionSafetyChecker,54)55from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents56from diffusers.schedulers import DDIMScheduler57from diffusers.utils import logging58 59 60"""61Installation instructions62python3 -m pip install --upgrade transformers diffusers>=0.16.063python3 -m pip install --upgrade tensorrt~=10.2.064python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com65python3 -m pip install onnxruntime66"""67 68TRT_LOGGER = trt.Logger(trt.Logger.ERROR)69logger = logging.get_logger(__name__)  # pylint: disable=invalid-name70 71# Map of numpy dtype -> torch dtype72numpy_to_torch_dtype_dict = {73    np.uint8: torch.uint8,74    np.int8: torch.int8,75    np.int16: torch.int16,76    np.int32: torch.int32,77    np.int64: torch.int64,78    np.float16: torch.float16,79    np.float32: torch.float32,80    np.float64: torch.float64,81    np.complex64: torch.complex64,82    np.complex128: torch.complex128,83}84if np.version.full_version >= "1.24.0":85    numpy_to_torch_dtype_dict[np.bool_] = torch.bool86else:87    numpy_to_torch_dtype_dict[np.bool] = torch.bool88 89# Map of torch dtype -> numpy dtype90torch_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()}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_all_tactics=False,127        timing_cache=None,128    ):129        logger.warning(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")130        p = Profile()131        if input_profile:132            for name, dims in input_profile.items():133                assert len(dims) == 3134                p.add(name, min=dims[0], opt=dims[1], max=dims[2])135 136        extra_build_args = {}137        if not enable_all_tactics:138            extra_build_args["tactic_sources"] = []139 140        engine = engine_from_network(141            network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]),142            config=CreateConfig(fp16=fp16, profiles=[p], load_timing_cache=timing_cache, **extra_build_args),143            save_timing_cache=timing_cache,144        )145        save_engine(engine, path=self.engine_path)146 147    def load(self):148        logger.warning(f"Loading TensorRT engine: {self.engine_path}")149        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))150 151    def activate(self):152        self.context = self.engine.create_execution_context()153 154    def allocate_buffers(self, shape_dict=None, device="cuda"):155        for binding in range(self.engine.num_io_tensors):156            name = self.engine.get_tensor_name(binding)157            if shape_dict and name in shape_dict:158                shape = shape_dict[name]159            else:160                shape = self.engine.get_tensor_shape(name)161            dtype = trt.nptype(self.engine.get_tensor_dtype(name))162            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:163                self.context.set_input_shape(name, shape)164            tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device)165            self.tensors[name] = tensor166 167    def infer(self, feed_dict, stream):168        for name, buf in feed_dict.items():169            self.tensors[name].copy_(buf)170        for name, tensor in self.tensors.items():171            self.context.set_tensor_address(name, tensor.data_ptr())172        noerror = self.context.execute_async_v3(stream)173        if not noerror:174            raise ValueError("ERROR: inference failed.")175 176        return self.tensors177 178 179class Optimizer:180    def __init__(self, onnx_graph):181        self.graph = gs.import_onnx(onnx_graph)182 183    def cleanup(self, return_onnx=False):184        self.graph.cleanup().toposort()185        if return_onnx:186            return gs.export_onnx(self.graph)187 188    def select_outputs(self, keep, names=None):189        self.graph.outputs = [self.graph.outputs[o] for o in keep]190        if names:191            for i, name in enumerate(names):192                self.graph.outputs[i].name = name193 194    def fold_constants(self, return_onnx=False):195        onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)196        self.graph = gs.import_onnx(onnx_graph)197        if return_onnx:198            return onnx_graph199 200    def infer_shapes(self, return_onnx=False):201        onnx_graph = gs.export_onnx(self.graph)202        if onnx_graph.ByteSize() > 2147483648:203            raise TypeError("ERROR: model size exceeds supported 2GB limit")204        else:205            onnx_graph = shape_inference.infer_shapes(onnx_graph)206 207        self.graph = gs.import_onnx(onnx_graph)208        if return_onnx:209            return onnx_graph210 211 212class BaseModel:213    def __init__(self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77):214        self.model = model215        self.name = "SD Model"216        self.fp16 = fp16217        self.device = device218 219        self.min_batch = 1220        self.max_batch = max_batch_size221        self.min_image_shape = 256  # min image resolution: 256x256222        self.max_image_shape = 1024  # max image resolution: 1024x1024223        self.min_latent_shape = self.min_image_shape // 8224        self.max_latent_shape = self.max_image_shape // 8225 226        self.embedding_dim = embedding_dim227        self.text_maxlen = text_maxlen228 229    def get_model(self):230        return self.model231 232    def get_input_names(self):233        pass234 235    def get_output_names(self):236        pass237 238    def get_dynamic_axes(self):239        return None240 241    def get_sample_input(self, batch_size, image_height, image_width):242        pass243 244    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):245        return None246 247    def get_shape_dict(self, batch_size, image_height, image_width):248        return None249 250    def optimize(self, onnx_graph):251        opt = Optimizer(onnx_graph)252        opt.cleanup()253        opt.fold_constants()254        opt.infer_shapes()255        onnx_opt_graph = opt.cleanup(return_onnx=True)256        return onnx_opt_graph257 258    def check_dims(self, batch_size, image_height, image_width):259        assert batch_size >= self.min_batch and batch_size <= self.max_batch260        assert image_height % 8 == 0 or image_width % 8 == 0261        latent_height = image_height // 8262        latent_width = image_width // 8263        assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape264        assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape265        return (latent_height, latent_width)266 267    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):268        min_batch = batch_size if static_batch else self.min_batch269        max_batch = batch_size if static_batch else self.max_batch270        latent_height = image_height // 8271        latent_width = image_width // 8272        min_image_height = image_height if static_shape else self.min_image_shape273        max_image_height = image_height if static_shape else self.max_image_shape274        min_image_width = image_width if static_shape else self.min_image_shape275        max_image_width = image_width if static_shape else self.max_image_shape276        min_latent_height = latent_height if static_shape else self.min_latent_shape277        max_latent_height = latent_height if static_shape else self.max_latent_shape278        min_latent_width = latent_width if static_shape else self.min_latent_shape279        max_latent_width = latent_width if static_shape else self.max_latent_shape280        return (281            min_batch,282            max_batch,283            min_image_height,284            max_image_height,285            min_image_width,286            max_image_width,287            min_latent_height,288            max_latent_height,289            min_latent_width,290            max_latent_width,291        )292 293 294def getOnnxPath(model_name, onnx_dir, opt=True):295    return os.path.join(onnx_dir, model_name + (".opt" if opt else "") + ".onnx")296 297 298def getEnginePath(model_name, engine_dir):299    return os.path.join(engine_dir, model_name + ".plan")300 301 302def build_engines(303    models: dict,304    engine_dir,305    onnx_dir,306    onnx_opset,307    opt_image_height,308    opt_image_width,309    opt_batch_size=1,310    force_engine_rebuild=False,311    static_batch=False,312    static_shape=True,313    enable_all_tactics=False,314    timing_cache=None,315):316    built_engines = {}317    if not os.path.isdir(onnx_dir):318        os.makedirs(onnx_dir)319    if not os.path.isdir(engine_dir):320        os.makedirs(engine_dir)321 322    # Export models to ONNX323    for model_name, model_obj in models.items():324        engine_path = getEnginePath(model_name, engine_dir)325        if force_engine_rebuild or not os.path.exists(engine_path):326            logger.warning("Building Engines...")327            logger.warning("Engine build can take a while to complete")328            onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)329            onnx_opt_path = getOnnxPath(model_name, onnx_dir)330            if force_engine_rebuild or not os.path.exists(onnx_opt_path):331                if force_engine_rebuild or not os.path.exists(onnx_path):332                    logger.warning(f"Exporting model: {onnx_path}")333                    model = model_obj.get_model()334                    with torch.inference_mode(), torch.autocast("cuda"):335                        inputs = model_obj.get_sample_input(opt_batch_size, opt_image_height, opt_image_width)336                        torch.onnx.export(337                            model,338                            inputs,339                            onnx_path,340                            export_params=True,341                            opset_version=onnx_opset,342                            do_constant_folding=True,343                            input_names=model_obj.get_input_names(),344                            output_names=model_obj.get_output_names(),345                            dynamic_axes=model_obj.get_dynamic_axes(),346                        )347                    del model348                    torch.cuda.empty_cache()349                    gc.collect()350                else:351                    logger.warning(f"Found cached model: {onnx_path}")352 353                # Optimize onnx354                if force_engine_rebuild or not os.path.exists(onnx_opt_path):355                    logger.warning(f"Generating optimizing model: {onnx_opt_path}")356                    onnx_opt_graph = model_obj.optimize(onnx.load(onnx_path))357                    onnx.save(onnx_opt_graph, onnx_opt_path)358                else:359                    logger.warning(f"Found cached optimized model: {onnx_opt_path} ")360 361    # Build TensorRT engines362    for model_name, model_obj in models.items():363        engine_path = getEnginePath(model_name, engine_dir)364        engine = Engine(engine_path)365        onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)366        onnx_opt_path = getOnnxPath(model_name, onnx_dir)367 368        if force_engine_rebuild or not os.path.exists(engine.engine_path):369            engine.build(370                onnx_opt_path,371                fp16=True,372                input_profile=model_obj.get_input_profile(373                    opt_batch_size,374                    opt_image_height,375                    opt_image_width,376                    static_batch=static_batch,377                    static_shape=static_shape,378                ),379                timing_cache=timing_cache,380            )381        built_engines[model_name] = engine382 383    # Load and activate TensorRT engines384    for model_name, model_obj in models.items():385        engine = built_engines[model_name]386        engine.load()387        engine.activate()388 389    return built_engines390 391 392def runEngine(engine, feed_dict, stream):393    return engine.infer(feed_dict, stream)394 395 396class CLIP(BaseModel):397    def __init__(self, model, device, max_batch_size, embedding_dim):398        super(CLIP, self).__init__(399            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim400        )401        self.name = "CLIP"402 403    def get_input_names(self):404        return ["input_ids"]405 406    def get_output_names(self):407        return ["text_embeddings", "pooler_output"]408 409    def get_dynamic_axes(self):410        return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}411 412    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):413        self.check_dims(batch_size, image_height, image_width)414        min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(415            batch_size, image_height, image_width, static_batch, static_shape416        )417        return {418            "input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]419        }420 421    def get_shape_dict(self, batch_size, image_height, image_width):422        self.check_dims(batch_size, image_height, image_width)423        return {424            "input_ids": (batch_size, self.text_maxlen),425            "text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),426        }427 428    def get_sample_input(self, batch_size, image_height, image_width):429        self.check_dims(batch_size, image_height, image_width)430        return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)431 432    def optimize(self, onnx_graph):433        opt = Optimizer(onnx_graph)434        opt.select_outputs([0])  # delete graph output#1435        opt.cleanup()436        opt.fold_constants()437        opt.infer_shapes()438        opt.select_outputs([0], names=["text_embeddings"])  # rename network output439        opt_onnx_graph = opt.cleanup(return_onnx=True)440        return opt_onnx_graph441 442 443def make_CLIP(model, device, max_batch_size, embedding_dim, inpaint=False):444    return CLIP(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)445 446 447class UNet(BaseModel):448    def __init__(449        self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77, unet_dim=4450    ):451        super(UNet, self).__init__(452            model=model,453            fp16=fp16,454            device=device,455            max_batch_size=max_batch_size,456            embedding_dim=embedding_dim,457            text_maxlen=text_maxlen,458        )459        self.unet_dim = unet_dim460        self.name = "UNet"461 462    def get_input_names(self):463        return ["sample", "timestep", "encoder_hidden_states"]464 465    def get_output_names(self):466        return ["latent"]467 468    def get_dynamic_axes(self):469        return {470            "sample": {0: "2B", 2: "H", 3: "W"},471            "encoder_hidden_states": {0: "2B"},472            "latent": {0: "2B", 2: "H", 3: "W"},473        }474 475    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):476        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)477        (478            min_batch,479            max_batch,480            _,481            _,482            _,483            _,484            min_latent_height,485            max_latent_height,486            min_latent_width,487            max_latent_width,488        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)489        return {490            "sample": [491                (2 * min_batch, self.unet_dim, min_latent_height, min_latent_width),492                (2 * batch_size, self.unet_dim, latent_height, latent_width),493                (2 * max_batch, self.unet_dim, max_latent_height, max_latent_width),494            ],495            "encoder_hidden_states": [496                (2 * min_batch, self.text_maxlen, self.embedding_dim),497                (2 * batch_size, self.text_maxlen, self.embedding_dim),498                (2 * max_batch, self.text_maxlen, self.embedding_dim),499            ],500        }501 502    def get_shape_dict(self, batch_size, image_height, image_width):503        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)504        return {505            "sample": (2 * batch_size, self.unet_dim, latent_height, latent_width),506            "encoder_hidden_states": (2 * batch_size, self.text_maxlen, self.embedding_dim),507            "latent": (2 * batch_size, 4, latent_height, latent_width),508        }509 510    def get_sample_input(self, batch_size, image_height, image_width):511        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)512        dtype = torch.float16 if self.fp16 else torch.float32513        return (514            torch.randn(515                2 * batch_size, self.unet_dim, latent_height, latent_width, dtype=torch.float32, device=self.device516            ),517            torch.tensor([1.0], dtype=torch.float32, device=self.device),518            torch.randn(2 * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),519        )520 521 522def make_UNet(model, device, max_batch_size, embedding_dim, inpaint=False):523    return UNet(524        model,525        fp16=True,526        device=device,527        max_batch_size=max_batch_size,528        embedding_dim=embedding_dim,529        unet_dim=(9 if inpaint else 4),530    )531 532 533class VAE(BaseModel):534    def __init__(self, model, device, max_batch_size, embedding_dim):535        super(VAE, self).__init__(536            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim537        )538        self.name = "VAE decoder"539 540    def get_input_names(self):541        return ["latent"]542 543    def get_output_names(self):544        return ["images"]545 546    def get_dynamic_axes(self):547        return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}548 549    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):550        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)551        (552            min_batch,553            max_batch,554            _,555            _,556            _,557            _,558            min_latent_height,559            max_latent_height,560            min_latent_width,561            max_latent_width,562        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)563        return {564            "latent": [565                (min_batch, 4, min_latent_height, min_latent_width),566                (batch_size, 4, latent_height, latent_width),567                (max_batch, 4, max_latent_height, max_latent_width),568            ]569        }570 571    def get_shape_dict(self, batch_size, image_height, image_width):572        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)573        return {574            "latent": (batch_size, 4, latent_height, latent_width),575            "images": (batch_size, 3, image_height, image_width),576        }577 578    def get_sample_input(self, batch_size, image_height, image_width):579        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)580        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)581 582 583def make_VAE(model, device, max_batch_size, embedding_dim, inpaint=False):584    return VAE(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)585 586 587class TorchVAEEncoder(torch.nn.Module):588    def __init__(self, model):589        super().__init__()590        self.vae_encoder = model591 592    def forward(self, x):593        return retrieve_latents(self.vae_encoder.encode(x))594 595 596class VAEEncoder(BaseModel):597    def __init__(self, model, device, max_batch_size, embedding_dim):598        super(VAEEncoder, self).__init__(599            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim600        )601        self.name = "VAE encoder"602 603    def get_model(self):604        vae_encoder = TorchVAEEncoder(self.model)605        return vae_encoder606 607    def get_input_names(self):608        return ["images"]609 610    def get_output_names(self):611        return ["latent"]612 613    def get_dynamic_axes(self):614        return {"images": {0: "B", 2: "8H", 3: "8W"}, "latent": {0: "B", 2: "H", 3: "W"}}615 616    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):617        assert batch_size >= self.min_batch and batch_size <= self.max_batch618        min_batch = batch_size if static_batch else self.min_batch619        max_batch = batch_size if static_batch else self.max_batch620        self.check_dims(batch_size, image_height, image_width)621        (622            min_batch,623            max_batch,624            min_image_height,625            max_image_height,626            min_image_width,627            max_image_width,628            _,629            _,630            _,631            _,632        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)633 634        return {635            "images": [636                (min_batch, 3, min_image_height, min_image_width),637                (batch_size, 3, image_height, image_width),638                (max_batch, 3, max_image_height, max_image_width),639            ]640        }641 642    def get_shape_dict(self, batch_size, image_height, image_width):643        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)644        return {645            "images": (batch_size, 3, image_height, image_width),646            "latent": (batch_size, 4, latent_height, latent_width),647        }648 649    def get_sample_input(self, batch_size, image_height, image_width):650        self.check_dims(batch_size, image_height, image_width)651        return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device)652 653 654def make_VAEEncoder(model, device, max_batch_size, embedding_dim, inpaint=False):655    return VAEEncoder(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)656 657 658class TensorRTStableDiffusionImg2ImgPipeline(DiffusionPipeline):659    r"""660    Pipeline for image-to-image generation using TensorRT accelerated Stable Diffusion.661 662    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the663    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)664 665    Args:666        vae ([`AutoencoderKL`]):667            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.668        text_encoder ([`CLIPTextModel`]):669            Frozen text-encoder. Stable Diffusion uses the text portion of670            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically671            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.672        tokenizer (`CLIPTokenizer`):673            Tokenizer of class674            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).675        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.676        scheduler ([`SchedulerMixin`]):677            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of678            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].679        safety_checker ([`StableDiffusionSafetyChecker`]):680            Classification module that estimates whether generated images could be considered offensive or harmful.681            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.682        feature_extractor ([`CLIPImageProcessor`]):683            Model that extracts features from generated images to be used as inputs for the `safety_checker`.684    """685 686    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]687 688    def __init__(689        self,690        vae: AutoencoderKL,691        text_encoder: CLIPTextModel,692        tokenizer: CLIPTokenizer,693        unet: UNet2DConditionModel,694        scheduler: DDIMScheduler,695        safety_checker: StableDiffusionSafetyChecker,696        feature_extractor: CLIPImageProcessor,697        image_encoder: CLIPVisionModelWithProjection = None,698        requires_safety_checker: bool = True,699        stages=["clip", "unet", "vae", "vae_encoder"],700        image_height: int = 512,701        image_width: int = 512,702        max_batch_size: int = 16,703        # ONNX export parameters704        onnx_opset: int = 17,705        onnx_dir: str = "onnx",706        # TensorRT engine build parameters707        engine_dir: str = "engine",708        force_engine_rebuild: bool = False,709        timing_cache: str = "timing_cache",710    ):711        super().__init__()712 713        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:714            deprecation_message = (715                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"716                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "717                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"718                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"719                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"720                " file"721            )722            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)723            new_config = dict(scheduler.config)724            new_config["steps_offset"] = 1725            scheduler._internal_dict = FrozenDict(new_config)726 727        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:728            deprecation_message = (729                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."730                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"731                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"732                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"733                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"734            )735            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)736            new_config = dict(scheduler.config)737            new_config["clip_sample"] = False738            scheduler._internal_dict = FrozenDict(new_config)739 740        if safety_checker is None and requires_safety_checker:741            logger.warning(742                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"743                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"744                " results in services or applications open to the public. Both the diffusers team and Hugging Face"745                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"746                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"747                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."748            )749 750        if safety_checker is not None and feature_extractor is None:751            raise ValueError(752                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"753                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."754            )755 756        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(757            version.parse(unet.config._diffusers_version).base_version758        ) < version.parse("0.9.0.dev0")759        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64760        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:761            deprecation_message = (762                "The configuration file of the unet has set the default `sample_size` to smaller than"763                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"764                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"765                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"766                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"767                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"768                " in the config might lead to incorrect results in future versions. If you have downloaded this"769                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"770                " the `unet/config.json` file"771            )772            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)773            new_config = dict(unet.config)774            new_config["sample_size"] = 64775            unet._internal_dict = FrozenDict(new_config)776 777        self.register_modules(778            vae=vae,779            text_encoder=text_encoder,780            tokenizer=tokenizer,781            unet=unet,782            scheduler=scheduler,783            safety_checker=safety_checker,784            feature_extractor=feature_extractor,785            image_encoder=image_encoder,786        )787 788        self.stages = stages789        self.image_height, self.image_width = image_height, image_width790        self.inpaint = False791        self.onnx_opset = onnx_opset792        self.onnx_dir = onnx_dir793        self.engine_dir = engine_dir794        self.force_engine_rebuild = force_engine_rebuild795        self.timing_cache = timing_cache796        self.build_static_batch = False797        self.build_dynamic_shape = False798 799        self.max_batch_size = max_batch_size800        # TODO: Restrict batch size to 4 for larger image dimensions as a WAR for TensorRT limitation.801        if self.build_dynamic_shape or self.image_height > 512 or self.image_width > 512:802            self.max_batch_size = 4803 804        self.stream = None  # loaded in loadResources()805        self.models = {}  # loaded in __loadModels()806        self.engine = {}  # loaded in build_engines()807 808        self.vae.forward = self.vae.decode809        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)810        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)811        self.register_to_config(requires_safety_checker=requires_safety_checker)812 813    def __loadModels(self):814        # Load pipeline models815        self.embedding_dim = self.text_encoder.config.hidden_size816        models_args = {817            "device": self.torch_device,818            "max_batch_size": self.max_batch_size,819            "embedding_dim": self.embedding_dim,820            "inpaint": self.inpaint,821        }822        if "clip" in self.stages:823            self.models["clip"] = make_CLIP(self.text_encoder, **models_args)824        if "unet" in self.stages:825            self.models["unet"] = make_UNet(self.unet, **models_args)826        if "vae" in self.stages:827            self.models["vae"] = make_VAE(self.vae, **models_args)828        if "vae_encoder" in self.stages:829            self.models["vae_encoder"] = make_VAEEncoder(self.vae, **models_args)830 831    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker832    def run_safety_checker(833        self, image: Union[torch.Tensor, PIL.Image.Image], device: torch.device, dtype: torch.dtype834    ) -> Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]:835        r"""836        Runs the safety checker on the given image.837        Args:838            image (Union[torch.Tensor, PIL.Image.Image]): The input image to be checked.839            device (torch.device): The device to run the safety checker on.840            dtype (torch.dtype): The data type of the input image.841        Returns:842            (image, has_nsfw_concept) Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]: A tuple containing the processed image and843            a boolean indicating whether the image has a NSFW (Not Safe for Work) concept.844        """845        if self.safety_checker is None:846            has_nsfw_concept = None847        else:848            if torch.is_tensor(image):849                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")850            else:851                feature_extractor_input = self.image_processor.numpy_to_pil(image)852            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)853            image, has_nsfw_concept = self.safety_checker(854                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)855            )856        return image, has_nsfw_concept857 858    @classmethod859    @validate_hf_hub_args860    def set_cached_folder(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):861        cache_dir = kwargs.pop("cache_dir", None)862        proxies = kwargs.pop("proxies", None)863        local_files_only = kwargs.pop("local_files_only", False)864        token = kwargs.pop("token", None)865        revision = kwargs.pop("revision", None)866 867        cls.cached_folder = (868            pretrained_model_name_or_path869            if os.path.isdir(pretrained_model_name_or_path)870            else snapshot_download(871                pretrained_model_name_or_path,872                cache_dir=cache_dir,873                proxies=proxies,874                local_files_only=local_files_only,875                token=token,876                revision=revision,877            )878        )879 880    def to(self, torch_device: Optional[Union[str, torch.device]] = None, silence_dtype_warnings: bool = False):881        super().to(torch_device, silence_dtype_warnings=silence_dtype_warnings)882 883        self.onnx_dir = os.path.join(self.cached_folder, self.onnx_dir)884        self.engine_dir = os.path.join(self.cached_folder, self.engine_dir)885        self.timing_cache = os.path.join(self.cached_folder, self.timing_cache)886 887        # set device888        self.torch_device = self._execution_device889        logger.warning(f"Running inference on device: {self.torch_device}")890 891        # load models892        self.__loadModels()893 894        # build engines895        self.engine = build_engines(896            self.models,897            self.engine_dir,898            self.onnx_dir,899            self.onnx_opset,900            opt_image_height=self.image_height,901            opt_image_width=self.image_width,902            force_engine_rebuild=self.force_engine_rebuild,903            static_batch=self.build_static_batch,904            static_shape=not self.build_dynamic_shape,905            timing_cache=self.timing_cache,906        )907 908        return self909 910    def __initialize_timesteps(self, timesteps, strength):911        self.scheduler.set_timesteps(timesteps)912        offset = self.scheduler.steps_offset if hasattr(self.scheduler, "steps_offset") else 0913        init_timestep = int(timesteps * strength) + offset914        init_timestep = min(init_timestep, timesteps)915        t_start = max(timesteps - init_timestep + offset, 0)916        timesteps = self.scheduler.timesteps[t_start:].to(self.torch_device)917        return timesteps, t_start918 919    def __preprocess_images(self, batch_size, images=()):920        init_images = []921        for image in images:922            image = image.to(self.torch_device).float()923            image = image.repeat(batch_size, 1, 1, 1)924            init_images.append(image)925        return tuple(init_images)926 927    def __encode_image(self, init_image):928        init_latents = runEngine(self.engine["vae_encoder"], {"images": init_image}, self.stream)["latent"]929        init_latents = 0.18215 * init_latents930        return init_latents931 932    def __encode_prompt(self, prompt, negative_prompt):933        r"""934        Encodes the prompt into text encoder hidden states.935 936        Args:937             prompt (`str` or `List[str]`, *optional*):938                prompt to be encoded939            negative_prompt (`str` or `List[str]`, *optional*):940                The prompt or prompts not to guide the image generation. If not defined, one has to pass941                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.942                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).943        """944        # Tokenize prompt945        text_input_ids = (946            self.tokenizer(947                prompt,948                padding="max_length",949                max_length=self.tokenizer.model_max_length,950                truncation=True,951                return_tensors="pt",952            )953            .input_ids.type(torch.int32)954            .to(self.torch_device)955        )956 957        # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt958        text_embeddings = runEngine(self.engine["clip"], {"input_ids": text_input_ids}, self.stream)[959            "text_embeddings"960        ].clone()961 962        # Tokenize negative prompt963        uncond_input_ids = (964            self.tokenizer(965                negative_prompt,966                padding="max_length",967                max_length=self.tokenizer.model_max_length,968                truncation=True,969                return_tensors="pt",970            )971            .input_ids.type(torch.int32)972            .to(self.torch_device)973        )974        uncond_embeddings = runEngine(self.engine["clip"], {"input_ids": uncond_input_ids}, self.stream)[975            "text_embeddings"976        ]977 978        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance979        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)980 981        return text_embeddings982 983    def __denoise_latent(984        self, latents, text_embeddings, timesteps=None, step_offset=0, mask=None, masked_image_latents=None985    ):986        if not isinstance(timesteps, torch.Tensor):987            timesteps = self.scheduler.timesteps988        for step_index, timestep in enumerate(timesteps):989            # Expand the latents if we are doing classifier free guidance990            latent_model_input = torch.cat([latents] * 2)991            latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)992            if isinstance(mask, torch.Tensor):993                latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)994 995            # Predict the noise residual996            timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep997 998            noise_pred = runEngine(999                self.engine["unet"],1000                {"sample": latent_model_input, "timestep": timestep_float, "encoder_hidden_states": text_embeddings},1001                self.stream,1002            )["latent"]1003 1004            # Perform guidance1005            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)1006            noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)1007 1008            latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample1009 1010        latents = 1.0 / 0.18215 * latents1011        return latents1012 1013    def __decode_latent(self, latents):1014        images = runEngine(self.engine["vae"], {"latent": latents}, self.stream)["images"]1015        images = (images / 2 + 0.5).clamp(0, 1)1016        return images.cpu().permute(0, 2, 3, 1).float().numpy()1017 1018    def __loadResources(self, image_height, image_width, batch_size):1019        self.stream = cudart.cudaStreamCreate()[1]1020 1021        # Allocate buffers for TensorRT engine bindings1022        for model_name, obj in self.models.items():1023            self.engine[model_name].allocate_buffers(1024                shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.torch_device1025            )1026 1027    @torch.no_grad()1028    def __call__(1029        self,1030        prompt: Union[str, List[str]] = None,1031        image: Union[torch.Tensor, PIL.Image.Image] = None,1032        strength: float = 0.8,1033        num_inference_steps: int = 50,1034        guidance_scale: float = 7.5,1035        negative_prompt: Optional[Union[str, List[str]]] = None,1036        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,1037    ):1038        r"""1039        Function invoked when calling the pipeline for generation.1040 1041        Args:1042            prompt (`str` or `List[str]`, *optional*):1043                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.1044                instead.1045            image (`PIL.Image.Image`):1046                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will1047                be masked out with `mask_image` and repainted according to `prompt`.1048            strength (`float`, *optional*, defaults to 0.8):1049                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`1050                will be used as a starting point, adding more noise to it the larger the `strength`. The number of1051                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will1052                be maximum and the denoising process will run for the full number of iterations specified in1053                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.1054            num_inference_steps (`int`, *optional*, defaults to 50):1055                The number of denoising steps. More denoising steps usually lead to a higher quality image at the1056                expense of slower inference.1057            guidance_scale (`float`, *optional*, defaults to 7.5):1058                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).1059                `guidance_scale` is defined as `w` of equation 2. of [Imagen1060                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1061                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1062                usually at the expense of lower image quality.1063            negative_prompt (`str` or `List[str]`, *optional*):1064                The prompt or prompts not to guide the image generation. If not defined, one has to pass1065                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.1066                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).1067            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):1068                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)1069                to make generation deterministic.1070 1071        """1072        self.generator = generator1073        self.denoising_steps = num_inference_steps1074        self._guidance_scale = guidance_scale1075 1076        # Pre-compute latent input scales and linear multistep coefficients1077        self.scheduler.set_timesteps(self.denoising_steps, device=self.torch_device)1078 1079        # Define call parameters1080        if prompt is not None and isinstance(prompt, str):1081            batch_size = 11082            prompt = [prompt]1083        elif prompt is not None and isinstance(prompt, list):1084            batch_size = len(prompt)1085        else:1086            raise ValueError(f"Expected prompt to be of type list or str but got {type(prompt)}")1087 1088        if negative_prompt is None:1089            negative_prompt = [""] * batch_size1090 1091        if negative_prompt is not None and isinstance(negative_prompt, str):1092            negative_prompt = [negative_prompt]1093 1094        assert len(prompt) == len(negative_prompt)1095 1096        if batch_size > self.max_batch_size:1097            raise ValueError(1098                f"Batch size {len(prompt)} is larger than allowed {self.max_batch_size}. If dynamic shape is used, then maximum batch size is 4"1099            )1100 1101        # load resources1102        self.__loadResources(self.image_height, self.image_width, batch_size)1103 1104        with torch.inference_mode(), torch.autocast("cuda"), trt.Runtime(TRT_LOGGER):1105            # Initialize timesteps1106            timesteps, t_start = self.__initialize_timesteps(self.denoising_steps, strength)1107            latent_timestep = timesteps[:1].repeat(batch_size)1108 1109            # Pre-process input image1110            if isinstance(image, PIL.Image.Image):1111                image = preprocess_image(image)1112            init_image = self.__preprocess_images(batch_size, (image,))[0]1113 1114            # VAE encode init image1115            init_latents = self.__encode_image(init_image)1116 1117            # Add noise to latents using timesteps1118            noise = torch.randn(1119                init_latents.shape, generator=self.generator, device=self.torch_device, dtype=torch.float321120            )1121            latents = self.scheduler.add_noise(init_latents, noise, latent_timestep)1122 1123            # CLIP text encoder1124            text_embeddings = self.__encode_prompt(prompt, negative_prompt)1125 1126            # UNet denoiser1127            latents = self.__denoise_latent(latents, text_embeddings, timesteps=timesteps, step_offset=t_start)1128 1129            # VAE decode latent1130            images = self.__decode_latent(latents)1131 1132        images, has_nsfw_concept = self.run_safety_checker(images, self.torch_device, text_embeddings.dtype)1133        images = self.numpy_to_pil(images)1134        return StableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept)1135