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_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 hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 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 hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample 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 = hasattr(unet.config, "_diffusers_version") and version.parse(761            version.parse(unet.config._diffusers_version).base_version762        ) < version.parse("0.9.0.dev0")763        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64764        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:765            deprecation_message = (766                "The configuration file of the unet has set the default `sample_size` to smaller than"767                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"768                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"769                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"770                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"771                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"772                " in the config might lead to incorrect results in future versions. If you have downloaded this"773                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"774                " the `unet/config.json` file"775            )776            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)777            new_config = dict(unet.config)778            new_config["sample_size"] = 64779            unet._internal_dict = FrozenDict(new_config)780 781        self.register_modules(782            vae=vae,783            text_encoder=text_encoder,784            tokenizer=tokenizer,785            unet=unet,786            scheduler=scheduler,787            safety_checker=safety_checker,788            feature_extractor=feature_extractor,789            image_encoder=image_encoder,790        )791 792        self.stages = stages793        self.image_height, self.image_width = image_height, image_width794        self.inpaint = True795        self.onnx_opset = onnx_opset796        self.onnx_dir = onnx_dir797        self.engine_dir = engine_dir798        self.force_engine_rebuild = force_engine_rebuild799        self.timing_cache = timing_cache800        self.build_static_batch = False801        self.build_dynamic_shape = False802 803        self.max_batch_size = max_batch_size804        # TODO: Restrict batch size to 4 for larger image dimensions as a WAR for TensorRT limitation.805        if self.build_dynamic_shape or self.image_height > 512 or self.image_width > 512:806            self.max_batch_size = 4807 808        self.stream = None  # loaded in loadResources()809        self.models = {}  # loaded in __loadModels()810        self.engine = {}  # loaded in build_engines()811 812        self.vae.forward = self.vae.decode813        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)814        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)815        self.register_to_config(requires_safety_checker=requires_safety_checker)816 817    def __loadModels(self):818        # Load pipeline models819        self.embedding_dim = self.text_encoder.config.hidden_size820        models_args = {821            "device": self.torch_device,822            "max_batch_size": self.max_batch_size,823            "embedding_dim": self.embedding_dim,824            "inpaint": self.inpaint,825        }826        if "clip" in self.stages:827            self.models["clip"] = make_CLIP(self.text_encoder, **models_args)828        if "unet" in self.stages:829            self.models["unet"] = make_UNet(self.unet, **models_args, unet_dim=self.unet.config.in_channels)830        if "vae" in self.stages:831            self.models["vae"] = make_VAE(self.vae, **models_args)832        if "vae_encoder" in self.stages:833            self.models["vae_encoder"] = make_VAEEncoder(self.vae, **models_args)834 835    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline836 837    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):838        if isinstance(generator, list):839            image_latents = [840                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])841                for i in range(image.shape[0])842            ]843            image_latents = torch.cat(image_latents, dim=0)844        else:845            image_latents = retrieve_latents(self.vae.encode(image), generator=generator)846 847        image_latents = self.vae.config.scaling_factor * image_latents848 849        return image_latents850 851    def prepare_latents(852        self,853        batch_size,854        num_channels_latents,855        height,856        width,857        dtype,858        device,859        generator,860        latents=None,861        image=None,862        timestep=None,863        is_strength_max=True,864        return_noise=False,865        return_image_latents=False,866    ):867        shape = (868            batch_size,869            num_channels_latents,870            int(height) // self.vae_scale_factor,871            int(width) // self.vae_scale_factor,872        )873        if isinstance(generator, list) and len(generator) != batch_size:874            raise ValueError(875                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"876                f" size of {batch_size}. Make sure the batch size matches the length of the generators."877            )878 879        if (image is None or timestep is None) and not is_strength_max:880            raise ValueError(881                "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise."882                "However, either the image or the noise timestep has not been provided."883            )884 885        if return_image_latents or (latents is None and not is_strength_max):886            image = image.to(device=device, dtype=dtype)887 888            if image.shape[1] == 4:889                image_latents = image890            else:891                image_latents = self._encode_vae_image(image=image, generator=generator)892            image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)893 894        if latents is None:895            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)896            # if strength is 1. then initialise the latents to noise, else initial to image + noise897            latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep)898            # if pure noise then scale the initial latents by the  Scheduler's init sigma899            latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents900        else:901            noise = latents.to(device)902            latents = noise * self.scheduler.init_noise_sigma903 904        outputs = (latents,)905 906        if return_noise:907            outputs += (noise,)908 909        if return_image_latents:910            outputs += (image_latents,)911 912        return outputs913 914    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker915    def run_safety_checker(916        self, image: Union[torch.Tensor, PIL.Image.Image], device: torch.device, dtype: torch.dtype917    ) -> Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]:918        r"""919        Runs the safety checker on the given image.920        Args:921            image (Union[torch.Tensor, PIL.Image.Image]): The input image to be checked.922            device (torch.device): The device to run the safety checker on.923            dtype (torch.dtype): The data type of the input image.924        Returns:925            (image, has_nsfw_concept) Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]: A tuple containing the processed image and926            a boolean indicating whether the image has a NSFW (Not Safe for Work) concept.927        """928        if self.safety_checker is None:929            has_nsfw_concept = None930        else:931            if torch.is_tensor(image):932                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")933            else:934                feature_extractor_input = self.image_processor.numpy_to_pil(image)935            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)936            image, has_nsfw_concept = self.safety_checker(937                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)938            )939        return image, has_nsfw_concept940 941    @classmethod942    @validate_hf_hub_args943    def set_cached_folder(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):944        cache_dir = kwargs.pop("cache_dir", None)945        proxies = kwargs.pop("proxies", None)946        local_files_only = kwargs.pop("local_files_only", False)947        token = kwargs.pop("token", None)948        revision = kwargs.pop("revision", None)949 950        cls.cached_folder = (951            pretrained_model_name_or_path952            if os.path.isdir(pretrained_model_name_or_path)953            else snapshot_download(954                pretrained_model_name_or_path,955                cache_dir=cache_dir,956                proxies=proxies,957                local_files_only=local_files_only,958                token=token,959                revision=revision,960            )961        )962 963    def to(self, torch_device: Optional[Union[str, torch.device]] = None, silence_dtype_warnings: bool = False):964        super().to(torch_device, silence_dtype_warnings=silence_dtype_warnings)965 966        self.onnx_dir = os.path.join(self.cached_folder, self.onnx_dir)967        self.engine_dir = os.path.join(self.cached_folder, self.engine_dir)968        self.timing_cache = os.path.join(self.cached_folder, self.timing_cache)969 970        # set device971        self.torch_device = self._execution_device972        logger.warning(f"Running inference on device: {self.torch_device}")973 974        # load models975        self.__loadModels()976 977        # build engines978        self.engine = build_engines(979            self.models,980            self.engine_dir,981            self.onnx_dir,982            self.onnx_opset,983            opt_image_height=self.image_height,984            opt_image_width=self.image_width,985            force_engine_rebuild=self.force_engine_rebuild,986            static_batch=self.build_static_batch,987            static_shape=not self.build_dynamic_shape,988            timing_cache=self.timing_cache,989        )990 991        return self992 993    def __initialize_timesteps(self, num_inference_steps, strength):994        self.scheduler.set_timesteps(num_inference_steps)995        offset = self.scheduler.config.steps_offset if hasattr(self.scheduler, "steps_offset") else 0996        init_timestep = int(num_inference_steps * strength) + offset997        init_timestep = min(init_timestep, num_inference_steps)998        t_start = max(num_inference_steps - init_timestep + offset, 0)999        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :].to(self.torch_device)1000        return timesteps, num_inference_steps - t_start1001 1002    def __preprocess_images(self, batch_size, images=()):1003        init_images = []1004        for image in images:1005            image = image.to(self.torch_device).float()1006            image = image.repeat(batch_size, 1, 1, 1)1007            init_images.append(image)1008        return tuple(init_images)1009 1010    def __encode_image(self, init_image):1011        init_latents = runEngine(self.engine["vae_encoder"], {"images": init_image}, self.stream)["latent"]1012        init_latents = 0.18215 * init_latents1013        return init_latents1014 1015    def __encode_prompt(self, prompt, negative_prompt):1016        r"""1017        Encodes the prompt into text encoder hidden states.1018 1019        Args:1020             prompt (`str` or `List[str]`, *optional*):1021                prompt to be encoded1022            negative_prompt (`str` or `List[str]`, *optional*):1023                The prompt or prompts not to guide the image generation. If not defined, one has to pass1024                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.1025                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).1026        """1027        # Tokenize prompt1028        text_input_ids = (1029            self.tokenizer(1030                prompt,1031                padding="max_length",1032                max_length=self.tokenizer.model_max_length,1033                truncation=True,1034                return_tensors="pt",1035            )1036            .input_ids.type(torch.int32)1037            .to(self.torch_device)1038        )1039 1040        # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt1041        text_embeddings = runEngine(self.engine["clip"], {"input_ids": text_input_ids}, self.stream)[1042            "text_embeddings"1043        ].clone()1044 1045        # Tokenize negative prompt1046        uncond_input_ids = (1047            self.tokenizer(1048                negative_prompt,1049                padding="max_length",1050                max_length=self.tokenizer.model_max_length,1051                truncation=True,1052                return_tensors="pt",1053            )1054            .input_ids.type(torch.int32)1055            .to(self.torch_device)1056        )1057        uncond_embeddings = runEngine(self.engine["clip"], {"input_ids": uncond_input_ids}, self.stream)[1058            "text_embeddings"1059        ]1060 1061        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance1062        text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)1063 1064        return text_embeddings1065 1066    def __denoise_latent(1067        self, latents, text_embeddings, timesteps=None, step_offset=0, mask=None, masked_image_latents=None1068    ):1069        if not isinstance(timesteps, torch.Tensor):1070            timesteps = self.scheduler.timesteps1071        for step_index, timestep in enumerate(timesteps):1072            # Expand the latents if we are doing classifier free guidance1073            latent_model_input = torch.cat([latents] * 2)1074            latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)1075            if isinstance(mask, torch.Tensor):1076                latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)1077 1078            # Predict the noise residual1079            timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep1080 1081            noise_pred = runEngine(1082                self.engine["unet"],1083                {"sample": latent_model_input, "timestep": timestep_float, "encoder_hidden_states": text_embeddings},1084                self.stream,1085            )["latent"]1086 1087            # Perform guidance1088            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)1089            noise_pred = noise_pred_uncond + self._guidance_scale * (noise_pred_text - noise_pred_uncond)1090 1091            latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample1092 1093        latents = 1.0 / 0.18215 * latents1094        return latents1095 1096    def __decode_latent(self, latents):1097        images = runEngine(self.engine["vae"], {"latent": latents}, self.stream)["images"]1098        images = (images / 2 + 0.5).clamp(0, 1)1099        return images.cpu().permute(0, 2, 3, 1).float().numpy()1100 1101    def __loadResources(self, image_height, image_width, batch_size):1102        self.stream = cudart.cudaStreamCreate()[1]1103 1104        # Allocate buffers for TensorRT engine bindings1105        for model_name, obj in self.models.items():1106            self.engine[model_name].allocate_buffers(1107                shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.torch_device1108            )1109 1110    @torch.no_grad()1111    def __call__(1112        self,1113        prompt: Union[str, List[str]] = None,1114        image: Union[torch.Tensor, PIL.Image.Image] = None,1115        mask_image: Union[torch.Tensor, PIL.Image.Image] = None,1116        strength: float = 1.0,1117        num_inference_steps: int = 50,1118        guidance_scale: float = 7.5,1119        negative_prompt: Optional[Union[str, List[str]]] = None,1120        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,1121    ):1122        r"""1123        Function invoked when calling the pipeline for generation.1124 1125        Args:1126            prompt (`str` or `List[str]`, *optional*):1127                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.1128                instead.1129            image (`PIL.Image.Image`):1130                `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will1131                be masked out with `mask_image` and repainted according to `prompt`.1132            mask_image (`PIL.Image.Image`):1133                `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be1134                repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted1135                to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L)1136                instead of 3, so the expected shape would be `(B, H, W, 1)`.1137            strength (`float`, *optional*, defaults to 0.8):1138                Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`1139                will be used as a starting point, adding more noise to it the larger the `strength`. The number of1140                denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will1141                be maximum and the denoising process will run for the full number of iterations specified in1142                `num_inference_steps`. A value of 1, therefore, essentially ignores `image`.1143            num_inference_steps (`int`, *optional*, defaults to 50):1144                The number of denoising steps. More denoising steps usually lead to a higher quality image at the1145                expense of slower inference.1146            guidance_scale (`float`, *optional*, defaults to 7.5):1147                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).1148                `guidance_scale` is defined as `w` of equation 2. of [Imagen1149                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >1150                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,1151                usually at the expense of lower image quality.1152            negative_prompt (`str` or `List[str]`, *optional*):1153                The prompt or prompts not to guide the image generation. If not defined, one has to pass1154                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.1155                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).1156            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):1157                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)1158                to make generation deterministic.1159 1160        """1161        self.generator = generator1162        self.denoising_steps = num_inference_steps1163        self._guidance_scale = guidance_scale1164 1165        # Pre-compute latent input scales and linear multistep coefficients1166        self.scheduler.set_timesteps(self.denoising_steps, device=self.torch_device)1167 1168        # Define call parameters1169        if prompt is not None and isinstance(prompt, str):1170            batch_size = 11171            prompt = [prompt]1172        elif prompt is not None and isinstance(prompt, list):1173            batch_size = len(prompt)1174        else:1175            raise ValueError(f"Expected prompt to be of type list or str but got {type(prompt)}")1176 1177        if negative_prompt is None:1178            negative_prompt = [""] * batch_size1179 1180        if negative_prompt is not None and isinstance(negative_prompt, str):1181            negative_prompt = [negative_prompt]1182 1183        assert len(prompt) == len(negative_prompt)1184 1185        if batch_size > self.max_batch_size:1186            raise ValueError(1187                f"Batch size {len(prompt)} is larger than allowed {self.max_batch_size}. If dynamic shape is used, then maximum batch size is 4"1188            )1189 1190        # Validate image dimensions1191        mask_width, mask_height = mask_image.size1192        if mask_height != self.image_height or mask_width != self.image_width:1193            raise ValueError(1194                f"Input image height and width {self.image_height} and {self.image_width} are not equal to "1195                f"the respective dimensions of the mask image {mask_height} and {mask_width}"1196            )1197 1198        # load resources1199        self.__loadResources(self.image_height, self.image_width, batch_size)1200 

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