diffusers/community-pipelines-mirror

1#2# Copyright 2025 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_inpaint import (56    prepare_mask_and_masked_image,57    retrieve_latents,58)59from diffusers.schedulers import DDIMScheduler60from diffusers.utils import logging61from diffusers.utils.torch_utils import randn_tensor62 63 64"""65Installation instructions66python3 -m pip install --upgrade transformers diffusers>=0.16.067python3 -m pip install --upgrade tensorrt~=10.2.068python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com69python3 -m pip install onnxruntime70"""71 72TRT_LOGGER = trt.Logger(trt.Logger.ERROR)73logger = logging.get_logger(__name__)  # pylint: disable=invalid-name74 75# Map of numpy dtype -> torch dtype76numpy_to_torch_dtype_dict = {77    np.uint8: torch.uint8,78    np.int8: torch.int8,79    np.int16: torch.int16,80    np.int32: torch.int32,81    np.int64: torch.int64,82    np.float16: torch.float16,83    np.float32: torch.float32,84    np.float64: torch.float64,85    np.complex64: torch.complex64,86    np.complex128: torch.complex128,87}88if np.version.full_version >= "1.24.0":89    numpy_to_torch_dtype_dict[np.bool_] = torch.bool90else:91    numpy_to_torch_dtype_dict[np.bool] = torch.bool92 93# Map of torch dtype -> numpy dtype94torch_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()}95 96 97def preprocess_image(image):98    """99    image: torch.Tensor100    """101    w, h = image.size102    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32103    image = image.resize((w, h))104    image = np.array(image).astype(np.float32) / 255.0105    image = image[None].transpose(0, 3, 1, 2)106    image = torch.from_numpy(image).contiguous()107    return 2.0 * image - 1.0108 109 110class Engine:111    def __init__(self, engine_path):112        self.engine_path = engine_path113        self.engine = None114        self.context = None115        self.buffers = OrderedDict()116        self.tensors = OrderedDict()117 118    def __del__(self):119        [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray)]120        del self.engine121        del self.context122        del self.buffers123        del self.tensors124 125    def build(126        self,127        onnx_path,128        fp16,129        input_profile=None,130        enable_all_tactics=False,131        timing_cache=None,132    ):133        logger.warning(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")134        p = Profile()135        if input_profile:136            for name, dims in input_profile.items():137                assert len(dims) == 3138                p.add(name, min=dims[0], opt=dims[1], max=dims[2])139 140        extra_build_args = {}141        if not enable_all_tactics:142            extra_build_args["tactic_sources"] = []143 144        engine = engine_from_network(145            network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]),146            config=CreateConfig(fp16=fp16, profiles=[p], load_timing_cache=timing_cache, **extra_build_args),147            save_timing_cache=timing_cache,148        )149        save_engine(engine, path=self.engine_path)150 151    def load(self):152        logger.warning(f"Loading TensorRT engine: {self.engine_path}")153        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))154 155    def activate(self):156        self.context = self.engine.create_execution_context()157 158    def allocate_buffers(self, shape_dict=None, device="cuda"):159        for binding in range(self.engine.num_io_tensors):160            name = self.engine.get_tensor_name(binding)161            if shape_dict and name in shape_dict:162                shape = shape_dict[name]163            else:164                shape = self.engine.get_tensor_shape(name)165            dtype = trt.nptype(self.engine.get_tensor_dtype(name))166            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:167                self.context.set_input_shape(name, shape)168            tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device)169            self.tensors[name] = tensor170 171    def infer(self, feed_dict, stream):172        for name, buf in feed_dict.items():173            self.tensors[name].copy_(buf)174        for name, tensor in self.tensors.items():175            self.context.set_tensor_address(name, tensor.data_ptr())176        noerror = self.context.execute_async_v3(stream)177        if not noerror:178            raise ValueError("ERROR: inference failed.")179 180        return self.tensors181 182 183class Optimizer:184    def __init__(self, onnx_graph):185        self.graph = gs.import_onnx(onnx_graph)186 187    def cleanup(self, return_onnx=False):188        self.graph.cleanup().toposort()189        if return_onnx:190            return gs.export_onnx(self.graph)191 192    def select_outputs(self, keep, names=None):193        self.graph.outputs = [self.graph.outputs[o] for o in keep]194        if names:195            for i, name in enumerate(names):196                self.graph.outputs[i].name = name197 198    def fold_constants(self, return_onnx=False):199        onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)200        self.graph = gs.import_onnx(onnx_graph)201        if return_onnx:202            return onnx_graph203 204    def infer_shapes(self, return_onnx=False):205        onnx_graph = gs.export_onnx(self.graph)206        if onnx_graph.ByteSize() > 2147483648:207            raise TypeError("ERROR: model size exceeds supported 2GB limit")208        else:209            onnx_graph = shape_inference.infer_shapes(onnx_graph)210 211        self.graph = gs.import_onnx(onnx_graph)212        if return_onnx:213            return onnx_graph214 215 216class BaseModel:217    def __init__(self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77):218        self.model = model219        self.name = "SD Model"220        self.fp16 = fp16221        self.device = device222 223        self.min_batch = 1224        self.max_batch = max_batch_size225        self.min_image_shape = 256  # min image resolution: 256x256226        self.max_image_shape = 1024  # max image resolution: 1024x1024227        self.min_latent_shape = self.min_image_shape // 8228        self.max_latent_shape = self.max_image_shape // 8229 230        self.embedding_dim = embedding_dim231        self.text_maxlen = text_maxlen232 233    def get_model(self):234        return self.model235 236    def get_input_names(self):237        pass238 239    def get_output_names(self):240        pass241 242    def get_dynamic_axes(self):243        return None244 245    def get_sample_input(self, batch_size, image_height, image_width):246        pass247 248    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):249        return None250 251    def get_shape_dict(self, batch_size, image_height, image_width):252        return None253 254    def optimize(self, onnx_graph):255        opt = Optimizer(onnx_graph)256        opt.cleanup()257        opt.fold_constants()258        opt.infer_shapes()259        onnx_opt_graph = opt.cleanup(return_onnx=True)260        return onnx_opt_graph261 262    def check_dims(self, batch_size, image_height, image_width):263        assert batch_size >= self.min_batch and batch_size <= self.max_batch264        assert image_height % 8 == 0 or image_width % 8 == 0265        latent_height = image_height // 8266        latent_width = image_width // 8267        assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape268        assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape269        return (latent_height, latent_width)270 271    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):272        min_batch = batch_size if static_batch else self.min_batch273        max_batch = batch_size if static_batch else self.max_batch274        latent_height = image_height // 8275        latent_width = image_width // 8276        min_image_height = image_height if static_shape else self.min_image_shape277        max_image_height = image_height if static_shape else self.max_image_shape278        min_image_width = image_width if static_shape else self.min_image_shape279        max_image_width = image_width if static_shape else self.max_image_shape280        min_latent_height = latent_height if static_shape else self.min_latent_shape281        max_latent_height = latent_height if static_shape else self.max_latent_shape282        min_latent_width = latent_width if static_shape else self.min_latent_shape283        max_latent_width = latent_width if static_shape else self.max_latent_shape284        return (285            min_batch,286            max_batch,287            min_image_height,288            max_image_height,289            min_image_width,290            max_image_width,291            min_latent_height,292            max_latent_height,293            min_latent_width,294            max_latent_width,295        )296 297 298def getOnnxPath(model_name, onnx_dir, opt=True):299    return os.path.join(onnx_dir, model_name + (".opt" if opt else "") + ".onnx")300 301 302def getEnginePath(model_name, engine_dir):303    return os.path.join(engine_dir, model_name + ".plan")304 305 306def build_engines(307    models: dict,308    engine_dir,309    onnx_dir,310    onnx_opset,311    opt_image_height,312    opt_image_width,313    opt_batch_size=1,314    force_engine_rebuild=False,315    static_batch=False,316    static_shape=True,317    enable_all_tactics=False,318    timing_cache=None,319):320    built_engines = {}321    if not os.path.isdir(onnx_dir):322        os.makedirs(onnx_dir)323    if not os.path.isdir(engine_dir):324        os.makedirs(engine_dir)325 326    # Export models to ONNX327    for model_name, model_obj in models.items():328        engine_path = getEnginePath(model_name, engine_dir)329        if force_engine_rebuild or not os.path.exists(engine_path):330            logger.warning("Building Engines...")331            logger.warning("Engine build can take a while to complete")332            onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)333            onnx_opt_path = getOnnxPath(model_name, onnx_dir)334            if force_engine_rebuild or not os.path.exists(onnx_opt_path):335                if force_engine_rebuild or not os.path.exists(onnx_path):336                    logger.warning(f"Exporting model: {onnx_path}")337                    model = model_obj.get_model()338                    with torch.inference_mode(), torch.autocast("cuda"):339                        inputs = model_obj.get_sample_input(opt_batch_size, opt_image_height, opt_image_width)340                        torch.onnx.export(341                            model,342                            inputs,343                            onnx_path,344                            export_params=True,345                            opset_version=onnx_opset,346                            do_constant_folding=True,347                            input_names=model_obj.get_input_names(),348                            output_names=model_obj.get_output_names(),349                            dynamic_axes=model_obj.get_dynamic_axes(),350                        )351                    del model352                    torch.cuda.empty_cache()353                    gc.collect()354                else:355                    logger.warning(f"Found cached model: {onnx_path}")356 357                # Optimize onnx358                if force_engine_rebuild or not os.path.exists(onnx_opt_path):359                    logger.warning(f"Generating optimizing model: {onnx_opt_path}")360                    onnx_opt_graph = model_obj.optimize(onnx.load(onnx_path))361                    onnx.save(onnx_opt_graph, onnx_opt_path)362                else:363                    logger.warning(f"Found cached optimized model: {onnx_opt_path} ")364 365    # Build TensorRT engines366    for model_name, model_obj in models.items():367        engine_path = getEnginePath(model_name, engine_dir)368        engine = Engine(engine_path)369        onnx_path = getOnnxPath(model_name, onnx_dir, opt=False)370        onnx_opt_path = getOnnxPath(model_name, onnx_dir)371 372        if force_engine_rebuild or not os.path.exists(engine.engine_path):373            engine.build(374                onnx_opt_path,375                fp16=True,376                input_profile=model_obj.get_input_profile(377                    opt_batch_size,378                    opt_image_height,379                    opt_image_width,380                    static_batch=static_batch,381                    static_shape=static_shape,382                ),383                timing_cache=timing_cache,384            )385        built_engines[model_name] = engine386 387    # Load and activate TensorRT engines388    for model_name, model_obj in models.items():389        engine = built_engines[model_name]390        engine.load()391        engine.activate()392 393    return built_engines394 395 396def runEngine(engine, feed_dict, stream):397    return engine.infer(feed_dict, stream)398 399 400class CLIP(BaseModel):401    def __init__(self, model, device, max_batch_size, embedding_dim):402        super(CLIP, self).__init__(403            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim404        )405        self.name = "CLIP"406 407    def get_input_names(self):408        return ["input_ids"]409 410    def get_output_names(self):411        return ["text_embeddings", "pooler_output"]412 413    def get_dynamic_axes(self):414        return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}415 416    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):417        self.check_dims(batch_size, image_height, image_width)418        min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(419            batch_size, image_height, image_width, static_batch, static_shape420        )421        return {422            "input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]423        }424 425    def get_shape_dict(self, batch_size, image_height, image_width):426        self.check_dims(batch_size, image_height, image_width)427        return {428            "input_ids": (batch_size, self.text_maxlen),429            "text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),430        }431 432    def get_sample_input(self, batch_size, image_height, image_width):433        self.check_dims(batch_size, image_height, image_width)434        return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)435 436    def optimize(self, onnx_graph):437        opt = Optimizer(onnx_graph)438        opt.select_outputs([0])  # delete graph output#1439        opt.cleanup()440        opt.fold_constants()441        opt.infer_shapes()442        opt.select_outputs([0], names=["text_embeddings"])  # rename network output443        opt_onnx_graph = opt.cleanup(return_onnx=True)444        return opt_onnx_graph445 446 447def make_CLIP(model, device, max_batch_size, embedding_dim, inpaint=False):448    return CLIP(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)449 450 451class UNet(BaseModel):452    def __init__(453        self, model, fp16=False, device="cuda", max_batch_size=16, embedding_dim=768, text_maxlen=77, unet_dim=4454    ):455        super(UNet, self).__init__(456            model=model,457            fp16=fp16,458            device=device,459            max_batch_size=max_batch_size,460            embedding_dim=embedding_dim,461            text_maxlen=text_maxlen,462        )463        self.unet_dim = unet_dim464        self.name = "UNet"465 466    def get_input_names(self):467        return ["sample", "timestep", "encoder_hidden_states"]468 469    def get_output_names(self):470        return ["latent"]471 472    def get_dynamic_axes(self):473        return {474            "sample": {0: "2B", 2: "H", 3: "W"},475            "encoder_hidden_states": {0: "2B"},476            "latent": {0: "2B", 2: "H", 3: "W"},477        }478 479    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):480        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)481        (482            min_batch,483            max_batch,484            _,485            _,486            _,487            _,488            min_latent_height,489            max_latent_height,490            min_latent_width,491            max_latent_width,492        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)493        return {494            "sample": [495                (2 * min_batch, self.unet_dim, min_latent_height, min_latent_width),496                (2 * batch_size, self.unet_dim, latent_height, latent_width),497                (2 * max_batch, self.unet_dim, max_latent_height, max_latent_width),498            ],499            "encoder_hidden_states": [500                (2 * min_batch, self.text_maxlen, self.embedding_dim),501                (2 * batch_size, self.text_maxlen, self.embedding_dim),502                (2 * max_batch, self.text_maxlen, self.embedding_dim),503            ],504        }505 506    def get_shape_dict(self, batch_size, image_height, image_width):507        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)508        return {509            "sample": (2 * batch_size, self.unet_dim, latent_height, latent_width),510            "encoder_hidden_states": (2 * batch_size, self.text_maxlen, self.embedding_dim),511            "latent": (2 * batch_size, 4, latent_height, latent_width),512        }513 514    def get_sample_input(self, batch_size, image_height, image_width):515        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)516        dtype = torch.float16 if self.fp16 else torch.float32517        return (518            torch.randn(519                2 * batch_size, self.unet_dim, latent_height, latent_width, dtype=torch.float32, device=self.device520            ),521            torch.tensor([1.0], dtype=torch.float32, device=self.device),522            torch.randn(2 * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),523        )524 525 526def make_UNet(model, device, max_batch_size, embedding_dim, inpaint=False, unet_dim=4):527    return UNet(528        model,529        fp16=True,530        device=device,531        max_batch_size=max_batch_size,532        embedding_dim=embedding_dim,533        unet_dim=unet_dim,534    )535 536 537class VAE(BaseModel):538    def __init__(self, model, device, max_batch_size, embedding_dim):539        super(VAE, self).__init__(540            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim541        )542        self.name = "VAE decoder"543 544    def get_input_names(self):545        return ["latent"]546 547    def get_output_names(self):548        return ["images"]549 550    def get_dynamic_axes(self):551        return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}552 553    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):554        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)555        (556            min_batch,557            max_batch,558            _,559            _,560            _,561            _,562            min_latent_height,563            max_latent_height,564            min_latent_width,565            max_latent_width,566        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)567        return {568            "latent": [569                (min_batch, 4, min_latent_height, min_latent_width),570                (batch_size, 4, latent_height, latent_width),571                (max_batch, 4, max_latent_height, max_latent_width),572            ]573        }574 575    def get_shape_dict(self, batch_size, image_height, image_width):576        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)577        return {578            "latent": (batch_size, 4, latent_height, latent_width),579            "images": (batch_size, 3, image_height, image_width),580        }581 582    def get_sample_input(self, batch_size, image_height, image_width):583        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)584        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)585 586 587def make_VAE(model, device, max_batch_size, embedding_dim, inpaint=False):588    return VAE(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)589 590 591class TorchVAEEncoder(torch.nn.Module):592    def __init__(self, model):593        super().__init__()594        self.vae_encoder = model595 596    def forward(self, x):597        return self.vae_encoder.encode(x).latent_dist.sample()598 599 600class VAEEncoder(BaseModel):601    def __init__(self, model, device, max_batch_size, embedding_dim):602        super(VAEEncoder, self).__init__(603            model=model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim604        )605        self.name = "VAE encoder"606 607    def get_model(self):608        vae_encoder = TorchVAEEncoder(self.model)609        return vae_encoder610 611    def get_input_names(self):612        return ["images"]613 614    def get_output_names(self):615        return ["latent"]616 617    def get_dynamic_axes(self):618        return {"images": {0: "B", 2: "8H", 3: "8W"}, "latent": {0: "B", 2: "H", 3: "W"}}619 620    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):621        assert batch_size >= self.min_batch and batch_size <= self.max_batch622        min_batch = batch_size if static_batch else self.min_batch623        max_batch = batch_size if static_batch else self.max_batch624        self.check_dims(batch_size, image_height, image_width)625        (626            min_batch,627            max_batch,628            min_image_height,629            max_image_height,630            min_image_width,631            max_image_width,632            _,633            _,634            _,635            _,636        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)637 638        return {639            "images": [640                (min_batch, 3, min_image_height, min_image_width),641                (batch_size, 3, image_height, image_width),642                (max_batch, 3, max_image_height, max_image_width),643            ]644        }645 646    def get_shape_dict(self, batch_size, image_height, image_width):647        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)648        return {649            "images": (batch_size, 3, image_height, image_width),650            "latent": (batch_size, 4, latent_height, latent_width),651        }652 653    def get_sample_input(self, batch_size, image_height, image_width):654        self.check_dims(batch_size, image_height, image_width)655        return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device)656 657 658def make_VAEEncoder(model, device, max_batch_size, embedding_dim, inpaint=False):659    return VAEEncoder(model, device=device, max_batch_size=max_batch_size, embedding_dim=embedding_dim)660 661 662class TensorRTStableDiffusionInpaintPipeline(DiffusionPipeline):663    r"""664    Pipeline for inpainting using TensorRT accelerated Stable Diffusion.665 666    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the667    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)668 669    Args:670        vae ([`AutoencoderKL`]):671            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.672        text_encoder ([`CLIPTextModel`]):673            Frozen text-encoder. Stable Diffusion uses the text portion of674            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically675            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.676        tokenizer (`CLIPTokenizer`):677            Tokenizer of class678            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).679        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.680        scheduler ([`SchedulerMixin`]):681            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of682            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].683        safety_checker ([`StableDiffusionSafetyChecker`]):684            Classification module that estimates whether generated images could be considered offensive or harmful.685            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.686        feature_extractor ([`CLIPImageProcessor`]):687            Model that extracts features from generated images to be used as inputs for the `safety_checker`.688    """689 690    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]691 692    def __init__(693        self,694        vae: AutoencoderKL,695        text_encoder: CLIPTextModel,696        tokenizer: CLIPTokenizer,697        unet: UNet2DConditionModel,698        scheduler: DDIMScheduler,699        safety_checker: StableDiffusionSafetyChecker,700        feature_extractor: CLIPImageProcessor,701        image_encoder: CLIPVisionModelWithProjection = None,702        requires_safety_checker: bool = True,703        stages=["clip", "unet", "vae", "vae_encoder"],704        image_height: int = 512,705        image_width: int = 512,706        max_batch_size: int = 16,707        # ONNX export parameters708        onnx_opset: int = 17,709        onnx_dir: str = "onnx",710        # TensorRT engine build parameters711        engine_dir: str = "engine",712        force_engine_rebuild: bool = False,713        timing_cache: str = "timing_cache",714    ):715        super().__init__()716 717        if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1:718            deprecation_message = (719                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"720                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "721                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"722                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"723                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"724                " file"725            )726            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)727            new_config = dict(scheduler.config)728            new_config["steps_offset"] = 1729            scheduler._internal_dict = FrozenDict(new_config)730 731        if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True:732            deprecation_message = (733                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."734                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"735                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"736                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"737                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"738            )739            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)740            new_config = dict(scheduler.config)741            new_config["clip_sample"] = False742            scheduler._internal_dict = FrozenDict(new_config)743 744        if safety_checker is None and requires_safety_checker:745            logger.warning(746                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"747                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"748                " results in services or applications open to the public. Both the diffusers team and Hugging Face"749                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"750                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"751                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."752            )753 754        if safety_checker is not None and feature_extractor is None:755            raise ValueError(756                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"757                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."758            )759 760        is_unet_version_less_0_9_0 = (761            unet is not None762            and hasattr(unet.config, "_diffusers_version")763            and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0")764        )765        is_unet_sample_size_less_64 = (766            unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64767        )768        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:769            deprecation_message = (770                "The configuration file of the unet has set the default `sample_size` to smaller than"771                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"772                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"773                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"774                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"775                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"776                " in the config might lead to incorrect results in future versions. If you have downloaded this"777                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"778                " the `unet/config.json` file"779            )780            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)781            new_config = dict(unet.config)782            new_config["sample_size"] = 64783            unet._internal_dict = FrozenDict(new_config)784 785        self.register_modules(786            vae=vae,787            text_encoder=text_encoder,788            tokenizer=tokenizer,789            unet=unet,790            scheduler=scheduler,791            safety_checker=safety_checker,792            feature_extractor=feature_extractor,793            image_encoder=image_encoder,794        )795 796        self.stages = stages797        self.image_height, self.image_width = image_height, image_width798        self.inpaint = True799        self.onnx_opset = onnx_opset800        self.onnx_dir = onnx_dir801        self.engine_dir = engine_dir802        self.force_engine_rebuild = force_engine_rebuild803        self.timing_cache = timing_cache804        self.build_static_batch = False805        self.build_dynamic_shape = False806 807        self.max_batch_size = max_batch_size808        # TODO: Restrict batch size to 4 for larger image dimensions as a WAR for TensorRT limitation.809        if self.build_dynamic_shape or self.image_height > 512 or self.image_width > 512:810            self.max_batch_size = 4811 812        self.stream = None  # loaded in loadResources()813        self.models = {}  # loaded in __loadModels()814        self.engine = {}  # loaded in build_engines()815 816        self.vae.forward = self.vae.decode817        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8818        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)819        self.register_to_config(requires_safety_checker=requires_safety_checker)820 821    def __loadModels(self):822        # Load pipeline models823        self.embedding_dim = self.text_encoder.config.hidden_size824        models_args = {825            "device": self.torch_device,826            "max_batch_size": self.max_batch_size,827            "embedding_dim": self.embedding_dim,828            "inpaint": self.inpaint,829        }830        if "clip" in self.stages:831            self.models["clip"] = make_CLIP(self.text_encoder, **models_args)832        if "unet" in self.stages:833            self.models["unet"] = make_UNet(self.unet, **models_args, unet_dim=self.unet.config.in_channels)834        if "vae" in self.stages:835            self.models["vae"] = make_VAE(self.vae, **models_args)836        if "vae_encoder" in self.stages:837            self.models["vae_encoder"] = make_VAEEncoder(self.vae, **models_args)838 839    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline840 841    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):842        if isinstance(generator, list):843            image_latents = [844                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])845                for i in range(image.shape[0])846            ]847            image_latents = torch.cat(image_latents, dim=0)848        else:849            image_latents = retrieve_latents(self.vae.encode(image), generator=generator)850 851        image_latents = self.vae.config.scaling_factor * image_latents852 853        return image_latents854 855    def prepare_latents(856        self,857        batch_size,858        num_channels_latents,859        height,860        width,861        dtype,862        device,863        generator,864        latents=None,865        image=None,866        timestep=None,867        is_strength_max=True,868        return_noise=False,869        return_image_latents=False,870    ):871        shape = (872            batch_size,873            num_channels_latents,874            int(height) // self.vae_scale_factor,875            int(width) // self.vae_scale_factor,876        )877        if isinstance(generator, list) and len(generator) != batch_size:878            raise ValueError(879                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"880                f" size of {batch_size}. Make sure the batch size matches the length of the generators."881            )882 883        if (image is None or timestep is None) and not is_strength_max:884            raise ValueError(885                "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise."886                "However, either the image or the noise timestep has not been provided."887            )888 889        if return_image_latents or (latents is None and not is_strength_max):890            image = image.to(device=device, dtype=dtype)891 892            if image.shape[1] == 4:893                image_latents = image894            else:895                image_latents = self._encode_vae_image(image=image, generator=generator)896            image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)897 898        if latents is None:899            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)900            # if strength is 1. then initialise the latents to noise, else initial to image + noise901            latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep)902            # if pure noise then scale the initial latents by the  Scheduler's init sigma903            latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents904        else:905            noise = latents.to(device)906            latents = noise * self.scheduler.init_noise_sigma907 908        outputs = (latents,)909 910        if return_noise:911            outputs += (noise,)912 913        if return_image_latents:914            outputs += (image_latents,)915 916        return outputs917 918    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker919    def run_safety_checker(920        self, image: Union[torch.Tensor, PIL.Image.Image], device: torch.device, dtype: torch.dtype921    ) -> Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]:922        r"""923        Runs the safety checker on the given image.924        Args:925            image (Union[torch.Tensor, PIL.Image.Image]): The input image to be checked.926            device (torch.device): The device to run the safety checker on.927            dtype (torch.dtype): The data type of the input image.928        Returns:929            (image, has_nsfw_concept) Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]: A tuple containing the processed image and930            a boolean indicating whether the image has a NSFW (Not Safe for Work) concept.931        """932        if self.safety_checker is None:933            has_nsfw_concept = None934        else:935            if torch.is_tensor(image):936                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")937            else:938                feature_extractor_input = self.image_processor.numpy_to_pil(image)939            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)940            image, has_nsfw_concept = self.safety_checker(941                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)942            )943        return image, has_nsfw_concept944 945    @classmethod946    @validate_hf_hub_args947    def set_cached_folder(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):948        cache_dir = kwargs.pop("cache_dir", None)949        proxies = kwargs.pop("proxies", None)950        local_files_only = kwargs.pop("local_files_only", False)951        token = kwargs.pop("token", None)952        revision = kwargs.pop("revision", None)953 954        cls.cached_folder = (955            pretrained_model_name_or_path956            if os.path.isdir(pretrained_model_name_or_path)957            else snapshot_download(958                pretrained_model_name_or_path,959                cache_dir=cache_dir,960                proxies=proxies,961                local_files_only=local_files_only,962                token=token,963                revision=revision,964            )965        )966 967    def to(self, torch_device: Optional[Union[str, torch.device]] = None, silence_dtype_warnings: bool = False):968        super().to(torch_device, silence_dtype_warnings=silence_dtype_warnings)969 970        self.onnx_dir = os.path.join(self.cached_folder, self.onnx_dir)971        self.engine_dir = os.path.join(self.cached_folder, self.engine_dir)972        self.timing_cache = os.path.join(self.cached_folder, self.timing_cache)973 974        # set device975        self.torch_device = self._execution_device976        logger.warning(f"Running inference on device: {self.torch_device}")977 978        # load models979        self.__loadModels()980 981        # build engines982        self.engine = build_engines(983            self.models,984            self.engine_dir,985            self.onnx_dir,986            self.onnx_opset,987            opt_image_height=self.image_height,988            opt_image_width=self.image_width,989            force_engine_rebuild=self.force_engine_rebuild,990            static_batch=self.build_static_batch,991            static_shape=not self.build_dynamic_shape,992            timing_cache=self.timing_cache,993        )994 995        return self996 997    def __initialize_timesteps(self, num_inference_steps, strength):998        self.scheduler.set_timesteps(num_inference_steps)999        offset = self.scheduler.config.steps_offset if hasattr(self.scheduler, "steps_offset") else 01000        init_timestep = int(num_inference_steps * strength) + offset1001        init_timestep = min(init_timestep, num_inference_steps)1002        t_start = max(num_inference_steps - init_timestep + offset, 0)1003        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :].to(self.torch_device)1004        return timesteps, num_inference_steps - t_start1005 1006    def __preprocess_images(self, batch_size, images=()):1007        init_images = []1008        for image in images:1009            image = image.to(self.torch_device).float()1010            image = image.repeat(batch_size, 1, 1, 1)1011            init_images.append(image)1012        return tuple(init_images)1013 1014    def __encode_image(self, init_image):1015        init_latents = runEngine(self.engine["vae_encoder"], {"images": init_image}, self.stream)["latent"]1016        init_latents = 0.18215 * init_latents1017        return init_latents1018 1019    def __encode_prompt(self, prompt, negative_prompt):1020        r"""1021        Encodes the prompt into text encoder hidden states.1022 1023        Args:1024             prompt (`str` or `List[str]`, *optional*):1025                prompt to be encoded1026            negative_prompt (`str` or `List[str]`, *optional*):1027                The prompt or prompts not to guide the image generation. If not defined, one has to pass1028                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.1029                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).1030        """1031        # Tokenize prompt1032        text_input_ids = (1033            self.tokenizer(1034                prompt,1035                padding="max_length",1036                max_length=self.tokenizer.model_max_length,1037                truncation=True,1038                return_tensors="pt",1039            )1040            .input_ids.type(torch.int32)1041            .to(self.torch_device)1042        )1043 1044        # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt1045        text_embeddings = runEngine(self.engine["clip"], {"input_ids": text_input_ids}, self.stream)[1046            "text_embeddings"1047        ].clone()1048 1049        # Tokenize negative prompt1050        uncond_input_ids = (1051            self.tokenizer(1052                negative_prompt,1053                padding="max_length",1054                max_length=self.tokenizer.model_max_length,1055                truncation=True,1056                return_tensors="pt",1057            )1058            .input_ids.type(torch.int32)1059            .to(self.torch_device)1060        )1061        uncond_embeddings = runEngine(self.engine["clip"], {"input_ids": uncond_input_ids}, self.stream)[1062            "text_embeddings"1063        ]1064 1065        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance1066        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)1067 1068        return text_embeddings1069 1070    def __denoise_latent(1071        self, latents, text_embeddings, timesteps=None, step_offset=0, mask=None, masked_image_latents=None1072    ):1073        if not isinstance(timesteps, torch.Tensor):1074            timesteps = self.scheduler.timesteps1075        for step_index, timestep in enumerate(timesteps):1076            # Expand the latents if we are doing classifier free guidance1077            latent_model_input = torch.cat([latents] * 2)1078            latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)1079            if isinstance(mask, torch.Tensor):1080                latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)1081 1082            # Predict the noise residual1083            timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep1084 1085            noise_pred = runEngine(1086                self.engine["unet"],1087                {"sample": latent_model_input, "timestep": timestep_float, "encoder_hidden_states": text_embeddings},1088                self.stream,1089            )["latent"]1090 1091            # Perform guidance1092            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)1093            noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)1094 1095            latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample1096 1097        latents = 1.0 / 0.18215 * latents1098        return latents1099 1100    def __decode_latent(self, latents):1101        images = runEngine(self.engine["vae"], {"latent": latents}, self.stream)["images"]1102        images = (images / 2 + 0.5).clamp(0, 1)1103        return images.cpu().permute(0, 2, 3, 1).float().numpy()1104 1105    def __loadResources(self, image_height, image_width, batch_size):1106        self.stream = cudart.cudaStreamCreate()[1]1107 1108        # Allocate buffers for TensorRT engine bindings1109        for model_name, obj in self.models.items():1110            self.engine[model_name].allocate_buffers(1111                shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.torch_device1112            )1113 1114    @torch.no_grad()1115    def __call__(1116        self,1117        prompt: Union[str, List[str]] = None,1118        image: Union[torch.Tensor, PIL.Image.Image] = None,1119        mask_image: Union[torch.Tensor, PIL.Image.Image] = None,1120        strength: float = 1.0,1121        num_inference_steps: int = 50,1122        guidance_scale: float = 7.5,1123        negative_prompt: Optional[Union[str, List[str]]] = None,1124        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,1125    ):1126        r"""1127        Function invoked when calling the pipeline for generation.1128 1129        Args:1130            prompt (`str` or `List[str]`, *optional*):1131                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.1132                instead.1133            image (`PIL.Image.Image`):1134                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will1135                be masked out with `mask_image` and repainted according to `prompt`.1136            mask_image (`PIL.Image.Image`):1137                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be1138                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted1139                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)1140                instead of 3, so the expected shape would be `(B, H, W, 1)`.1141            strength (`float`, *optional*, defaults to 0.8):1142                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`1143                will be used as a starting point, adding more noise to it the larger the `strength`. The number of1144                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will1145                be maximum and the denoising process will run for the full number of iterations specified in1146                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.1147            num_inference_steps (`int`, *optional*, defaults to 50):1148                The number of denoising steps. More denoising steps usually lead to a higher quality image at the1149                expense of slower inference.1150            guidance_scale (`float`, *optional*, defaults to 7.5):1151                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598).1152                `guidance_scale` is defined as `w` of equation 2. of [Imagen1153                Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale >1154                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1155                usually at the expense of lower image quality.1156            negative_prompt (`str` or `List[str]`, *optional*):1157                The prompt or prompts not to guide the image generation. If not defined, one has to pass1158                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.1159                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).1160            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):1161                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)1162                to make generation deterministic.1163 1164        """1165        self.generator = generator1166        self.denoising_steps = num_inference_steps1167        self._guidance_scale = guidance_scale1168 1169        # Pre-compute latent input scales and linear multistep coefficients1170        self.scheduler.set_timesteps(self.denoising_steps, device=self.torch_device)1171 1172        # Define call parameters1173        if prompt is not None and isinstance(prompt, str):1174            batch_size = 11175            prompt = [prompt]1176        elif prompt is not None and isinstance(prompt, list):1177            batch_size = len(prompt)1178        else:1179            raise ValueError(f"Expected prompt to be of type list or str but got {type(prompt)}")1180 1181        if negative_prompt is None:1182            negative_prompt = [""] * batch_size1183 1184        if negative_prompt is not None and isinstance(negative_prompt, str):1185            negative_prompt = [negative_prompt]1186 1187        assert len(prompt) == len(negative_prompt)1188 1189        if batch_size > self.max_batch_size:1190            raise ValueError(1191                f"Batch size {len(prompt)} is larger than allowed {self.max_batch_size}. If dynamic shape is used, then maximum batch size is 4"1192            )1193 1194        # Validate image dimensions1195        mask_width, mask_height = mask_image.size1196        if mask_height != self.image_height or mask_width != self.image_width:1197            raise ValueError(1198                f"Input image height and width {self.image_height} and {self.image_width} are not equal to "1199                f"the respective dimensions of the mask image {mask_height} and {mask_width}"1200            )

Showing the first 1,200 of 1273 lines. Download the file for the rest.