Delete utilities.py

utilities.py

@@ -1,537 +0,0 @@
-#
-# Copyright 2022 The HuggingFace Inc. team.
-# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-from collections import OrderedDict
-from copy import copy
-import numpy as np
-import os
-import math
-from PIL import Image
-from polygraphy.backend.common import bytes_from_path
-from polygraphy.backend.trt import CreateConfig, Profile
-from polygraphy.backend.trt import engine_from_bytes, engine_from_network, network_from_onnx_path, save_engine
-from polygraphy.backend.trt import util as trt_util
-from polygraphy import cuda
-import random
-from scipy import integrate
-import tensorrt as trt
-import torch
-
-TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
-
-class Engine():
-    def __init__(
-        self,
-        model_name,
-        engine_dir,
-    ):
-        self.engine_path = os.path.join(engine_dir, model_name+'.plan')
-        self.engine = None
-        self.context = None
-        self.buffers = OrderedDict()
-        self.tensors = OrderedDict()
-
-    def __del__(self):
-        [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray) ]
-        del self.engine
-        del self.context
-        del self.buffers
-        del self.tensors
-
-    def build(self, onnx_path, fp16, input_profile=None, enable_preview=False):
-        print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
-        p = Profile()
-        if input_profile:
-            for name, dims in input_profile.items():
-                assert len(dims) == 3
-                p.add(name, min=dims[0], opt=dims[1], max=dims[2])
-
-        preview_features = []
-        if enable_preview:
-            trt_version = [int(i) for i in trt.__version__.split(".")]
-            # FASTER_DYNAMIC_SHAPES_0805 should only be used for TRT 8.5.1 or above.
-            if trt_version[0] > 8 or \
-                (trt_version[0] == 8 and (trt_version[1] > 5 or (trt_version[1] == 5 and trt_version[2] >= 1))):
-                preview_features = [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
-
-        engine = engine_from_network(network_from_onnx_path(onnx_path), config=CreateConfig(fp16=fp16, profiles=[p],
-            preview_features=preview_features))
-        save_engine(engine, path=self.engine_path)
-
-    def activate(self):
-        print(f"Loading TensorRT engine: {self.engine_path}")
-        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
-        self.context = self.engine.create_execution_context()
-
-    def allocate_buffers(self, shape_dict=None, device='cuda'):
-        for idx in range(trt_util.get_bindings_per_profile(self.engine)):
-            binding = self.engine[idx]
-            if shape_dict and binding in shape_dict:
-                shape = shape_dict[binding]
-            else:
-                shape = self.engine.get_binding_shape(binding)
-            dtype = trt_util.np_dtype_from_trt(self.engine.get_binding_dtype(binding))
-            if self.engine.binding_is_input(binding):
-                self.context.set_binding_shape(idx, shape)
-            # Workaround to convert np dtype to torch
-            np_type_tensor = np.empty(shape=[], dtype=dtype)
-            torch_type_tensor = torch.from_numpy(np_type_tensor)
-            tensor = torch.empty(tuple(shape), dtype=torch_type_tensor.dtype).to(device=device)
-            self.tensors[binding] = tensor
-            self.buffers[binding] = cuda.DeviceView(ptr=tensor.data_ptr(), shape=shape, dtype=dtype)
-
-    def infer(self, feed_dict, stream):
-        start_binding, end_binding = trt_util.get_active_profile_bindings(self.context)
-        # shallow copy of ordered dict
-        device_buffers = copy(self.buffers)
-        for name, buf in feed_dict.items():
-            assert isinstance(buf, cuda.DeviceView)
-            device_buffers[name] = buf
-        bindings = [0] * start_binding + [buf.ptr for buf in device_buffers.values()]
-        noerror = self.context.execute_async_v2(bindings=bindings, stream_handle=stream.ptr)
-        if not noerror:
-            raise ValueError(f"ERROR: inference failed.")
-
-        return self.tensors
-
-class LMSDiscreteScheduler():
-    def __init__(
-        self,
-        device = 'cuda',
-        beta_start = 0.00085,
-        beta_end = 0.012,
-        num_train_timesteps = 1000,
-    ):
-        self.num_train_timesteps = num_train_timesteps
-        self.order = 4
-
-        self.beta_start = beta_start
-        self.beta_end = beta_end
-        betas = (torch.linspace(beta_start**0.5, beta_end**0.5, self.num_train_timesteps, dtype=torch.float32) ** 2)
-        alphas = 1.0 - betas
-        self.alphas_cumprod = torch.cumprod(alphas, dim=0)
-
-        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
-        sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
-        self.sigmas = torch.from_numpy(sigmas)
-
-        # standard deviation of the initial noise distribution
-        self.init_noise_sigma = self.sigmas.max()
-
-        self.device = device
-
-    def set_timesteps(self, steps):
-        self.num_inference_steps = steps
-
-        timesteps = np.linspace(0, self.num_train_timesteps - 1, steps, dtype=float)[::-1].copy()
-        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
-        sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
-        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
-        self.sigmas = torch.from_numpy(sigmas).to(device=self.device)
-
-        # Move all timesteps to correct device beforehand
-        self.timesteps = torch.from_numpy(timesteps).to(device=self.device).float()
-        self.derivatives = []
-
-    def scale_model_input(self, sample: torch.FloatTensor, idx, *args, **kwargs) -> torch.FloatTensor:
-        return sample * self.latent_scales[idx]
-
-    def configure(self):
-        order = self.order
-        self.lms_coeffs = []
-        self.latent_scales = [1./((sigma**2 + 1) ** 0.5) for sigma in self.sigmas]
-
-        def get_lms_coefficient(order, t, current_order):
-            """
-            Compute a linear multistep coefficient.
-            """
-            def lms_derivative(tau):
-                prod = 1.0
-                for k in range(order):
-                    if current_order == k:
-                        continue
-                    prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
-                return prod
-            integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0]
-            return integrated_coeff
-
-        for step_index in range(self.num_inference_steps):
-            order = min(step_index + 1, order)
-            self.lms_coeffs.append([get_lms_coefficient(order, step_index, curr_order) for curr_order in range(order)])
-
-    def step(self, output, latents, idx, timestep):
-        # compute the previous noisy sample x_t -> x_t-1
-        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
-        sigma = self.sigmas[idx]
-        pred_original_sample = latents - sigma * output
-        # 2. Convert to an ODE derivative
-        derivative = (latents - pred_original_sample) / sigma
-        self.derivatives.append(derivative)
-        if len(self.derivatives) > self.order:
-            self.derivatives.pop(0)
-        # 3. Compute previous sample based on the derivatives path
-        prev_sample = latents + sum(
-            coeff * derivative for coeff, derivative in zip(self.lms_coeffs[idx], reversed(self.derivatives))
-        )
-
-        return prev_sample
-
-class DPMScheduler():
-    def __init__(
-        self,
-        beta_start = 0.00085,
-        beta_end = 0.012,
-        num_train_timesteps = 1000,
-        solver_order = 2,
-        predict_epsilon = True,
-        thresholding = False,
-        dynamic_thresholding_ratio = 0.995,
-        sample_max_value = 1.0,
-        algorithm_type = "dpmsolver++",
-        solver_type = "midpoint",
-        lower_order_final = True,
-        device = 'cuda',
-    ):
-        # this schedule is very specific to the latent diffusion model.
-        self.betas = (
-            torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
-        )
-
-        self.device = device
-        self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
-        # Currently we only support VP-type noise schedule
-        self.alpha_t = torch.sqrt(self.alphas_cumprod)
-        self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
-        self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
-
-        # standard deviation of the initial noise distribution
-        self.init_noise_sigma = 1.0
-
-        self.algorithm_type = algorithm_type
-        self.predict_epsilon = predict_epsilon
-        self.thresholding = thresholding
-        self.dynamic_thresholding_ratio = dynamic_thresholding_ratio
-        self.sample_max_value = sample_max_value
-        self.lower_order_final = lower_order_final
-
-        # settings for DPM-Solver
-        if algorithm_type not in ["dpmsolver", "dpmsolver++"]:
-            raise NotImplementedError(f"{algorithm_type} does is not implemented for {self.__class__}")
-        if solver_type not in ["midpoint", "heun"]:
-            raise NotImplementedError(f"{solver_type} does is not implemented for {self.__class__}")
-
-        # setable values
-        self.num_inference_steps = None
-        self.solver_order = solver_order
-        self.num_train_timesteps = num_train_timesteps
-        self.solver_type = solver_type
-
-        self.first_order_first_coef = []
-        self.first_order_second_coef = []
-
-        self.second_order_first_coef = []
-        self.second_order_second_coef = []
-        self.second_order_third_coef = []
-
-        self.third_order_first_coef = []
-        self.third_order_second_coef = []
-        self.third_order_third_coef = []
-        self.third_order_fourth_coef = []
-
-    def scale_model_input(self, sample: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor:
-        return sample
-
-    def configure(self):
-        lower_order_nums = 0
-        for step_index in range(self.num_inference_steps):
-            step_idx = step_index
-            timestep = self.timesteps[step_idx]
-
-            prev_timestep = 0 if step_idx == len(self.timesteps) - 1 else self.timesteps[step_idx + 1]
-
-            self.dpm_solver_first_order_coefs_precompute(timestep, prev_timestep)
-
-            timestep_list = [self.timesteps[step_index - 1], timestep]
-            self.multistep_dpm_solver_second_order_coefs_precompute(timestep_list, prev_timestep)
-                
-            timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep]
-            self.multistep_dpm_solver_third_order_coefs_precompute(timestep_list, prev_timestep)
-
-            if lower_order_nums < self.solver_order:
-                lower_order_nums += 1
-
-    def dpm_solver_first_order_coefs_precompute(self, timestep, prev_timestep):
-        lambda_t, lambda_s = self.lambda_t[prev_timestep], self.lambda_t[timestep]
-        alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep]
-        sigma_t, sigma_s = self.sigma_t[prev_timestep], self.sigma_t[timestep]
-        h = lambda_t - lambda_s
-        if self.algorithm_type == "dpmsolver++":
-            self.first_order_first_coef.append(sigma_t / sigma_s)
-            self.first_order_second_coef.append(alpha_t * (torch.exp(-h) - 1.0))
-        elif self.algorithm_type == "dpmsolver":
-            self.first_order_first_coef.append(alpha_t / alpha_s)
-            self.first_order_second_coef.append(sigma_t * (torch.exp(h) - 1.0))
-
-    def multistep_dpm_solver_second_order_coefs_precompute(self, timestep_list, prev_timestep):
-        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
-        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
-        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
-        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
-        h = lambda_t - lambda_s0
-        if self.algorithm_type == "dpmsolver++":
-            # See https://arxiv.org/abs/2211.01095 for detailed derivations
-            if self.solver_type == "midpoint": 
-                self.second_order_first_coef.append(sigma_t / sigma_s0)
-                self.second_order_second_coef.append((alpha_t * (torch.exp(-h) - 1.0)))
-                self.second_order_third_coef.append(0.5 * (alpha_t * (torch.exp(-h) - 1.0)))
-            elif self.solver_type == "heun":
-                self.second_order_first_coef.append(sigma_t / sigma_s0)
-                self.second_order_second_coef.append((alpha_t * (torch.exp(-h) - 1.0)))
-                self.second_order_third_coef.append(alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0))
-        elif self.algorithm_type == "dpmsolver":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            if self.solver_type == "midpoint":
-                self.second_order_first_coef.append(alpha_t / alpha_s0)
-                self.second_order_second_coef.append((sigma_t * (torch.exp(h) - 1.0)))
-                self.second_order_third_coef.append(0.5 * (sigma_t * (torch.exp(h) - 1.0)))
-            elif self.solver_type == "heun":
-                self.second_order_first_coef.append(alpha_t / alpha_s0)
-                self.second_order_second_coef.append((sigma_t * (torch.exp(h) - 1.0)))
-                self.second_order_third_coef.append((sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)))
-
-    def multistep_dpm_solver_third_order_coefs_precompute(self, timestep_list, prev_timestep):
-        t, s0 = prev_timestep, timestep_list[-1]
-        lambda_t, lambda_s0 = (
-            self.lambda_t[t],
-            self.lambda_t[s0]
-        )
-        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
-        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
-        h = lambda_t - lambda_s0
-        if self.algorithm_type == "dpmsolver++":
-            self.third_order_first_coef.append(sigma_t / sigma_s0)
-            self.third_order_second_coef.append(alpha_t * (torch.exp(-h) - 1.0))
-            self.third_order_third_coef.append(alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0))
-            self.third_order_fourth_coef.append(alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5))
-        elif self.algorithm_type == "dpmsolver":
-            self.third_order_first_coef.append(alpha_t / alpha_s0)
-            self.third_order_second_coef.append(sigma_t * (torch.exp(h) - 1.0))
-            self.third_order_third_coef.append(sigma_t * ((torch.exp(h) - 1.0) / h - 1.0))
-            self.third_order_fourth_coef.append(sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5))
-
-    def set_timesteps(self, num_inference_steps):
-        self.num_inference_steps = num_inference_steps
-        timesteps = (
-            np.linspace(0, self.num_train_timesteps - 1, num_inference_steps + 1)
-            .round()[::-1][:-1]
-            .copy()
-            .astype(np.int32)
-        )
-        self.timesteps = torch.from_numpy(timesteps).to(self.device)
-        self.model_outputs = [
-            None,
-        ] * self.solver_order
-        self.lower_order_nums = 0
-
-    def convert_model_output(
-        self, model_output, timestep, sample
-    ):
-        # DPM-Solver++ needs to solve an integral of the data prediction model.
-        if self.algorithm_type == "dpmsolver++":
-            if self.predict_epsilon:
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
-                x0_pred = (sample - sigma_t * model_output) / alpha_t
-            else:
-                x0_pred = model_output
-            if self.thresholding:
-                # Dynamic thresholding in https://arxiv.org/abs/2205.11487
-                dynamic_max_val = torch.quantile(
-                    torch.abs(x0_pred).reshape((x0_pred.shape[0], -1)), self.dynamic_thresholding_ratio, dim=1
-                )
-                dynamic_max_val = torch.maximum(
-                    dynamic_max_val,
-                    self.sample_max_value * torch.ones_like(dynamic_max_val).to(dynamic_max_val.device),
-                )[(...,) + (None,) * (x0_pred.ndim - 1)]
-                x0_pred = torch.clamp(x0_pred, -dynamic_max_val, dynamic_max_val) / dynamic_max_val
-            return x0_pred
-        # DPM-Solver needs to solve an integral of the noise prediction model.
-        elif self.algorithm_type == "dpmsolver":
-            if self.predict_epsilon:
-                return model_output
-            else:
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
-                epsilon = (sample - alpha_t * model_output) / sigma_t
-                return epsilon
-
-    def dpm_solver_first_order_update(
-        self,
-        idx,
-        model_output,
-        sample
-    ):
-        first_coef = self.first_order_first_coef[idx]
-        second_coef = self.first_order_second_coef[idx]
-
-        if self.algorithm_type == "dpmsolver++":
-            x_t = first_coef * sample - second_coef * model_output
-        elif self.algorithm_type == "dpmsolver":
-            x_t = first_coef * sample - second_coef * model_output
-        return x_t
-
-    def multistep_dpm_solver_second_order_update(
-        self,
-        idx,
-        model_output_list,
-        timestep_list,
-        prev_timestep,
-        sample
-    ):
-        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
-        m0, m1 = model_output_list[-1], model_output_list[-2]
-        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
-        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
-        r0 = h_0 / h
-        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
-
-        first_coef = self.second_order_first_coef[idx]
-        second_coef = self.second_order_second_coef[idx]
-        third_coef = self.second_order_third_coef[idx]
-
-        if self.algorithm_type == "dpmsolver++":
-            # See https://arxiv.org/abs/2211.01095 for detailed derivations
-            if self.solver_type == "midpoint":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    - third_coef * D1
-                )
-            elif self.solver_type == "heun":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    + third_coef * D1
-                )
-        elif self.algorithm_type == "dpmsolver":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            if self.solver_type == "midpoint":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    - third_coef * D1
-                )
-            elif self.solver_type == "heun":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    - third_coef * D1
-                )
-        return x_t
-
-    def multistep_dpm_solver_third_order_update(
-        self,
-        idx,
-        model_output_list,
-        timestep_list,
-        prev_timestep,
-        sample
-    ):
-        t, s0, s1, s2 = prev_timestep, timestep_list[-1], timestep_list[-2], timestep_list[-3]
-        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
-        lambda_t, lambda_s0, lambda_s1, lambda_s2 = (
-            self.lambda_t[t],
-            self.lambda_t[s0],
-            self.lambda_t[s1],
-            self.lambda_t[s2],
-        )
-        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
-        r0, r1 = h_0 / h, h_1 / h
-        D0 = m0
-        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
-        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
-        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
-
-        first_coef = self.third_order_first_coef[idx]
-        second_coef = self.third_order_second_coef[idx]
-        third_coef = self.third_order_third_coef[idx]
-        fourth_coef = self.third_order_fourth_coef[idx]
-
-        if self.algorithm_type == "dpmsolver++":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            x_t = (
-                first_coef * sample
-                - second_coef * D0
-                + third_coef * D1
-                - fourth_coef * D2
-            )
-        elif self.algorithm_type == "dpmsolver":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            x_t = (
-                first_coef * sample
-                - second_coef * D0
-                - third_coef * D1
-                - fourth_coef * D2
-            )
-        return x_t
-
-    def step(self, output, latents, step_index, timestep):
-        if self.num_inference_steps is None:
-            raise ValueError(
-                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
-            )
-
-        prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1]
-        lower_order_final = (
-            (step_index == len(self.timesteps) - 1) and self.lower_order_final and len(self.timesteps) < 15
-        )
-        lower_order_second = (
-            (step_index == len(self.timesteps) - 2) and self.lower_order_final and len(self.timesteps) < 15
-        )
-
-        output = self.convert_model_output(output, timestep, latents)
-        for i in range(self.solver_order - 1):
-            self.model_outputs[i] = self.model_outputs[i + 1]
-        self.model_outputs[-1] = output
-
-        if self.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
-            prev_sample = self.dpm_solver_first_order_update(step_index, output, latents)
-        elif self.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
-            timestep_list = [self.timesteps[step_index - 1], timestep]
-            prev_sample = self.multistep_dpm_solver_second_order_update(
-                step_index, self.model_outputs, timestep_list, prev_timestep, latents
-            )
-        else:
-            timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep]
-            prev_sample = self.multistep_dpm_solver_third_order_update(
-                step_index, self.model_outputs, timestep_list, prev_timestep, latents
-            )
-
-        if self.lower_order_nums < self.solver_order:
-            self.lower_order_nums += 1
-
-        return prev_sample
-
-def save_image(images, image_path_dir, image_name_prefix):
-    """
-    Save the generated images to png files.
-    """
-    images = ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy()
-    for i in range(images.shape[0]):
-        image_path  = os.path.join(image_path_dir, image_name_prefix+str(i+1)+'-'+str(random.randint(1000,9999))+'.png')
-        print(f"Saving image {i+1} / {images.shape[0]} to: {image_path}")
-        Image.fromarray(images[i]).save(image_path)