feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
This commit is contained in:
project6-dev
5
ixformer_sdk/contrib/DeepCache/__init__.py
Normal file
5
ixformer_sdk/contrib/DeepCache/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from .sd.pipeline_stable_diffusion import StableDiffusionPipeline
|
||||
from .sdxl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
|
||||
from .sdxl.pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipeline
|
||||
|
||||
from .sd.pipeline_text_to_video_zero import TextToVideoZeroPipeline
|
||||
0
ixformer_sdk/contrib/DeepCache/ddpm/__init__.py
Normal file
0
ixformer_sdk/contrib/DeepCache/ddpm/__init__.py
Normal file
180
ixformer_sdk/contrib/DeepCache/ddpm/ddim.py
Normal file
180
ixformer_sdk/contrib/DeepCache/ddpm/ddim.py
Normal file
@@ -0,0 +1,180 @@
|
||||
import argparse
|
||||
import traceback
|
||||
import shutil
|
||||
import logging
|
||||
import yaml
|
||||
import random
|
||||
import sys
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from ddpm.utils.logging import Logger, EmptyLogger
|
||||
from ddpm.utils.tools import set_random_seed
|
||||
from accelerate import Accelerator, DistributedDataParallelKwargs
|
||||
|
||||
torch.set_printoptions(sci_mode=False)
|
||||
|
||||
def dict2namespace(config):
|
||||
namespace = argparse.Namespace()
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
new_value = dict2namespace(value)
|
||||
else:
|
||||
new_value = value
|
||||
setattr(namespace, key, new_value)
|
||||
return namespace
|
||||
|
||||
def parse_args_and_config():
|
||||
parser = argparse.ArgumentParser(description=globals()["__doc__"])
|
||||
|
||||
parser.add_argument(
|
||||
"--config", type=str, required=True, help="Path to the config file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed", type=int, default=1234, help="Random seed")
|
||||
parser.add_argument(
|
||||
"--exp", type=str, default="exp", help="Path for saving running related data."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--test", action="store_true", help="Whether to test the model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample", action="store_true", help="Whether to produce samples from the model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_folder", type=str, default="images", help="folder name for storing the sampled images"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--fid", action="store_true"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--interpolation", action="store_true"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resume_training", action="store_true", help="Whether to resume training"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ni", action="store_true", help="No interaction. Suitable for Slurm Job launcher",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_pretrained", action="store_true"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_type", type=str, default="generalized", help="sampling approach (generalized or ddpm_noisy)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip_type", type=str, default="uniform", help="skip according to (uniform or quadratic)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--timesteps", type=int, default=1000, help="number of steps involved"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eta", type=float, default=0.0, help="eta used to control the variances of sigma",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dyn", action="store_true", help="whether to activate the dynamic train/inference"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sequence", action="store_true"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--select_step", type=int, default=None
|
||||
)
|
||||
parser.add_argument(
|
||||
"--select_depth", type=int, default=None
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--cache", action="store_true"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache_interval", type=int, default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--non_uniform", action="store_true"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pow", type=float, default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--center", type=int, default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--branch", type=int, default=None,
|
||||
)
|
||||
args = parser.parse_args()
|
||||
# parse config file
|
||||
with open(args.config, "r") as f:
|
||||
config = yaml.safe_load(f)
|
||||
new_config = dict2namespace(config)
|
||||
new_config.select_step = args.select_step
|
||||
new_config.select_depth = args.select_depth
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
return args, new_config
|
||||
|
||||
|
||||
def main():
|
||||
args, config = parse_args_and_config()
|
||||
|
||||
if args.dyn:
|
||||
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||||
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
|
||||
else:
|
||||
accelerator = Accelerator()
|
||||
args.accelerator = accelerator
|
||||
|
||||
#log_root_dir = "{}_runtime_log".format(args.config[8:-4])
|
||||
log_root_dir = "runtime_log"
|
||||
dataset = args.config[8:-4]
|
||||
if args.cache:
|
||||
if args.non_uniform:
|
||||
sub_dir_name = "{}_{}_cache_{}_pow_{}_center_{}".format(dataset, args.exp, args.cache_interval, args.pow, args.center)
|
||||
else:
|
||||
sub_dir_name = "{}_{}_cache_{}".format(dataset, args.exp, args.cache_interval)
|
||||
else:
|
||||
sub_dir_name = "{}".format(args.exp)
|
||||
|
||||
if accelerator.is_main_process:
|
||||
logger = Logger(
|
||||
root_dir=log_root_dir,
|
||||
sub_name=sub_dir_name,
|
||||
config=args.__dict__,
|
||||
append=(args.sample == True)
|
||||
)
|
||||
args.logger = logger
|
||||
|
||||
args.logger.log("Writing log file to {}".format(args.logger.sub_dir))
|
||||
args.logger.log("Exp instance PID = {}".format(os.getpid()))
|
||||
else:
|
||||
args.logger = EmptyLogger(
|
||||
root_dir=log_root_dir,
|
||||
sub_name=sub_dir_name,
|
||||
)
|
||||
|
||||
args.image_folder = args.logger.setup_image_folder("{}".format(args.image_folder))
|
||||
|
||||
args.seed += accelerator.process_index
|
||||
# set random seed
|
||||
set_random_seed(args.seed)
|
||||
try:
|
||||
if args.cache:
|
||||
from ddpm.runners.deepcache import Diffusion
|
||||
runner = Diffusion(args, config)
|
||||
runner.sample()
|
||||
else:
|
||||
from ddpm.runners.diffusion import Diffusion
|
||||
runner = Diffusion(args, config)
|
||||
runner.sample()
|
||||
except Exception:
|
||||
logging.error(traceback.format_exc())
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
361
ixformer_sdk/contrib/DeepCache/ddpm/fid.py
Normal file
361
ixformer_sdk/contrib/DeepCache/ddpm/fid.py
Normal file
@@ -0,0 +1,361 @@
|
||||
"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
|
||||
|
||||
The FID metric calculates the distance between two distributions of images.
|
||||
Typically, we have summary statistics (mean & covariance matrix) of one
|
||||
of these distributions, while the 2nd distribution is given by a GAN.
|
||||
|
||||
When run as a stand-alone program, it compares the distribution of
|
||||
images that are stored as PNG/JPEG at a specified location with a
|
||||
distribution given by summary statistics (in pickle format).
|
||||
|
||||
The FID is calculated by assuming that X_1 and X_2 are the activations of
|
||||
the pool_3 layer of the inception net for generated samples and real world
|
||||
samples respectively.
|
||||
|
||||
See --help to see further details.
|
||||
|
||||
Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
|
||||
of Tensorflow
|
||||
|
||||
Copyright 2018 Institute of Bioinformatics, JKU Linz
|
||||
|
||||
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.
|
||||
"""
|
||||
import os
|
||||
import pathlib
|
||||
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms as TF
|
||||
from PIL import Image
|
||||
from scipy import linalg
|
||||
from torch.nn.functional import adaptive_avg_pool2d
|
||||
|
||||
try:
|
||||
from tqdm import tqdm
|
||||
except ImportError:
|
||||
# If tqdm is not available, provide a mock version of it
|
||||
def tqdm(x):
|
||||
return x
|
||||
|
||||
from pytorch_fid.inception import InceptionV3
|
||||
|
||||
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--batch-size', type=int, default=50,
|
||||
help='Batch size to use')
|
||||
parser.add_argument('--dataset_name', type=str, default=None)
|
||||
parser.add_argument('--num-workers', type=int,
|
||||
help=('Number of processes to use for data loading. '
|
||||
'Defaults to `min(8, num_cpus)`'))
|
||||
parser.add_argument('--device', type=str, default=None,
|
||||
help='Device to use. Like cuda, cuda:0 or cpu')
|
||||
parser.add_argument('--dims', type=int, default=2048,
|
||||
choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
|
||||
help=('Dimensionality of Inception features to use. '
|
||||
'By default, uses pool3 features'))
|
||||
parser.add_argument('--num_samples', type=int, default=None,
|
||||
help=('Number of samples for FID estimation'))
|
||||
parser.add_argument('--res', type=int, default=None,
|
||||
help=('Resolutions of samples for FID estimation'))
|
||||
parser.add_argument('--save-stats', action='store_true',
|
||||
help=('Generate an npz archive from a directory of samples. '
|
||||
'The first path is used as input and the second as output.'))
|
||||
|
||||
parser.add_argument('--path', type=str, nargs=2,
|
||||
help=('Paths to the generated images or '
|
||||
'to .npz statistic files'))
|
||||
|
||||
|
||||
IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
|
||||
'tif', 'tiff', 'webp'}
|
||||
|
||||
|
||||
class ImagePathDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, files, transforms=None):
|
||||
self.files = files
|
||||
self.transforms = transforms
|
||||
|
||||
def __len__(self):
|
||||
return len(self.files)
|
||||
|
||||
def __getitem__(self, i):
|
||||
path = self.files[i]
|
||||
img = Image.open(path).convert('RGB')
|
||||
if self.transforms is not None:
|
||||
img = self.transforms(img)
|
||||
return img
|
||||
|
||||
|
||||
def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
|
||||
num_workers=1, res=None, dataset_name=None):
|
||||
"""Calculates the activations of the pool_3 layer for all images.
|
||||
|
||||
Params:
|
||||
-- files : List of image files paths
|
||||
-- model : Instance of inception model
|
||||
-- batch_size : Batch size of images for the model to process at once.
|
||||
Make sure that the number of samples is a multiple of
|
||||
the batch size, otherwise some samples are ignored. This
|
||||
behavior is retained to match the original FID score
|
||||
implementation.
|
||||
-- dims : Dimensionality of features returned by Inception
|
||||
-- device : Device to run calculations
|
||||
-- num_workers : Number of parallel dataloader workers
|
||||
|
||||
Returns:
|
||||
-- A numpy array of dimension (num images, dims) that contains the
|
||||
activations of the given tensor when feeding inception with the
|
||||
query tensor.
|
||||
"""
|
||||
model.eval()
|
||||
|
||||
if batch_size > len(files):
|
||||
print(('Warning: batch size is bigger than the data size. '
|
||||
'Setting batch size to data size'))
|
||||
batch_size = len(files)
|
||||
|
||||
if res is None:
|
||||
trans = TF.ToTensor()
|
||||
else:
|
||||
if dataset_name == 'celeba':
|
||||
from switchable_diffusion.datasets import Crop
|
||||
print("In crop image: {}, {}".format(res, dataset_name))
|
||||
cx = 89
|
||||
cy = 121
|
||||
x1 = cy - 64
|
||||
x2 = cy + 64
|
||||
y1 = cx - 64
|
||||
y2 = cx + 64
|
||||
trans = TF.Compose([
|
||||
Crop(x1, x2, y1, y2),
|
||||
TF.Resize(res),
|
||||
TF.ToTensor(),
|
||||
])
|
||||
else:
|
||||
trans = TF.Compose([
|
||||
TF.Resize(res),
|
||||
TF.CenterCrop(res),
|
||||
TF.ToTensor()
|
||||
])
|
||||
|
||||
dataset = ImagePathDataset(files, transforms=trans)
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
num_workers=num_workers)
|
||||
|
||||
pred_arr = np.empty((len(files), dims))
|
||||
|
||||
start_idx = 0
|
||||
|
||||
for batch in tqdm(dataloader):
|
||||
batch = batch.to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(batch)[0]
|
||||
|
||||
# If model output is not scalar, apply global spatial average pooling.
|
||||
# This happens if you choose a dimensionality not equal 2048.
|
||||
if pred.size(2) != 1 or pred.size(3) != 1:
|
||||
pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
|
||||
|
||||
pred = pred.squeeze(3).squeeze(2).cpu().numpy()
|
||||
|
||||
pred_arr[start_idx:start_idx + pred.shape[0]] = pred
|
||||
|
||||
start_idx = start_idx + pred.shape[0]
|
||||
|
||||
return pred_arr
|
||||
|
||||
|
||||
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
|
||||
"""Numpy implementation of the Frechet Distance.
|
||||
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
|
||||
and X_2 ~ N(mu_2, C_2) is
|
||||
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
|
||||
|
||||
Stable version by Dougal J. Sutherland.
|
||||
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representative data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representative data set.
|
||||
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1)
|
||||
+ np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
|
||||
device='cpu', num_workers=1, res=None, dataset_name=None):
|
||||
"""Calculation of the statistics used by the FID.
|
||||
Params:
|
||||
-- files : List of image files paths
|
||||
-- model : Instance of inception model
|
||||
-- batch_size : The images numpy array is split into batches with
|
||||
batch size batch_size. A reasonable batch size
|
||||
depends on the hardware.
|
||||
-- dims : Dimensionality of features returned by Inception
|
||||
-- device : Device to run calculations
|
||||
-- num_workers : Number of parallel dataloader workers
|
||||
|
||||
Returns:
|
||||
-- mu : The mean over samples of the activations of the pool_3 layer of
|
||||
the inception model.
|
||||
-- sigma : The covariance matrix of the activations of the pool_3 layer of
|
||||
the inception model.
|
||||
"""
|
||||
act = get_activations(files, model, batch_size, dims, device, num_workers, res=res, dataset_name=dataset_name)
|
||||
mu = np.mean(act, axis=0)
|
||||
sigma = np.cov(act, rowvar=False)
|
||||
return mu, sigma
|
||||
|
||||
|
||||
def compute_statistics_of_path(path, model, batch_size, dims, device,
|
||||
num_workers=1, num_samples=None, res=None, dataset_name=None):
|
||||
if path.endswith('.npz'):
|
||||
with np.load(path) as f:
|
||||
m, s = f['mu'][:], f['sigma'][:]
|
||||
else:
|
||||
path = pathlib.Path(path)
|
||||
|
||||
files = sorted([file for ext in IMAGE_EXTENSIONS
|
||||
for file in path.glob('**/*.{}'.format(ext))])
|
||||
if num_samples is not None:
|
||||
#import random
|
||||
#files = random.sample(files, num_samples)
|
||||
files = files[:num_samples]
|
||||
print("Found %d files." % len(files))
|
||||
m, s = calculate_activation_statistics(files, model, batch_size,
|
||||
dims, device, num_workers, res=res, dataset_name=dataset_name)
|
||||
|
||||
return m, s
|
||||
|
||||
|
||||
def calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None):
|
||||
"""Calculates the FID of two paths"""
|
||||
for p in paths:
|
||||
if not os.path.exists(p):
|
||||
raise RuntimeError('Invalid path: %s' % p)
|
||||
|
||||
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
|
||||
|
||||
model = InceptionV3([block_idx]).to(device)
|
||||
|
||||
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
|
||||
dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
|
||||
m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,
|
||||
dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
|
||||
fid_value = calculate_frechet_distance(m1, s1, m2, s2)
|
||||
|
||||
return fid_value
|
||||
|
||||
|
||||
def save_fid_stats(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None):
|
||||
"""Calculates the FID of two paths"""
|
||||
if not os.path.exists(paths[0]):
|
||||
raise RuntimeError('Invalid path: %s' % paths[0])
|
||||
|
||||
if os.path.exists(paths[1]):
|
||||
raise RuntimeError('Existing output file: %s' % paths[1])
|
||||
|
||||
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
|
||||
|
||||
model = InceptionV3([block_idx]).to(device)
|
||||
|
||||
print(f"Saving statistics for {paths[0]}")
|
||||
|
||||
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
|
||||
dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
|
||||
|
||||
np.savez_compressed(paths[1], mu=m1, sigma=s1)
|
||||
|
||||
|
||||
def main():
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.device is None:
|
||||
device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
|
||||
else:
|
||||
device = torch.device(args.device)
|
||||
|
||||
if args.num_workers is None:
|
||||
try:
|
||||
num_cpus = len(os.sched_getaffinity(0))
|
||||
except AttributeError:
|
||||
# os.sched_getaffinity is not available under Windows, use
|
||||
# os.cpu_count instead (which may not return the *available* number
|
||||
# of CPUs).
|
||||
num_cpus = os.cpu_count()
|
||||
|
||||
num_workers = min(num_cpus, 8) if num_cpus is not None else 0
|
||||
else:
|
||||
num_workers = args.num_workers
|
||||
|
||||
if args.save_stats:
|
||||
save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers, num_samples=args.num_samples, res=args.res, dataset_name=args.dataset_name)
|
||||
return
|
||||
|
||||
fid_value = calculate_fid_given_paths(args.path,
|
||||
args.batch_size,
|
||||
device,
|
||||
args.dims,
|
||||
num_workers,
|
||||
num_samples=args.num_samples,
|
||||
res = args.res, dataset_name=args.dataset_name)
|
||||
print('FID: ', fid_value)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
559
ixformer_sdk/contrib/DeepCache/flops.py
Normal file
559
ixformer_sdk/contrib/DeepCache/flops.py
Normal file
@@ -0,0 +1,559 @@
|
||||
'''
|
||||
This opcounter is adapted from https://github.com/sovrasov/flops-counter.pytorch and https://github.com/Lyken17/pytorch-OpCounter
|
||||
|
||||
Copyright (C) 2021 Sovrasov V. - All Rights Reserved
|
||||
* You may use, distribute and modify this code under the
|
||||
* terms of the MIT license.
|
||||
* You should have received a copy of the MIT license with
|
||||
* this file. If not visit https://opensource.org/licenses/MIT
|
||||
'''
|
||||
import os
|
||||
import yaml
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
has_timm = False
|
||||
|
||||
from diffusers.models.lora import LoRACompatibleLinear, LoRACompatibleConv
|
||||
|
||||
@torch.no_grad()
|
||||
def count_ops_and_params(model, example_inputs, layer_wise=False):
|
||||
global CUSTOM_MODULES_MAPPING
|
||||
ori_model = model
|
||||
model = copy.deepcopy(model) # deepcopy to avoid changing the original model
|
||||
flops_model = add_flops_counting_methods(model)
|
||||
flops_model.eval()
|
||||
flops_model.start_flops_count(ost=sys.stdout, verbose=False,
|
||||
ignore_list=[])
|
||||
if isinstance(example_inputs, (tuple, list)):
|
||||
_ = flops_model(*example_inputs)
|
||||
elif isinstance(example_inputs, dict):
|
||||
_ = flops_model(**example_inputs)
|
||||
else:
|
||||
_ = flops_model(example_inputs)
|
||||
flops_count, params_count, _layer_flops, _layer_params = flops_model.compute_average_flops_cost()
|
||||
layer_flops = {}
|
||||
layer_params = {}
|
||||
|
||||
for m_name, m in model.named_modules():
|
||||
layer_flops[m_name] = _layer_flops.get(m)
|
||||
layer_params[m_name] = _layer_params.get(m)
|
||||
if layer_wise:
|
||||
space = 30 - len(m_name)
|
||||
print("Layer {}: {} MACs = {:.4f} G, Params = {:.4f} M, MACs% = {:.2f}".format(
|
||||
m_name, ' ' * space, layer_flops[m_name]/1e9, layer_params[m_name] / 1e6, 100 * layer_flops[m_name] / flops_count
|
||||
))
|
||||
|
||||
flops_model.stop_flops_count()
|
||||
CUSTOM_MODULES_MAPPING = {}
|
||||
#if layer_wise:
|
||||
# return flops_count, params_count, layer_flops, layer_params
|
||||
return flops_count, params_count
|
||||
|
||||
def empty_flops_counter_hook(module, input, output):
|
||||
module.__flops__ += 0
|
||||
|
||||
|
||||
def upsample_flops_counter_hook(module, input, output):
|
||||
output_size = output[0]
|
||||
batch_size = output_size.shape[0]
|
||||
output_elements_count = batch_size
|
||||
for val in output_size.shape[1:]:
|
||||
output_elements_count *= val
|
||||
module.__flops__ += int(output_elements_count)
|
||||
|
||||
|
||||
def relu_flops_counter_hook(module, input, output):
|
||||
active_elements_count = output.numel()
|
||||
module.__flops__ += int(active_elements_count)
|
||||
|
||||
|
||||
def linear_flops_counter_hook(module, input, output):
|
||||
input = input[0]
|
||||
# pytorch checks dimensions, so here we don't care much
|
||||
output_last_dim = output.shape[-1]
|
||||
bias_flops = output_last_dim if module.bias is not None else 0
|
||||
module.__flops__ += int(np.prod(input.shape) * output_last_dim + bias_flops)
|
||||
|
||||
|
||||
def pool_flops_counter_hook(module, input, output):
|
||||
input = input[0]
|
||||
module.__flops__ += int(np.prod(input.shape))
|
||||
|
||||
|
||||
def bn_flops_counter_hook(module, input, output):
|
||||
input = input[0]
|
||||
|
||||
batch_flops = np.prod(input.shape)
|
||||
if module.affine:
|
||||
batch_flops *= 2
|
||||
module.__flops__ += int(batch_flops)
|
||||
|
||||
def ln_flops_counter_hook(module, input, output):
|
||||
input = input[0]
|
||||
batch_flops = np.prod(input.shape)
|
||||
if module.elementwise_affine:
|
||||
batch_flops *= 2
|
||||
module.__flops__ += int(batch_flops)
|
||||
|
||||
def conv_flops_counter_hook(conv_module, input, output):
|
||||
# Can have multiple inputs, getting the first one
|
||||
input = input[0]
|
||||
|
||||
batch_size = input.shape[0]
|
||||
output_dims = list(output.shape[2:])
|
||||
|
||||
kernel_dims = list(conv_module.kernel_size)
|
||||
in_channels = conv_module.in_channels
|
||||
out_channels = conv_module.out_channels
|
||||
groups = conv_module.groups
|
||||
|
||||
filters_per_channel = out_channels // groups
|
||||
conv_per_position_flops = int(np.prod(kernel_dims)) * \
|
||||
in_channels * filters_per_channel
|
||||
|
||||
active_elements_count = batch_size * int(np.prod(output_dims))
|
||||
|
||||
overall_conv_flops = conv_per_position_flops * active_elements_count
|
||||
|
||||
bias_flops = 0
|
||||
|
||||
if conv_module.bias is not None:
|
||||
|
||||
bias_flops = out_channels * active_elements_count
|
||||
|
||||
overall_flops = overall_conv_flops + bias_flops
|
||||
|
||||
conv_module.__flops__ += int(overall_flops)
|
||||
|
||||
|
||||
def rnn_flops(flops, rnn_module, w_ih, w_hh, input_size):
|
||||
# matrix matrix mult ih state and internal state
|
||||
flops += w_ih.shape[0]*w_ih.shape[1]
|
||||
# matrix matrix mult hh state and internal state
|
||||
flops += w_hh.shape[0]*w_hh.shape[1]
|
||||
if isinstance(rnn_module, (nn.RNN, nn.RNNCell)):
|
||||
# add both operations
|
||||
flops += rnn_module.hidden_size
|
||||
elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)):
|
||||
# hadamard of r
|
||||
flops += rnn_module.hidden_size
|
||||
# adding operations from both states
|
||||
flops += rnn_module.hidden_size*3
|
||||
# last two hadamard product and add
|
||||
flops += rnn_module.hidden_size*3
|
||||
elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)):
|
||||
# adding operations from both states
|
||||
flops += rnn_module.hidden_size*4
|
||||
# two hadamard product and add for C state
|
||||
flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
|
||||
# final hadamard
|
||||
flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
|
||||
return flops
|
||||
|
||||
|
||||
def rnn_flops_counter_hook(rnn_module, input, output):
|
||||
"""
|
||||
Takes into account batch goes at first position, contrary
|
||||
to pytorch common rule (but actually it doesn't matter).
|
||||
If sigmoid and tanh are hard, only a comparison FLOPS should be accurate
|
||||
"""
|
||||
flops = 0
|
||||
# input is a tuple containing a sequence to process and (optionally) hidden state
|
||||
inp = input[0]
|
||||
batch_size = inp[0].shape[0]
|
||||
seq_length = inp[0].shape[1]
|
||||
num_layers = rnn_module.num_layers
|
||||
|
||||
for i in range(num_layers):
|
||||
w_ih = rnn_module.__getattr__('weight_ih_l' + str(i))
|
||||
w_hh = rnn_module.__getattr__('weight_hh_l' + str(i))
|
||||
if i == 0:
|
||||
input_size = rnn_module.input_size
|
||||
else:
|
||||
input_size = rnn_module.hidden_size
|
||||
flops = rnn_flops(flops, rnn_module, w_ih, w_hh, input_size)
|
||||
if rnn_module.bias:
|
||||
b_ih = rnn_module.__getattr__('bias_ih_l' + str(i))
|
||||
b_hh = rnn_module.__getattr__('bias_hh_l' + str(i))
|
||||
flops += b_ih.shape[0] + b_hh.shape[0]
|
||||
|
||||
flops *= batch_size
|
||||
flops *= seq_length
|
||||
if rnn_module.bidirectional:
|
||||
flops *= 2
|
||||
rnn_module.__flops__ += int(flops)
|
||||
|
||||
|
||||
def rnn_cell_flops_counter_hook(rnn_cell_module, input, output):
|
||||
flops = 0
|
||||
inp = input[0]
|
||||
batch_size = inp.shape[0]
|
||||
w_ih = rnn_cell_module.__getattr__('weight_ih')
|
||||
w_hh = rnn_cell_module.__getattr__('weight_hh')
|
||||
input_size = inp.shape[1]
|
||||
flops = rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size)
|
||||
if rnn_cell_module.bias:
|
||||
b_ih = rnn_cell_module.__getattr__('bias_ih')
|
||||
b_hh = rnn_cell_module.__getattr__('bias_hh')
|
||||
flops += b_ih.shape[0] + b_hh.shape[0]
|
||||
|
||||
flops *= batch_size
|
||||
rnn_cell_module.__flops__ += int(flops)
|
||||
|
||||
|
||||
def multihead_attention_counter_hook(multihead_attention_module, input, output):
|
||||
flops = 0
|
||||
q, k, v = input
|
||||
|
||||
batch_first = multihead_attention_module.batch_first \
|
||||
if hasattr(multihead_attention_module, 'batch_first') else False
|
||||
if batch_first:
|
||||
batch_size = q.shape[0]
|
||||
len_idx = 1
|
||||
else:
|
||||
batch_size = q.shape[1]
|
||||
len_idx = 0
|
||||
|
||||
dim_idx = 2
|
||||
|
||||
qdim = q.shape[dim_idx]
|
||||
kdim = k.shape[dim_idx]
|
||||
vdim = v.shape[dim_idx]
|
||||
|
||||
qlen = q.shape[len_idx]
|
||||
klen = k.shape[len_idx]
|
||||
vlen = v.shape[len_idx]
|
||||
|
||||
num_heads = multihead_attention_module.num_heads
|
||||
assert qdim == multihead_attention_module.embed_dim
|
||||
|
||||
if multihead_attention_module.kdim is None:
|
||||
assert kdim == qdim
|
||||
if multihead_attention_module.vdim is None:
|
||||
assert vdim == qdim
|
||||
|
||||
flops = 0
|
||||
|
||||
# Q scaling
|
||||
flops += qlen * qdim
|
||||
# Initial projections
|
||||
flops += (
|
||||
(qlen * qdim * qdim) # QW
|
||||
+ (klen * kdim * kdim) # KW
|
||||
+ (vlen * vdim * vdim) # VW
|
||||
)
|
||||
if multihead_attention_module.in_proj_bias is not None:
|
||||
flops += (qlen + klen + vlen) * qdim
|
||||
# attention heads: scale, matmul, softmax, matmul
|
||||
qk_head_dim = qdim // num_heads
|
||||
v_head_dim = vdim // num_heads
|
||||
|
||||
head_flops = (
|
||||
(qlen * klen * qk_head_dim) # QK^T
|
||||
+ (qlen * klen) # softmax
|
||||
+ (qlen * klen * v_head_dim) # AV
|
||||
)
|
||||
flops += num_heads * head_flops
|
||||
# final projection, bias is always enabled
|
||||
flops += qlen * vdim * (vdim + 1)
|
||||
flops *= batch_size
|
||||
multihead_attention_module.__flops__ += int(flops)
|
||||
|
||||
def timm_multihead_attention_counter_hook(multihead_attention_module, input, output):
|
||||
flops = 0
|
||||
|
||||
q, k, v = input[0], input[0], input[0]
|
||||
input_dim = input[0].shape[2]
|
||||
input_len = input[0].shape[1]
|
||||
batch_size = input[0].shape[0]
|
||||
|
||||
kdim = qdim = vdim = multihead_attention_module.qkv.out_features//3
|
||||
qlen = klen = vlen = input_len
|
||||
|
||||
num_heads = multihead_attention_module.num_heads
|
||||
assert qdim == multihead_attention_module.head_dim * multihead_attention_module.num_heads
|
||||
|
||||
flops = 0
|
||||
# Q scaling
|
||||
flops += qlen * qdim
|
||||
# Initial projections
|
||||
flops += (
|
||||
(qlen * input_dim * qdim) # QW
|
||||
+ (klen * input_dim * kdim) # KW
|
||||
+ (vlen * input_dim * vdim) # VW
|
||||
)
|
||||
|
||||
if multihead_attention_module.qkv.bias is not None:
|
||||
flops += (qlen + klen + vlen) * qdim
|
||||
# attention heads: scale, matmul, softmax, matmul
|
||||
qk_head_dim = qdim // num_heads
|
||||
v_head_dim = vdim // num_heads
|
||||
|
||||
head_flops = (
|
||||
(qlen * klen * qk_head_dim) # QK^T
|
||||
+ (qlen * klen) # softmax
|
||||
+ (qlen * klen * v_head_dim) # AV
|
||||
)
|
||||
flops += num_heads * head_flops
|
||||
# final projection, bias is always enabled
|
||||
flops += qlen * vdim * (vdim + 1)
|
||||
flops *= batch_size
|
||||
multihead_attention_module.__flops__ += int(flops)
|
||||
|
||||
|
||||
|
||||
CUSTOM_MODULES_MAPPING = {}
|
||||
|
||||
MODULES_MAPPING = {
|
||||
# convolutions
|
||||
nn.Conv1d: conv_flops_counter_hook,
|
||||
nn.Conv2d: conv_flops_counter_hook,
|
||||
nn.Conv3d: conv_flops_counter_hook,
|
||||
LoRACompatibleConv: conv_flops_counter_hook,
|
||||
# activations
|
||||
nn.ReLU: relu_flops_counter_hook,
|
||||
nn.PReLU: relu_flops_counter_hook,
|
||||
nn.ELU: relu_flops_counter_hook,
|
||||
nn.LeakyReLU: relu_flops_counter_hook,
|
||||
nn.ReLU6: relu_flops_counter_hook,
|
||||
# poolings
|
||||
nn.MaxPool1d: pool_flops_counter_hook,
|
||||
nn.AvgPool1d: pool_flops_counter_hook,
|
||||
nn.AvgPool2d: pool_flops_counter_hook,
|
||||
nn.MaxPool2d: pool_flops_counter_hook,
|
||||
nn.MaxPool3d: pool_flops_counter_hook,
|
||||
nn.AvgPool3d: pool_flops_counter_hook,
|
||||
nn.AdaptiveMaxPool1d: pool_flops_counter_hook,
|
||||
nn.AdaptiveAvgPool1d: pool_flops_counter_hook,
|
||||
nn.AdaptiveMaxPool2d: pool_flops_counter_hook,
|
||||
nn.AdaptiveAvgPool2d: pool_flops_counter_hook,
|
||||
nn.AdaptiveMaxPool3d: pool_flops_counter_hook,
|
||||
nn.AdaptiveAvgPool3d: pool_flops_counter_hook,
|
||||
# BNs
|
||||
nn.BatchNorm1d: bn_flops_counter_hook,
|
||||
nn.BatchNorm2d: bn_flops_counter_hook,
|
||||
nn.BatchNorm3d: bn_flops_counter_hook,
|
||||
|
||||
nn.InstanceNorm1d: bn_flops_counter_hook,
|
||||
nn.InstanceNorm2d: bn_flops_counter_hook,
|
||||
nn.InstanceNorm3d: bn_flops_counter_hook,
|
||||
nn.GroupNorm: bn_flops_counter_hook,
|
||||
nn.LayerNorm: ln_flops_counter_hook,
|
||||
# FC
|
||||
nn.Linear: linear_flops_counter_hook,
|
||||
LoRACompatibleLinear: linear_flops_counter_hook,
|
||||
# Upscale
|
||||
nn.Upsample: upsample_flops_counter_hook,
|
||||
# Deconvolution
|
||||
nn.ConvTranspose1d: conv_flops_counter_hook,
|
||||
nn.ConvTranspose2d: conv_flops_counter_hook,
|
||||
nn.ConvTranspose3d: conv_flops_counter_hook,
|
||||
# RNN
|
||||
nn.RNN: rnn_flops_counter_hook,
|
||||
nn.GRU: rnn_flops_counter_hook,
|
||||
nn.LSTM: rnn_flops_counter_hook,
|
||||
nn.RNNCell: rnn_cell_flops_counter_hook,
|
||||
nn.LSTMCell: rnn_cell_flops_counter_hook,
|
||||
nn.GRUCell: rnn_cell_flops_counter_hook,
|
||||
nn.MultiheadAttention: multihead_attention_counter_hook
|
||||
}
|
||||
|
||||
if has_timm:
|
||||
MODULES_MAPPING.update(
|
||||
{
|
||||
timm.models.vision_transformer.Attention: timm_multihead_attention_counter_hook,
|
||||
}
|
||||
)
|
||||
|
||||
if hasattr(nn, 'GELU'):
|
||||
MODULES_MAPPING[nn.GELU] = relu_flops_counter_hook
|
||||
|
||||
|
||||
import sys
|
||||
from functools import partial
|
||||
import torch.nn as nn
|
||||
import copy
|
||||
|
||||
def accumulate_flops(self, layer_flops):
|
||||
if is_supported_instance(self):
|
||||
layer_flops[self] = self.__flops__
|
||||
return self.__flops__
|
||||
else:
|
||||
sum = 0
|
||||
for m in self.children():
|
||||
sum += m.accumulate_flops(layer_flops)
|
||||
layer_flops[self] = sum
|
||||
return sum
|
||||
|
||||
|
||||
def get_model_parameters_number(model):
|
||||
params_num = sum(p.numel() for p in model.parameters())
|
||||
return params_num
|
||||
|
||||
|
||||
def add_flops_counting_methods(net_main_module):
|
||||
# adding additional methods to the existing module object,
|
||||
# this is done this way so that each function has access to self object
|
||||
net_main_module.start_flops_count = start_flops_count.__get__(net_main_module)
|
||||
net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module)
|
||||
net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module)
|
||||
net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__(
|
||||
net_main_module)
|
||||
|
||||
net_main_module.reset_flops_count()
|
||||
|
||||
return net_main_module
|
||||
|
||||
def compute_average_flops_cost(self):
|
||||
"""
|
||||
A method that will be available after add_flops_counting_methods() is called
|
||||
on a desired net object.
|
||||
Returns current mean flops consumption per image.
|
||||
"""
|
||||
|
||||
for m in self.modules():
|
||||
m.accumulate_flops = accumulate_flops.__get__(m)
|
||||
|
||||
layer_flops = {}
|
||||
flops_sum = self.accumulate_flops(layer_flops)
|
||||
|
||||
for m in self.modules():
|
||||
if hasattr(m, 'accumulate_flops'):
|
||||
del m.accumulate_flops
|
||||
|
||||
layer_params = {}
|
||||
for m in self.modules():
|
||||
layer_params[m] = get_model_parameters_number(m)
|
||||
|
||||
params_sum = get_model_parameters_number(self)
|
||||
return flops_sum / self.__batch_counter__, params_sum, layer_flops, layer_params
|
||||
|
||||
|
||||
def start_flops_count(self, **kwargs):
|
||||
"""
|
||||
A method that will be available after add_flops_counting_methods() is called
|
||||
on a desired net object.
|
||||
Activates the computation of mean flops consumption per image.
|
||||
Call it before you run the network.
|
||||
"""
|
||||
add_batch_counter_hook_function(self)
|
||||
|
||||
seen_types = set()
|
||||
|
||||
def add_flops_counter_hook_function(module, ost, verbose, ignore_list):
|
||||
if type(module) in ignore_list:
|
||||
seen_types.add(type(module))
|
||||
if is_supported_instance(module):
|
||||
module.__params__ = 0
|
||||
elif is_supported_instance(module):
|
||||
if hasattr(module, '__flops_handle__'):
|
||||
return
|
||||
if type(module) in CUSTOM_MODULES_MAPPING:
|
||||
handle = module.register_forward_hook(
|
||||
CUSTOM_MODULES_MAPPING[type(module)])
|
||||
else:
|
||||
handle = module.register_forward_hook(MODULES_MAPPING[type(module)])
|
||||
module.__flops_handle__ = handle
|
||||
seen_types.add(type(module))
|
||||
else:
|
||||
if verbose and not type(module) in (nn.Sequential, nn.ModuleList) and \
|
||||
not type(module) in seen_types:
|
||||
print('Warning: module ' + type(module).__name__ +
|
||||
' is treated as a zero-op.', file=ost)
|
||||
seen_types.add(type(module))
|
||||
|
||||
self.apply(partial(add_flops_counter_hook_function, **kwargs))
|
||||
|
||||
|
||||
def stop_flops_count(self):
|
||||
"""
|
||||
A method that will be available after add_flops_counting_methods() is called
|
||||
on a desired net object.
|
||||
Stops computing the mean flops consumption per image.
|
||||
Call whenever you want to pause the computation.
|
||||
"""
|
||||
remove_batch_counter_hook_function(self)
|
||||
self.apply(remove_flops_counter_hook_function)
|
||||
self.apply(remove_flops_counter_variables)
|
||||
|
||||
|
||||
def reset_flops_count(self):
|
||||
"""
|
||||
A method that will be available after add_flops_counting_methods() is called
|
||||
on a desired net object.
|
||||
Resets statistics computed so far.
|
||||
"""
|
||||
add_batch_counter_variables_or_reset(self)
|
||||
self.apply(add_flops_counter_variable_or_reset)
|
||||
|
||||
|
||||
# ---- Internal functions
|
||||
def batch_counter_hook(module, input, output):
|
||||
batch_size = 1
|
||||
if len(input) > 0:
|
||||
# Can have multiple inputs, getting the first one
|
||||
input = input[0]
|
||||
batch_size = len(input)
|
||||
else:
|
||||
pass
|
||||
print('Warning! No positional inputs found for a module,'
|
||||
' assuming batch size is 1.')
|
||||
module.__batch_counter__ += batch_size
|
||||
|
||||
|
||||
def add_batch_counter_variables_or_reset(module):
|
||||
|
||||
module.__batch_counter__ = 0
|
||||
|
||||
|
||||
def add_batch_counter_hook_function(module):
|
||||
if hasattr(module, '__batch_counter_handle__'):
|
||||
return
|
||||
|
||||
handle = module.register_forward_hook(batch_counter_hook)
|
||||
module.__batch_counter_handle__ = handle
|
||||
|
||||
|
||||
def remove_batch_counter_hook_function(module):
|
||||
if hasattr(module, '__batch_counter_handle__'):
|
||||
module.__batch_counter_handle__.remove()
|
||||
del module.__batch_counter_handle__
|
||||
|
||||
|
||||
def add_flops_counter_variable_or_reset(module):
|
||||
if is_supported_instance(module):
|
||||
if hasattr(module, '__flops__') or hasattr(module, '__params__'):
|
||||
print('Warning: variables __flops__ or __params__ are already '
|
||||
'defined for the module' + type(module).__name__ +
|
||||
' ptflops can affect your code!')
|
||||
module.__ptflops_backup_flops__ = module.__flops__
|
||||
module.__ptflops_backup_params__ = module.__params__
|
||||
module.__flops__ = 0
|
||||
module.__params__ = get_model_parameters_number(module)
|
||||
|
||||
|
||||
def is_supported_instance(module):
|
||||
if type(module) in MODULES_MAPPING or type(module) in CUSTOM_MODULES_MAPPING:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def remove_flops_counter_hook_function(module):
|
||||
if is_supported_instance(module):
|
||||
if hasattr(module, '__flops_handle__'):
|
||||
module.__flops_handle__.remove()
|
||||
del module.__flops_handle__
|
||||
|
||||
|
||||
def remove_flops_counter_variables(module):
|
||||
if is_supported_instance(module):
|
||||
if hasattr(module, '__flops__'):
|
||||
del module.__flops__
|
||||
if hasattr(module, '__ptflops_backup_flops__'):
|
||||
module.__flops__ = module.__ptflops_backup_flops__
|
||||
if hasattr(module, '__params__'):
|
||||
del module.__params__
|
||||
if hasattr(module, '__ptflops_backup_params__'):
|
||||
module.__params__ = module.__ptflops_backup_params__
|
||||
|
||||
0
ixformer_sdk/contrib/DeepCache/sd/__init__.py
Normal file
0
ixformer_sdk/contrib/DeepCache/sd/__init__.py
Normal file
812
ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py
Normal file
812
ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py
Normal file
@@ -0,0 +1,812 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# 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.
|
||||
import time
|
||||
import inspect
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from packaging import version
|
||||
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from diffusers.configuration_utils import FrozenDict
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
|
||||
from diffusers.models import AutoencoderKL
|
||||
from diffusers.models.lora import adjust_lora_scale_text_encoder
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils import (
|
||||
deprecate,
|
||||
logging,
|
||||
replace_example_docstring,
|
||||
)
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
||||
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
|
||||
|
||||
from .unet_2d_condition import UNet2DConditionModel
|
||||
from .pipeline_utils import DiffusionPipeline
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
EXAMPLE_DOC_STRING = """
|
||||
Examples:
|
||||
```py
|
||||
>>> import torch
|
||||
>>> from diffusers import StableDiffusionPipeline
|
||||
|
||||
>>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
|
||||
>>> pipe = pipe.to("cuda")
|
||||
|
||||
>>> prompt = "a photo of an astronaut riding a horse on mars"
|
||||
>>> image = pipe(prompt).images[0]
|
||||
```
|
||||
"""
|
||||
|
||||
def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100):
|
||||
samples = []
|
||||
|
||||
while len(samples) < sample_size:
|
||||
# Sample from a Gaussian centered at n/2
|
||||
sample = int(np.random.normal(loc=n/2, scale=std_dev))
|
||||
|
||||
# Check if the sample is in bounds
|
||||
if 1 <= sample < n and sample not in samples:
|
||||
samples.append(sample)
|
||||
|
||||
return samples
|
||||
|
||||
def sample_from_quad(total_numbers, n_samples, pow=1.2):
|
||||
while pow > 1:
|
||||
# Generate linearly spaced values between 0 and a max value
|
||||
x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1)
|
||||
|
||||
# Raise these values to the power of 1.5 to get a non-linear distribution
|
||||
indices = np.unique(np.int32(x_values**pow))[:-1]
|
||||
if len(indices) == n_samples:
|
||||
break
|
||||
pow -=0.02
|
||||
if pow <= 1:
|
||||
raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
|
||||
return indices, pow
|
||||
|
||||
def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2):
|
||||
while pow > 1:
|
||||
# Generate linearly spaced values between 0 and a max value
|
||||
x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1)
|
||||
indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]]
|
||||
if len(indices) == n_samples:
|
||||
break
|
||||
pow -=0.02
|
||||
if pow <= 1:
|
||||
raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
|
||||
return indices, pow
|
||||
|
||||
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
|
||||
"""
|
||||
Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
|
||||
Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
|
||||
"""
|
||||
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
|
||||
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
|
||||
# rescale the results from guidance (fixes overexposure)
|
||||
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
|
||||
# mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
|
||||
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
|
||||
return noise_cfg
|
||||
|
||||
|
||||
class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin):
|
||||
r"""
|
||||
Pipeline for text-to-image generation using Stable Diffusion.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
|
||||
implemented for all pipelines (downloading, saving, running on a particular device, etc.).
|
||||
|
||||
The pipeline also inherits the following loading methods:
|
||||
- [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
|
||||
- [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
|
||||
- [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
|
||||
- [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
|
||||
|
||||
Args:
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
|
||||
text_encoder ([`~transformers.CLIPTextModel`]):
|
||||
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
|
||||
tokenizer ([`~transformers.CLIPTokenizer`]):
|
||||
A `CLIPTokenizer` to tokenize text.
|
||||
unet ([`UNet2DConditionModel`]):
|
||||
A `UNet2DConditionModel` to denoise the encoded image latents.
|
||||
scheduler ([`SchedulerMixin`]):
|
||||
A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
|
||||
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
|
||||
safety_checker ([`StableDiffusionSafetyChecker`]):
|
||||
Classification module that estimates whether generated images could be considered offensive or harmful.
|
||||
Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
|
||||
about a model's potential harms.
|
||||
feature_extractor ([`~transformers.CLIPImageProcessor`]):
|
||||
A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
|
||||
"""
|
||||
model_cpu_offload_seq = "text_encoder->unet->vae"
|
||||
_optional_components = ["safety_checker", "feature_extractor"]
|
||||
_exclude_from_cpu_offload = ["safety_checker"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder: CLIPTextModel,
|
||||
tokenizer: CLIPTokenizer,
|
||||
unet: UNet2DConditionModel,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
safety_checker: StableDiffusionSafetyChecker,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
requires_safety_checker: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
|
||||
deprecation_message = (
|
||||
f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
|
||||
f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
|
||||
"to update the config accordingly as leaving `steps_offset` might led to incorrect results"
|
||||
" in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
|
||||
" it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
|
||||
" file"
|
||||
)
|
||||
deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
|
||||
new_config = dict(scheduler.config)
|
||||
new_config["steps_offset"] = 1
|
||||
scheduler._internal_dict = FrozenDict(new_config)
|
||||
|
||||
if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
|
||||
deprecation_message = (
|
||||
f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
|
||||
" `clip_sample` should be set to False in the configuration file. Please make sure to update the"
|
||||
" config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
|
||||
" future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
|
||||
" nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
|
||||
)
|
||||
deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
|
||||
new_config = dict(scheduler.config)
|
||||
new_config["clip_sample"] = False
|
||||
scheduler._internal_dict = FrozenDict(new_config)
|
||||
|
||||
if safety_checker is None and requires_safety_checker:
|
||||
logger.warning(
|
||||
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
|
||||
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
|
||||
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
|
||||
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
|
||||
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
|
||||
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
|
||||
)
|
||||
|
||||
if safety_checker is not None and feature_extractor is None:
|
||||
raise ValueError(
|
||||
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
|
||||
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
|
||||
)
|
||||
|
||||
is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
|
||||
version.parse(unet.config._diffusers_version).base_version
|
||||
) < version.parse("0.9.0.dev0")
|
||||
is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
|
||||
if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
|
||||
deprecation_message = (
|
||||
"The configuration file of the unet has set the default `sample_size` to smaller than"
|
||||
" 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
|
||||
" following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
|
||||
" CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
|
||||
" \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
|
||||
" configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
|
||||
" in the config might lead to incorrect results in future versions. If you have downloaded this"
|
||||
" checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
|
||||
" the `unet/config.json` file"
|
||||
)
|
||||
deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
|
||||
new_config = dict(unet.config)
|
||||
new_config["sample_size"] = 64
|
||||
unet._internal_dict = FrozenDict(new_config)
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
text_encoder=text_encoder,
|
||||
tokenizer=tokenizer,
|
||||
unet=unet,
|
||||
scheduler=scheduler,
|
||||
safety_checker=safety_checker,
|
||||
feature_extractor=feature_extractor,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
self.register_to_config(requires_safety_checker=requires_safety_checker)
|
||||
|
||||
def enable_vae_slicing(self):
|
||||
r"""
|
||||
Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
|
||||
compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
|
||||
"""
|
||||
self.vae.enable_slicing()
|
||||
|
||||
def disable_vae_slicing(self):
|
||||
r"""
|
||||
Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
|
||||
computing decoding in one step.
|
||||
"""
|
||||
self.vae.disable_slicing()
|
||||
|
||||
def enable_vae_tiling(self):
|
||||
r"""
|
||||
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
|
||||
compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
|
||||
processing larger images.
|
||||
"""
|
||||
self.vae.enable_tiling()
|
||||
|
||||
def disable_vae_tiling(self):
|
||||
r"""
|
||||
Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
|
||||
computing decoding in one step.
|
||||
"""
|
||||
self.vae.disable_tiling()
|
||||
|
||||
def _encode_prompt(
|
||||
self,
|
||||
prompt,
|
||||
device,
|
||||
num_images_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
negative_prompt=None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
lora_scale: Optional[float] = None,
|
||||
):
|
||||
deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
|
||||
deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
|
||||
|
||||
prompt_embeds_tuple = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
negative_prompt=negative_prompt,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
# concatenate for backwards comp
|
||||
prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
|
||||
|
||||
return prompt_embeds
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt,
|
||||
device,
|
||||
num_images_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
negative_prompt=None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
lora_scale: Optional[float] = None,
|
||||
):
|
||||
r"""
|
||||
Encodes the prompt into text encoder hidden states.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
prompt to be encoded
|
||||
device: (`torch.device`):
|
||||
torch device
|
||||
num_images_per_prompt (`int`):
|
||||
number of images that should be generated per prompt
|
||||
do_classifier_free_guidance (`bool`):
|
||||
whether to use classifier free guidance or not
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
||||
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
|
||||
less than `1`).
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
||||
provided, text embeddings will be generated from `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
|
||||
argument.
|
||||
lora_scale (`float`, *optional*):
|
||||
A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
|
||||
"""
|
||||
# set lora scale so that monkey patched LoRA
|
||||
# function of text encoder can correctly access it
|
||||
if lora_scale is not None and isinstance(self, LoraLoaderMixin):
|
||||
self._lora_scale = lora_scale
|
||||
|
||||
# dynamically adjust the LoRA scale
|
||||
adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
|
||||
|
||||
if prompt is not None and isinstance(prompt, str):
|
||||
batch_size = 1
|
||||
elif prompt is not None and isinstance(prompt, list):
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
if prompt_embeds is None:
|
||||
# textual inversion: procecss multi-vector tokens if necessary
|
||||
if isinstance(self, TextualInversionLoaderMixin):
|
||||
prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=self.tokenizer.model_max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
|
||||
|
||||
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
|
||||
text_input_ids, untruncated_ids
|
||||
):
|
||||
removed_text = self.tokenizer.batch_decode(
|
||||
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
|
||||
)
|
||||
logger.warning(
|
||||
"The following part of your input was truncated because CLIP can only handle sequences up to"
|
||||
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
|
||||
)
|
||||
|
||||
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
|
||||
attention_mask = text_inputs.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
|
||||
prompt_embeds = self.text_encoder(
|
||||
text_input_ids.to(device),
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
prompt_embeds = prompt_embeds[0]
|
||||
|
||||
if self.text_encoder is not None:
|
||||
prompt_embeds_dtype = self.text_encoder.dtype
|
||||
elif self.unet is not None:
|
||||
prompt_embeds_dtype = self.unet.dtype
|
||||
else:
|
||||
prompt_embeds_dtype = prompt_embeds.dtype
|
||||
|
||||
prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
|
||||
|
||||
bs_embed, seq_len, _ = prompt_embeds.shape
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
|
||||
|
||||
# get unconditional embeddings for classifier free guidance
|
||||
if do_classifier_free_guidance and negative_prompt_embeds is None:
|
||||
uncond_tokens: List[str]
|
||||
if negative_prompt is None:
|
||||
uncond_tokens = [""] * batch_size
|
||||
elif prompt is not None and type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif isinstance(negative_prompt, str):
|
||||
uncond_tokens = [negative_prompt]
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
else:
|
||||
uncond_tokens = negative_prompt
|
||||
|
||||
# textual inversion: procecss multi-vector tokens if necessary
|
||||
if isinstance(self, TextualInversionLoaderMixin):
|
||||
uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
|
||||
|
||||
max_length = prompt_embeds.shape[1]
|
||||
uncond_input = self.tokenizer(
|
||||
uncond_tokens,
|
||||
padding="max_length",
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
|
||||
attention_mask = uncond_input.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
|
||||
negative_prompt_embeds = self.text_encoder(
|
||||
uncond_input.input_ids.to(device),
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
negative_prompt_embeds = negative_prompt_embeds[0]
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
||||
seq_len = negative_prompt_embeds.shape[1]
|
||||
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
|
||||
|
||||
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
||||
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
|
||||
|
||||
return prompt_embeds, negative_prompt_embeds
|
||||
|
||||
def run_safety_checker(self, image, device, dtype):
|
||||
if self.safety_checker is None:
|
||||
has_nsfw_concept = None
|
||||
else:
|
||||
if torch.is_tensor(image):
|
||||
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
|
||||
else:
|
||||
feature_extractor_input = self.image_processor.numpy_to_pil(image)
|
||||
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
|
||||
image, has_nsfw_concept = self.safety_checker(
|
||||
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
|
||||
)
|
||||
return image, has_nsfw_concept
|
||||
|
||||
def decode_latents(self, latents):
|
||||
deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
|
||||
deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
|
||||
|
||||
latents = 1 / self.vae.config.scaling_factor * latents
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return image
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
height,
|
||||
width,
|
||||
callback_steps,
|
||||
negative_prompt=None,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
):
|
||||
if height % 8 != 0 or width % 8 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
||||
|
||||
if (callback_steps is None) or (
|
||||
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
|
||||
):
|
||||
raise ValueError(
|
||||
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
|
||||
f" {type(callback_steps)}."
|
||||
)
|
||||
|
||||
if prompt is not None and prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
||||
" only forward one of the two."
|
||||
)
|
||||
elif prompt is None and prompt_embeds is None:
|
||||
raise ValueError(
|
||||
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
||||
)
|
||||
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
|
||||
if negative_prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
|
||||
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
||||
)
|
||||
|
||||
if prompt_embeds is not None and negative_prompt_embeds is not None:
|
||||
if prompt_embeds.shape != negative_prompt_embeds.shape:
|
||||
raise ValueError(
|
||||
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
|
||||
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
|
||||
f" {negative_prompt_embeds.shape}."
|
||||
)
|
||||
|
||||
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
|
||||
shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
@torch.no_grad()
|
||||
@replace_example_docstring(EXAMPLE_DOC_STRING)
|
||||
def __call__(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: int = 1,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
guidance_rescale: float = 0.0,
|
||||
cache_interval: int = 1,
|
||||
cache_layer_id: int = None,
|
||||
cache_block_id: int = None,
|
||||
uniform: bool = True,
|
||||
pow: float = None,
|
||||
center: int = None,
|
||||
output_all_sequence: bool = False,
|
||||
):
|
||||
r"""
|
||||
The call function to the pipeline for generation.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
|
||||
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The width in pixels of the generated image.
|
||||
num_inference_steps (`int`, *optional*, defaults to 50):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference.
|
||||
guidance_scale (`float`, *optional*, defaults to 7.5):
|
||||
A higher guidance scale value encourages the model to generate images closely linked to the text
|
||||
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
|
||||
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
|
||||
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
eta (`float`, *optional*, defaults to 0.0):
|
||||
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
|
||||
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
||||
generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor is generated by sampling using the supplied random `generator`.
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
|
||||
provided, text embeddings are generated from the `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
|
||||
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
||||
plain tuple.
|
||||
callback (`Callable`, *optional*):
|
||||
A function that calls every `callback_steps` steps during inference. The function is called with the
|
||||
following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
||||
callback_steps (`int`, *optional*, defaults to 1):
|
||||
The frequency at which the `callback` function is called. If not specified, the callback is called at
|
||||
every step.
|
||||
cross_attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
|
||||
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
guidance_rescale (`float`, *optional*, defaults to 0.7):
|
||||
Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
|
||||
Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
|
||||
using zero terminal SNR.
|
||||
|
||||
Examples:
|
||||
|
||||
Returns:
|
||||
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
|
||||
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
|
||||
otherwise a `tuple` is returned where the first element is a list with the generated images and the
|
||||
second element is a list of `bool`s indicating whether the corresponding generated image contains
|
||||
"not-safe-for-work" (nsfw) content.
|
||||
"""
|
||||
# 0. Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(
|
||||
prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
|
||||
)
|
||||
|
||||
# 2. Define call parameters
|
||||
if prompt is not None and isinstance(prompt, str):
|
||||
batch_size = 1
|
||||
elif prompt is not None and isinstance(prompt, list):
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
device = self._execution_device
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# 3. Encode input prompt
|
||||
text_encoder_lora_scale = (
|
||||
cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
|
||||
)
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt,
|
||||
device,
|
||||
num_images_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
negative_prompt,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
lora_scale=text_encoder_lora_scale,
|
||||
)
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
if do_classifier_free_guidance:
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
|
||||
|
||||
# 4. Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
# 5. Prepare latent variables
|
||||
num_channels_latents = self.unet.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
# 7. Denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
|
||||
prv_features = None
|
||||
latents_list = [latents]
|
||||
|
||||
if cache_interval == 1:
|
||||
interval_seq = list(range(num_inference_steps))
|
||||
else:
|
||||
if uniform:
|
||||
interval_seq = list(range(0, num_inference_steps, cache_interval))
|
||||
else:
|
||||
num_slow_step = num_inference_steps//cache_interval
|
||||
if num_inference_steps%cache_interval != 0:
|
||||
num_slow_step += 1
|
||||
|
||||
interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
|
||||
#interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
|
||||
|
||||
interval_seq = sorted(interval_seq)
|
||||
#print(interval_seq, len(interval_seq), pow)
|
||||
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
#print("[INFO] Update Feature Interval = {}, Update Layer Number = {}, Update Block Number = {}".format(cache_interval, cache_layer_id, cache_block_id))
|
||||
for i, t in enumerate(timesteps):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
if i in interval_seq:
|
||||
prv_features = None
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred, prv_features = self.unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
replicate_prv_feature=prv_features,
|
||||
quick_replicate= cache_interval>1,
|
||||
cache_layer_id=cache_layer_id,
|
||||
cache_block_id=cache_block_id,
|
||||
return_dict=False,
|
||||
)
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
|
||||
if do_classifier_free_guidance and guidance_rescale > 0.0:
|
||||
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
|
||||
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
|
||||
latents_list.append(latents)
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
callback(i, t, latents)
|
||||
|
||||
if not output_type == "latent":
|
||||
if output_all_sequence:
|
||||
image = [self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] for latents in latents_list]
|
||||
has_nsfw_concept = None #self.run_safety_checker(images[0], device, prompt_embeds.dtype)
|
||||
num_img = len(image)
|
||||
else:
|
||||
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
|
||||
has_nsfw_concept = None
|
||||
num_img = image.shape[0]
|
||||
else:
|
||||
image = latents
|
||||
has_nsfw_concept = None
|
||||
|
||||
if has_nsfw_concept is None:
|
||||
do_denormalize = [True] * num_img
|
||||
else:
|
||||
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
|
||||
|
||||
if output_all_sequence:
|
||||
image = [self.image_processor.postprocess(img, output_type=output_type, do_denormalize=do_denormalize) for img in image]
|
||||
else:
|
||||
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
if not return_dict:
|
||||
return (image, has_nsfw_concept,)
|
||||
|
||||
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
||||
741
ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py
Normal file
741
ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py
Normal file
@@ -0,0 +1,741 @@
|
||||
import copy
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.functional import grid_sample
|
||||
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from diffusers.models import AutoencoderKL
|
||||
from .unet_2d_condition import UNet2DConditionModel
|
||||
from .pipeline_stable_diffusion import StableDiffusionPipeline, StableDiffusionSafetyChecker
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100):
|
||||
samples = []
|
||||
|
||||
while len(samples) < sample_size:
|
||||
# Sample from a Gaussian centered at n/2
|
||||
sample = int(np.random.normal(loc=n/2, scale=std_dev))
|
||||
|
||||
# Check if the sample is in bounds
|
||||
if 1 <= sample < n and sample not in samples:
|
||||
samples.append(sample)
|
||||
|
||||
return samples
|
||||
|
||||
def sample_from_quad(total_numbers, n_samples, pow=1.2):
|
||||
while pow > 1:
|
||||
# Generate linearly spaced values between 0 and a max value
|
||||
x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1)
|
||||
|
||||
# Raise these values to the power of 1.5 to get a non-linear distribution
|
||||
indices = np.unique(np.int32(x_values**pow))[:-1]
|
||||
if len(indices) == n_samples:
|
||||
break
|
||||
pow -=0.02
|
||||
if pow <= 1:
|
||||
raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
|
||||
return indices, pow
|
||||
|
||||
def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2):
|
||||
while pow > 1:
|
||||
# Generate linearly spaced values between 0 and a max value
|
||||
x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1)
|
||||
indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]]
|
||||
if len(indices) == n_samples:
|
||||
break
|
||||
pow -=0.02
|
||||
if pow <= 1:
|
||||
raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
|
||||
return indices, pow
|
||||
|
||||
def rearrange_0(tensor, f):
|
||||
F, C, H, W = tensor.size()
|
||||
tensor = torch.permute(torch.reshape(tensor, (F // f, f, C, H, W)), (0, 2, 1, 3, 4))
|
||||
return tensor
|
||||
|
||||
|
||||
def rearrange_1(tensor):
|
||||
B, C, F, H, W = tensor.size()
|
||||
return torch.reshape(torch.permute(tensor, (0, 2, 1, 3, 4)), (B * F, C, H, W))
|
||||
|
||||
|
||||
def rearrange_3(tensor, f):
|
||||
F, D, C = tensor.size()
|
||||
return torch.reshape(tensor, (F // f, f, D, C))
|
||||
|
||||
|
||||
def rearrange_4(tensor):
|
||||
B, F, D, C = tensor.size()
|
||||
return torch.reshape(tensor, (B * F, D, C))
|
||||
|
||||
|
||||
class CrossFrameAttnProcessor:
|
||||
"""
|
||||
Cross frame attention processor. Each frame attends the first frame.
|
||||
|
||||
Args:
|
||||
batch_size: The number that represents actual batch size, other than the frames.
|
||||
For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to
|
||||
2, due to classifier-free guidance.
|
||||
"""
|
||||
|
||||
def __init__(self, batch_size=2):
|
||||
self.batch_size = batch_size
|
||||
|
||||
def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
|
||||
batch_size, sequence_length, _ = hidden_states.shape
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
is_cross_attention = encoder_hidden_states is not None
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
# Cross Frame Attention
|
||||
if not is_cross_attention:
|
||||
video_length = key.size()[0] // self.batch_size
|
||||
first_frame_index = [0] * video_length
|
||||
|
||||
# rearrange keys to have batch and frames in the 1st and 2nd dims respectively
|
||||
key = rearrange_3(key, video_length)
|
||||
key = key[:, first_frame_index]
|
||||
# rearrange values to have batch and frames in the 1st and 2nd dims respectively
|
||||
value = rearrange_3(value, video_length)
|
||||
value = value[:, first_frame_index]
|
||||
|
||||
# rearrange back to original shape
|
||||
key = rearrange_4(key)
|
||||
value = rearrange_4(value)
|
||||
|
||||
query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CrossFrameAttnProcessor2_0:
|
||||
"""
|
||||
Cross frame attention processor with scaled_dot_product attention of Pytorch 2.0.
|
||||
|
||||
Args:
|
||||
batch_size: The number that represents actual batch size, other than the frames.
|
||||
For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to
|
||||
2, due to classifier-free guidance.
|
||||
"""
|
||||
|
||||
def __init__(self, batch_size=2):
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
self.batch_size = batch_size
|
||||
|
||||
def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
inner_dim = hidden_states.shape[-1]
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
is_cross_attention = encoder_hidden_states is not None
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
# Cross Frame Attention
|
||||
if not is_cross_attention:
|
||||
video_length = max(1, key.size()[0] // self.batch_size)
|
||||
first_frame_index = [0] * video_length
|
||||
|
||||
# rearrange keys to have batch and frames in the 1st and 2nd dims respectively
|
||||
key = rearrange_3(key, video_length)
|
||||
key = key[:, first_frame_index]
|
||||
# rearrange values to have batch and frames in the 1st and 2nd dims respectively
|
||||
value = rearrange_3(value, video_length)
|
||||
value = value[:, first_frame_index]
|
||||
|
||||
# rearrange back to original shape
|
||||
key = rearrange_4(key)
|
||||
value = rearrange_4(value)
|
||||
|
||||
head_dim = inner_dim // attn.heads
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextToVideoPipelineOutput(BaseOutput):
|
||||
r"""
|
||||
Output class for zero-shot text-to-video pipeline.
|
||||
|
||||
Args:
|
||||
images (`[List[PIL.Image.Image]`, `np.ndarray`]):
|
||||
List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
|
||||
num_channels)`.
|
||||
nsfw_content_detected (`[List[bool]]`):
|
||||
List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
|
||||
`None` if safety checking could not be performed.
|
||||
"""
|
||||
|
||||
images: Union[List[PIL.Image.Image], np.ndarray]
|
||||
nsfw_content_detected: Optional[List[bool]]
|
||||
|
||||
|
||||
def coords_grid(batch, ht, wd, device):
|
||||
# Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py
|
||||
coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
|
||||
def warp_single_latent(latent, reference_flow):
|
||||
"""
|
||||
Warp latent of a single frame with given flow
|
||||
|
||||
Args:
|
||||
latent: latent code of a single frame
|
||||
reference_flow: flow which to warp the latent with
|
||||
|
||||
Returns:
|
||||
warped: warped latent
|
||||
"""
|
||||
_, _, H, W = reference_flow.size()
|
||||
_, _, h, w = latent.size()
|
||||
coords0 = coords_grid(1, H, W, device=latent.device).to(latent.dtype)
|
||||
|
||||
coords_t0 = coords0 + reference_flow
|
||||
coords_t0[:, 0] /= W
|
||||
coords_t0[:, 1] /= H
|
||||
|
||||
coords_t0 = coords_t0 * 2.0 - 1.0
|
||||
coords_t0 = F.interpolate(coords_t0, size=(h, w), mode="bilinear")
|
||||
coords_t0 = torch.permute(coords_t0, (0, 2, 3, 1))
|
||||
|
||||
warped = grid_sample(latent, coords_t0, mode="nearest", padding_mode="reflection")
|
||||
return warped
|
||||
|
||||
|
||||
def create_motion_field(motion_field_strength_x, motion_field_strength_y, frame_ids, device, dtype):
|
||||
"""
|
||||
Create translation motion field
|
||||
|
||||
Args:
|
||||
motion_field_strength_x: motion strength along x-axis
|
||||
motion_field_strength_y: motion strength along y-axis
|
||||
frame_ids: indexes of the frames the latents of which are being processed.
|
||||
This is needed when we perform chunk-by-chunk inference
|
||||
device: device
|
||||
dtype: dtype
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
seq_length = len(frame_ids)
|
||||
reference_flow = torch.zeros((seq_length, 2, 512, 512), device=device, dtype=dtype)
|
||||
for fr_idx in range(seq_length):
|
||||
reference_flow[fr_idx, 0, :, :] = motion_field_strength_x * (frame_ids[fr_idx])
|
||||
reference_flow[fr_idx, 1, :, :] = motion_field_strength_y * (frame_ids[fr_idx])
|
||||
return reference_flow
|
||||
|
||||
|
||||
def create_motion_field_and_warp_latents(motion_field_strength_x, motion_field_strength_y, frame_ids, latents):
|
||||
"""
|
||||
Creates translation motion and warps the latents accordingly
|
||||
|
||||
Args:
|
||||
motion_field_strength_x: motion strength along x-axis
|
||||
motion_field_strength_y: motion strength along y-axis
|
||||
frame_ids: indexes of the frames the latents of which are being processed.
|
||||
This is needed when we perform chunk-by-chunk inference
|
||||
latents: latent codes of frames
|
||||
|
||||
Returns:
|
||||
warped_latents: warped latents
|
||||
"""
|
||||
motion_field = create_motion_field(
|
||||
motion_field_strength_x=motion_field_strength_x,
|
||||
motion_field_strength_y=motion_field_strength_y,
|
||||
frame_ids=frame_ids,
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
warped_latents = latents.clone().detach()
|
||||
for i in range(len(warped_latents)):
|
||||
warped_latents[i] = warp_single_latent(latents[i][None], motion_field[i][None])
|
||||
return warped_latents
|
||||
|
||||
|
||||
class TextToVideoZeroPipeline(StableDiffusionPipeline):
|
||||
r"""
|
||||
Pipeline for zero-shot text-to-video generation using Stable Diffusion.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
|
||||
implemented for all pipelines (downloading, saving, running on a particular device, etc.).
|
||||
|
||||
Args:
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
|
||||
text_encoder ([`CLIPTextModel`]):
|
||||
Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
|
||||
tokenizer (`CLIPTokenizer`):
|
||||
A [`~transformers.CLIPTokenizer`] to tokenize text.
|
||||
unet ([`UNet2DConditionModel`]):
|
||||
A [`UNet3DConditionModel`] to denoise the encoded video latents.
|
||||
scheduler ([`SchedulerMixin`]):
|
||||
A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
|
||||
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
|
||||
safety_checker ([`StableDiffusionSafetyChecker`]):
|
||||
Classification module that estimates whether generated images could be considered offensive or harmful.
|
||||
Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
|
||||
about a model's potential harms.
|
||||
feature_extractor ([`CLIPImageProcessor`]):
|
||||
A [`CLIPImageProcessor`] to extract features from generated images; used as inputs to the `safety_checker`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder: CLIPTextModel,
|
||||
tokenizer: CLIPTokenizer,
|
||||
unet: UNet2DConditionModel,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
safety_checker: StableDiffusionSafetyChecker,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
requires_safety_checker: bool = True,
|
||||
):
|
||||
super().__init__(
|
||||
vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker
|
||||
)
|
||||
processor = (
|
||||
CrossFrameAttnProcessor2_0(batch_size=2)
|
||||
if hasattr(F, "scaled_dot_product_attention")
|
||||
else CrossFrameAttnProcessor(batch_size=2)
|
||||
)
|
||||
self.unet.set_attn_processor(processor)
|
||||
|
||||
def forward_loop(self, x_t0, t0, t1, generator):
|
||||
"""
|
||||
Perform DDPM forward process from time t0 to t1. This is the same as adding noise with corresponding variance.
|
||||
|
||||
Args:
|
||||
x_t0:
|
||||
Latent code at time t0.
|
||||
t0:
|
||||
Timestep at t0.
|
||||
t1:
|
||||
Timestamp at t1.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
||||
generation deterministic.
|
||||
|
||||
Returns:
|
||||
x_t1:
|
||||
Forward process applied to x_t0 from time t0 to t1.
|
||||
"""
|
||||
eps = randn_tensor(x_t0.size(), generator=generator, dtype=x_t0.dtype, device=x_t0.device)
|
||||
alpha_vec = torch.prod(self.scheduler.alphas[t0:t1])
|
||||
x_t1 = torch.sqrt(alpha_vec) * x_t0 + torch.sqrt(1 - alpha_vec) * eps
|
||||
return x_t1
|
||||
|
||||
def backward_loop(
|
||||
self,
|
||||
latents,
|
||||
timesteps,
|
||||
prompt_embeds,
|
||||
guidance_scale,
|
||||
callback,
|
||||
callback_steps,
|
||||
num_warmup_steps,
|
||||
extra_step_kwargs,
|
||||
prv_features,
|
||||
interval_seq,
|
||||
cache_interval,
|
||||
cache_block_id,
|
||||
cache_layer_id,
|
||||
cross_attention_kwargs=None,
|
||||
):
|
||||
"""
|
||||
Perform backward process given list of time steps.
|
||||
|
||||
Args:
|
||||
latents:
|
||||
Latents at time timesteps[0].
|
||||
timesteps:
|
||||
Time steps along which to perform backward process.
|
||||
prompt_embeds:
|
||||
Pre-generated text embeddings.
|
||||
guidance_scale:
|
||||
A higher guidance scale value encourages the model to generate images closely linked to the text
|
||||
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
|
||||
callback (`Callable`, *optional*):
|
||||
A function that calls every `callback_steps` steps during inference. The function is called with the
|
||||
following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
||||
callback_steps (`int`, *optional*, defaults to 1):
|
||||
The frequency at which the `callback` function is called. If not specified, the callback is called at
|
||||
every step.
|
||||
extra_step_kwargs:
|
||||
Extra_step_kwargs.
|
||||
cross_attention_kwargs:
|
||||
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
|
||||
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
num_warmup_steps:
|
||||
number of warmup steps.
|
||||
|
||||
Returns:
|
||||
latents:
|
||||
Latents of backward process output at time timesteps[-1].
|
||||
"""
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
num_steps = (len(timesteps) - num_warmup_steps) // self.scheduler.order
|
||||
with self.progress_bar(total=num_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
########
|
||||
if i in interval_seq:
|
||||
prv_features = None
|
||||
# predict the noise residual
|
||||
noise_pred, prv_features = self.unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
replicate_prv_feature=prv_features,
|
||||
quick_replicate= cache_interval>1,
|
||||
cache_layer_id=cache_layer_id,
|
||||
cache_block_id=cache_block_id,
|
||||
return_dict=False,
|
||||
)
|
||||
########
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
step_idx = i // getattr(self.scheduler, "order", 1)
|
||||
callback(step_idx, t, latents)
|
||||
return latents.clone().detach()
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: Union[str, List[str]],
|
||||
video_length: Optional[int] = 8,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
num_videos_per_prompt: Optional[int] = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
motion_field_strength_x: float = 12,
|
||||
motion_field_strength_y: float = 12,
|
||||
output_type: Optional[str] = "tensor",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: Optional[int] = 1,
|
||||
t0: int = 44,
|
||||
t1: int = 47,
|
||||
frame_ids: Optional[List[int]] = None,
|
||||
########
|
||||
cache_interval: int = 1,
|
||||
cache_layer_id: int = None,
|
||||
cache_block_id: int = None,
|
||||
uniform: bool = True,
|
||||
pow: float = None,
|
||||
center: int = None,
|
||||
output_all_sequence: bool = False,
|
||||
########
|
||||
):
|
||||
"""
|
||||
The call function to the pipeline for generation.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
|
||||
video_length (`int`, *optional*, defaults to 8):
|
||||
The number of generated video frames.
|
||||
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The width in pixels of the generated image.
|
||||
num_inference_steps (`int`, *optional*, defaults to 50):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference.
|
||||
guidance_scale (`float`, *optional*, defaults to 7.5):
|
||||
A higher guidance scale value encourages the model to generate images closely linked to the text
|
||||
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide what to not include in video generation. If not defined, you need to
|
||||
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
|
||||
num_videos_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of videos to generate per prompt.
|
||||
eta (`float`, *optional*, defaults to 0.0):
|
||||
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
|
||||
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
||||
generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for video
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor is generated by sampling using the supplied random `generator`.
|
||||
output_type (`str`, *optional*, defaults to `"numpy"`):
|
||||
The output format of the generated video. Choose between `"latent"` and `"numpy"`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a
|
||||
[`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`] instead of
|
||||
a plain tuple.
|
||||
callback (`Callable`, *optional*):
|
||||
A function that calls every `callback_steps` steps during inference. The function is called with the
|
||||
following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
||||
callback_steps (`int`, *optional*, defaults to 1):
|
||||
The frequency at which the `callback` function is called. If not specified, the callback is called at
|
||||
every step.
|
||||
motion_field_strength_x (`float`, *optional*, defaults to 12):
|
||||
Strength of motion in generated video along x-axis. See the [paper](https://arxiv.org/abs/2303.13439),
|
||||
Sect. 3.3.1.
|
||||
motion_field_strength_y (`float`, *optional*, defaults to 12):
|
||||
Strength of motion in generated video along y-axis. See the [paper](https://arxiv.org/abs/2303.13439),
|
||||
Sect. 3.3.1.
|
||||
t0 (`int`, *optional*, defaults to 44):
|
||||
Timestep t0. Should be in the range [0, num_inference_steps - 1]. See the
|
||||
[paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1.
|
||||
t1 (`int`, *optional*, defaults to 47):
|
||||
Timestep t0. Should be in the range [t0 + 1, num_inference_steps - 1]. See the
|
||||
[paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1.
|
||||
frame_ids (`List[int]`, *optional*):
|
||||
Indexes of the frames that are being generated. This is used when generating longer videos
|
||||
chunk-by-chunk.
|
||||
|
||||
Returns:
|
||||
[`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`]:
|
||||
The output contains a `ndarray` of the generated video, when `output_type` != `"latent"`, otherwise a
|
||||
latent code of generated videos and a list of `bool`s indicating whether the corresponding generated
|
||||
video contains "not-safe-for-work" (nsfw) content..
|
||||
"""
|
||||
assert video_length > 0
|
||||
if frame_ids is None:
|
||||
frame_ids = list(range(video_length))
|
||||
assert len(frame_ids) == video_length
|
||||
|
||||
assert num_videos_per_prompt == 1
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
if isinstance(negative_prompt, str):
|
||||
negative_prompt = [negative_prompt]
|
||||
|
||||
# Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
# Check inputs. Raise error if not correct
|
||||
self.check_inputs(prompt, height, width, callback_steps)
|
||||
|
||||
# Define call parameters
|
||||
batch_size = 1 if isinstance(prompt, str) else len(prompt)
|
||||
device = self._execution_device
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# Encode input prompt
|
||||
prompt_embeds = self._encode_prompt(
|
||||
prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
|
||||
)
|
||||
|
||||
# Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
# Prepare latent variables
|
||||
num_channels_latents = self.unet.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_videos_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
# Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
|
||||
prv_features = None #record cache feature ****
|
||||
latents_list = [latents]
|
||||
|
||||
if cache_interval == 1:
|
||||
interval_seq = list(range(num_inference_steps))
|
||||
else:
|
||||
if uniform:
|
||||
interval_seq = list(range(0, num_inference_steps, cache_interval))
|
||||
else:
|
||||
num_slow_step = num_inference_steps//cache_interval
|
||||
if num_inference_steps%cache_interval != 0:
|
||||
num_slow_step += 1
|
||||
|
||||
interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
|
||||
#interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
|
||||
|
||||
interval_seq = sorted(interval_seq)
|
||||
|
||||
# Perform the first backward process up to time T_1
|
||||
x_1_t1 = self.backward_loop(
|
||||
timesteps=timesteps[: -t1 - 1],
|
||||
prompt_embeds=prompt_embeds,
|
||||
latents=latents,
|
||||
guidance_scale=guidance_scale,
|
||||
callback=callback,
|
||||
callback_steps=callback_steps,
|
||||
extra_step_kwargs=extra_step_kwargs,
|
||||
num_warmup_steps=num_warmup_steps,
|
||||
prv_features=prv_features,
|
||||
interval_seq=interval_seq,
|
||||
cache_interval=cache_interval,
|
||||
cache_block_id=cache_block_id,
|
||||
cache_layer_id=cache_layer_id,
|
||||
)
|
||||
scheduler_copy = copy.deepcopy(self.scheduler)
|
||||
|
||||
# Perform the second backward process up to time T_0
|
||||
x_1_t0 = self.backward_loop(
|
||||
timesteps=timesteps[-t1 - 1 : -t0 - 1],
|
||||
prompt_embeds=prompt_embeds,
|
||||
latents=x_1_t1,
|
||||
guidance_scale=guidance_scale,
|
||||
callback=callback,
|
||||
callback_steps=callback_steps,
|
||||
extra_step_kwargs=extra_step_kwargs,
|
||||
num_warmup_steps=0,
|
||||
prv_features=prv_features,
|
||||
interval_seq=interval_seq,
|
||||
cache_interval=cache_interval,
|
||||
cache_block_id=cache_block_id,
|
||||
cache_layer_id=cache_layer_id,
|
||||
)
|
||||
|
||||
# Propagate first frame latents at time T_0 to remaining frames
|
||||
x_2k_t0 = x_1_t0.repeat(video_length - 1, 1, 1, 1)
|
||||
|
||||
# Add motion in latents at time T_0
|
||||
x_2k_t0 = create_motion_field_and_warp_latents(
|
||||
motion_field_strength_x=motion_field_strength_x,
|
||||
motion_field_strength_y=motion_field_strength_y,
|
||||
latents=x_2k_t0,
|
||||
frame_ids=frame_ids[1:],
|
||||
)
|
||||
|
||||
# Perform forward process up to time T_1
|
||||
x_2k_t1 = self.forward_loop(
|
||||
x_t0=x_2k_t0,
|
||||
t0=timesteps[-t0 - 1].item(),
|
||||
t1=timesteps[-t1 - 1].item(),
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
# Perform backward process from time T_1 to 0
|
||||
x_1k_t1 = torch.cat([x_1_t1, x_2k_t1])
|
||||
b, l, d = prompt_embeds.size()
|
||||
prompt_embeds = prompt_embeds[:, None].repeat(1, video_length, 1, 1).reshape(b * video_length, l, d)
|
||||
|
||||
self.scheduler = scheduler_copy
|
||||
x_1k_0 = self.backward_loop(
|
||||
timesteps=timesteps[-t1 - 1 :],
|
||||
prompt_embeds=prompt_embeds,
|
||||
latents=x_1k_t1,
|
||||
guidance_scale=guidance_scale,
|
||||
callback=callback,
|
||||
callback_steps=callback_steps,
|
||||
extra_step_kwargs=extra_step_kwargs,
|
||||
num_warmup_steps=0,
|
||||
prv_features=prv_features,
|
||||
interval_seq=interval_seq,
|
||||
cache_interval=cache_interval,
|
||||
cache_block_id=cache_block_id,
|
||||
cache_layer_id=cache_layer_id,
|
||||
)
|
||||
latents = x_1k_0
|
||||
|
||||
# manually for max memory savings
|
||||
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
|
||||
self.unet.to("cpu")
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if output_type == "latent":
|
||||
image = latents
|
||||
has_nsfw_concept = None
|
||||
else:
|
||||
image = self.decode_latents(latents)
|
||||
# Run safety checker
|
||||
image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if not return_dict:
|
||||
return (image, has_nsfw_concept)
|
||||
|
||||
return TextToVideoPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
||||
1839
ixformer_sdk/contrib/DeepCache/sd/pipeline_utils.py
Normal file
1839
ixformer_sdk/contrib/DeepCache/sd/pipeline_utils.py
Normal file
File diff suppressed because it is too large
Load Diff
3296
ixformer_sdk/contrib/DeepCache/sd/unet_2d_blocks.py
Normal file
3296
ixformer_sdk/contrib/DeepCache/sd/unet_2d_blocks.py
Normal file
File diff suppressed because it is too large
Load Diff
1257
ixformer_sdk/contrib/DeepCache/sd/unet_2d_condition.py
Normal file
1257
ixformer_sdk/contrib/DeepCache/sd/unet_2d_condition.py
Normal file
File diff suppressed because it is too large
Load Diff
0
ixformer_sdk/contrib/DeepCache/sdxl/__init__.py
Normal file
0
ixformer_sdk/contrib/DeepCache/sdxl/__init__.py
Normal file
1100
ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl.py
Normal file
1100
ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl.py
Normal file
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
1839
ixformer_sdk/contrib/DeepCache/sdxl/pipeline_utils.py
Normal file
1839
ixformer_sdk/contrib/DeepCache/sdxl/pipeline_utils.py
Normal file
File diff suppressed because it is too large
Load Diff
3339
ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_blocks.py
Normal file
3339
ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_blocks.py
Normal file
File diff suppressed because it is too large
Load Diff
1259
ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_condition.py
Normal file
1259
ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_condition.py
Normal file
File diff suppressed because it is too large
Load Diff
0
ixformer_sdk/contrib/DeepCache/svd/__init__.py
Normal file
0
ixformer_sdk/contrib/DeepCache/svd/__init__.py
Normal file
@@ -0,0 +1,659 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# 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.
|
||||
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
|
||||
from diffusers.schedulers import EulerDiscreteScheduler
|
||||
from diffusers.utils import BaseOutput, logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from .pipeline_utils import DiffusionPipeline
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def _append_dims(x, target_dims):
|
||||
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
||||
dims_to_append = target_dims - x.ndim
|
||||
if dims_to_append < 0:
|
||||
raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
|
||||
return x[(...,) + (None,) * dims_to_append]
|
||||
|
||||
|
||||
def tensor2vid(video: torch.Tensor, processor, output_type="np"):
|
||||
# Based on:
|
||||
# https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78
|
||||
|
||||
batch_size, channels, num_frames, height, width = video.shape
|
||||
outputs = []
|
||||
for batch_idx in range(batch_size):
|
||||
batch_vid = video[batch_idx].permute(1, 0, 2, 3)
|
||||
batch_output = processor.postprocess(batch_vid, output_type)
|
||||
|
||||
outputs.append(batch_output)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
@dataclass
|
||||
class StableVideoDiffusionPipelineOutput(BaseOutput):
|
||||
r"""
|
||||
Output class for zero-shot text-to-video pipeline.
|
||||
|
||||
Args:
|
||||
frames (`[List[PIL.Image.Image]`, `np.ndarray`]):
|
||||
List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
|
||||
num_channels)`.
|
||||
"""
|
||||
|
||||
frames: Union[List[PIL.Image.Image], np.ndarray]
|
||||
|
||||
|
||||
class StableVideoDiffusionPipeline(DiffusionPipeline):
|
||||
r"""
|
||||
Pipeline to generate video from an input image using Stable Video Diffusion.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
|
||||
implemented for all pipelines (downloading, saving, running on a particular device, etc.).
|
||||
|
||||
Args:
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
|
||||
image_encoder ([`~transformers.CLIPVisionModelWithProjection`]):
|
||||
Frozen CLIP image-encoder ([laion/CLIP-ViT-H-14-laion2B-s32B-b79K](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)).
|
||||
unet ([`UNetSpatioTemporalConditionModel`]):cache_interval=5, cache_branch=0,
|
||||
A `UNetSpatioTemporalConditionModel` to denoise the encoded image latents.
|
||||
scheduler ([`EulerDiscreteScheduler`]):
|
||||
A scheduler to be used in combination with `unet` to denoise the encoded image latents.
|
||||
feature_extractor ([`~transformers.CLIPImageProcessor`]):
|
||||
A `CLIPImageProcessor` to extract features from generated images.
|
||||
"""
|
||||
|
||||
model_cpu_offload_seq = "image_encoder->unet->vae"
|
||||
_callback_tensor_inputs = ["latents"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKLTemporalDecoder,
|
||||
image_encoder: CLIPVisionModelWithProjection,
|
||||
unet: UNetSpatioTemporalConditionModel,
|
||||
scheduler: EulerDiscreteScheduler,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
image_encoder=image_encoder,
|
||||
unet=unet,
|
||||
scheduler=scheduler,
|
||||
feature_extractor=feature_extractor,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
|
||||
def _encode_image(self, image, device, num_videos_per_prompt, do_classifier_free_guidance):
|
||||
dtype = next(self.image_encoder.parameters()).dtype
|
||||
|
||||
if not isinstance(image, torch.Tensor):
|
||||
image = self.image_processor.pil_to_numpy(image)
|
||||
image = self.image_processor.numpy_to_pt(image)
|
||||
|
||||
# We normalize the image before resizing to match with the original implementation.
|
||||
# Then we unnormalize it after resizing.
|
||||
image = image * 2.0 - 1.0
|
||||
image = _resize_with_antialiasing(image, (224, 224))
|
||||
image = (image + 1.0) / 2.0
|
||||
|
||||
# Normalize the image with for CLIP input
|
||||
image = self.feature_extractor(
|
||||
images=image,
|
||||
do_normalize=True,
|
||||
do_center_crop=False,
|
||||
do_resize=False,
|
||||
do_rescale=False,
|
||||
return_tensors="pt",
|
||||
).pixel_values
|
||||
|
||||
image = image.to(device=device, dtype=dtype)
|
||||
image_embeddings = self.image_encoder(image).image_embeds
|
||||
image_embeddings = image_embeddings.unsqueeze(1)
|
||||
|
||||
# duplicate image embeddings for each generation per prompt, using mps friendly method
|
||||
bs_embed, seq_len, _ = image_embeddings.shape
|
||||
image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1)
|
||||
image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
negative_image_embeddings = torch.zeros_like(image_embeddings)
|
||||
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
image_embeddings = torch.cat([negative_image_embeddings, image_embeddings])
|
||||
|
||||
return image_embeddings
|
||||
|
||||
def _encode_vae_image(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
device,
|
||||
num_videos_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
):
|
||||
image = image.to(device=device)
|
||||
image_latents = self.vae.encode(image).latent_dist.mode()
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
negative_image_latents = torch.zeros_like(image_latents)
|
||||
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
image_latents = torch.cat([negative_image_latents, image_latents])
|
||||
|
||||
# duplicate image_latents for each generation per prompt, using mps friendly method
|
||||
image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1)
|
||||
|
||||
return image_latents
|
||||
|
||||
def _get_add_time_ids(
|
||||
self,
|
||||
fps,
|
||||
motion_bucket_id,
|
||||
noise_aug_strength,
|
||||
dtype,
|
||||
batch_size,
|
||||
num_videos_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
):
|
||||
add_time_ids = [fps, motion_bucket_id, noise_aug_strength]
|
||||
|
||||
passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids)
|
||||
expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
|
||||
|
||||
if expected_add_embed_dim != passed_add_embed_dim:
|
||||
raise ValueError(
|
||||
f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
|
||||
)
|
||||
|
||||
add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
|
||||
add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
add_time_ids = torch.cat([add_time_ids, add_time_ids])
|
||||
|
||||
return add_time_ids
|
||||
|
||||
def decode_latents(self, latents, num_frames, decode_chunk_size=14):
|
||||
# [batch, frames, channels, height, width] -> [batch*frames, channels, height, width]
|
||||
latents = latents.flatten(0, 1)
|
||||
|
||||
latents = 1 / self.vae.config.scaling_factor * latents
|
||||
|
||||
accepts_num_frames = "num_frames" in set(inspect.signature(self.vae.forward).parameters.keys())
|
||||
|
||||
# decode decode_chunk_size frames at a time to avoid OOM
|
||||
frames = []
|
||||
for i in range(0, latents.shape[0], decode_chunk_size):
|
||||
num_frames_in = latents[i : i + decode_chunk_size].shape[0]
|
||||
decode_kwargs = {}
|
||||
if accepts_num_frames:
|
||||
# we only pass num_frames_in if it's expected
|
||||
decode_kwargs["num_frames"] = num_frames_in
|
||||
|
||||
frame = self.vae.decode(latents[i : i + decode_chunk_size], **decode_kwargs).sample
|
||||
frames.append(frame)
|
||||
frames = torch.cat(frames, dim=0)
|
||||
|
||||
# [batch*frames, channels, height, width] -> [batch, channels, frames, height, width]
|
||||
frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
|
||||
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
frames = frames.float()
|
||||
return frames
|
||||
|
||||
def check_inputs(self, image, height, width):
|
||||
if (
|
||||
not isinstance(image, torch.Tensor)
|
||||
and not isinstance(image, PIL.Image.Image)
|
||||
and not isinstance(image, list)
|
||||
):
|
||||
raise ValueError(
|
||||
"`image` has to be of type `torch.FloatTensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
|
||||
f" {type(image)}"
|
||||
)
|
||||
|
||||
if height % 8 != 0 or width % 8 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size,
|
||||
num_frames,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
dtype,
|
||||
device,
|
||||
generator,
|
||||
latents=None,
|
||||
):
|
||||
shape = (
|
||||
batch_size,
|
||||
num_frames,
|
||||
num_channels_latents // 2,
|
||||
height // self.vae_scale_factor,
|
||||
width // self.vae_scale_factor,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
@property
|
||||
def do_classifier_free_guidance(self):
|
||||
return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.FloatTensor],
|
||||
height: int = 576,
|
||||
width: int = 1024,
|
||||
num_frames: Optional[int] = None,
|
||||
num_inference_steps: int = 25,
|
||||
min_guidance_scale: float = 1.0,
|
||||
max_guidance_scale: float = 3.0,
|
||||
fps: int = 7,
|
||||
motion_bucket_id: int = 127,
|
||||
noise_aug_strength: int = 0.02,
|
||||
decode_chunk_size: Optional[int] = None,
|
||||
num_videos_per_prompt: Optional[int] = 1,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
cache_interval: Optional[int] = 1,
|
||||
cache_branch: Optional[int] = None,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
r"""
|
||||
The call function to the pipeline for generation.
|
||||
|
||||
Args:
|
||||
image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
|
||||
Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
|
||||
[`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
|
||||
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The height in pixels of the generated image.
|
||||
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
|
||||
The width in pixels of the generated image.
|
||||
num_frames (`int`, *optional*):
|
||||
The number of video frames to generate. Defaults to 14 for `stable-video-diffusion-img2vid` and to 25 for `stable-video-diffusion-img2vid-xt`
|
||||
num_inference_steps (`int`, *optional*, defaults to 25):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference. This parameter is modulated by `strength`.
|
||||
min_guidance_scale (`float`, *optional*, defaults to 1.0):
|
||||
The minimum guidance scale. Used for the classifier free guidance with first frame.
|
||||
max_guidance_scale (`float`, *optional*, defaults to 3.0):
|
||||
The maximum guidance scale. Used for the classifier free guidance with last frame.
|
||||
fps (`int`, *optional*, defaults to 7):
|
||||
Frames per second. The rate at which the generated images shall be exported to a video after generation.
|
||||
Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training.
|
||||
motion_bucket_id (`int`, *optional*, defaults to 127):
|
||||
The motion bucket ID. Used as conditioning for the generation. The higher the number the more motion will be in the video.
|
||||
noise_aug_strength (`int`, *optional*, defaults to 0.02):
|
||||
The amount of noise added to the init image, the higher it is the less the video will look like the init image. Increase it for more motion.
|
||||
decode_chunk_size (`int`, *optional*):
|
||||
The number of frames to decode at a time. The higher the chunk size, the higher the temporal consistency
|
||||
between frames, but also the higher the memory consumption. By default, the decoder will decode all frames at once
|
||||
for maximal quality. Reduce `decode_chunk_size` to reduce memory usage.
|
||||
num_videos_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
|
||||
generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor is generated by sampling using the supplied random `generator`.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
|
||||
callback_on_step_end (`Callable`, *optional*):
|
||||
A function that calls at the end of each denoising steps during the inference. The function is called
|
||||
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
|
||||
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
|
||||
`callback_on_step_end_tensor_inputs`.
|
||||
callback_on_step_end_tensor_inputs (`List`, *optional*):
|
||||
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
|
||||
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
|
||||
`._callback_tensor_inputs` attribute of your pipeline class.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
||||
plain tuple.
|
||||
|
||||
Returns:
|
||||
[`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`:
|
||||
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is returned,
|
||||
otherwise a `tuple` is returned where the first element is a list of list with the generated frames.
|
||||
|
||||
Examples:
|
||||
|
||||
```py
|
||||
from diffusers import StableVideoDiffusionPipeline
|
||||
from diffusers.utils import load_image, export_to_video
|
||||
|
||||
pipe = StableVideoDiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16")
|
||||
pipe.to("cuda")
|
||||
|
||||
image = load_image("https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200")
|
||||
image = image.resize((1024, 576))
|
||||
|
||||
frames = pipe(image, num_frames=25, decode_chunk_size=8).frames[0]
|
||||
export_to_video(frames, "generated.mp4", fps=7)
|
||||
```
|
||||
"""
|
||||
# 0. Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
num_frames = num_frames if num_frames is not None else self.unet.config.num_frames
|
||||
decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(image, height, width)
|
||||
|
||||
# 2. Define call parameters
|
||||
if isinstance(image, PIL.Image.Image):
|
||||
batch_size = 1
|
||||
elif isinstance(image, list):
|
||||
batch_size = len(image)
|
||||
else:
|
||||
batch_size = image.shape[0]
|
||||
device = self._execution_device
|
||||
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
||||
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
||||
# corresponds to doing no classifier free guidance.
|
||||
do_classifier_free_guidance = max_guidance_scale > 1.0
|
||||
|
||||
# 3. Encode input image
|
||||
image_embeddings = self._encode_image(image, device, num_videos_per_prompt, do_classifier_free_guidance)
|
||||
|
||||
# NOTE: Stable Diffusion Video was conditioned on fps - 1, which
|
||||
# is why it is reduced here.
|
||||
# See: https://github.com/Stability-AI/generative-models/blob/ed0997173f98eaf8f4edf7ba5fe8f15c6b877fd3/scripts/sampling/simple_video_sample.py#L188
|
||||
fps = fps - 1
|
||||
|
||||
# 4. Encode input image using VAE
|
||||
image = self.image_processor.preprocess(image, height=height, width=width)
|
||||
noise = randn_tensor(image.shape, generator=generator, device=image.device, dtype=image.dtype)
|
||||
image = image + noise_aug_strength * noise
|
||||
|
||||
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
|
||||
if needs_upcasting:
|
||||
self.vae.to(dtype=torch.float32)
|
||||
|
||||
image_latents = self._encode_vae_image(image, device, num_videos_per_prompt, do_classifier_free_guidance)
|
||||
image_latents = image_latents.to(image_embeddings.dtype)
|
||||
|
||||
# cast back to fp16 if needed
|
||||
if needs_upcasting:
|
||||
self.vae.to(dtype=torch.float16)
|
||||
|
||||
# Repeat the image latents for each frame so we can concatenate them with the noise
|
||||
# image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width]
|
||||
image_latents = image_latents.unsqueeze(1).repeat(1, num_frames, 1, 1, 1)
|
||||
|
||||
# 5. Get Added Time IDs
|
||||
added_time_ids = self._get_add_time_ids(
|
||||
fps,
|
||||
motion_bucket_id,
|
||||
noise_aug_strength,
|
||||
image_embeddings.dtype,
|
||||
batch_size,
|
||||
num_videos_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
)
|
||||
added_time_ids = added_time_ids.to(device)
|
||||
|
||||
# 4. Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
# 5. Prepare latent variables
|
||||
num_channels_latents = self.unet.config.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_videos_per_prompt,
|
||||
num_frames,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
image_embeddings.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 7. Prepare guidance scale
|
||||
guidance_scale = torch.linspace(min_guidance_scale, max_guidance_scale, num_frames).unsqueeze(0)
|
||||
guidance_scale = guidance_scale.to(device, latents.dtype)
|
||||
guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1)
|
||||
guidance_scale = _append_dims(guidance_scale, latents.ndim)
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
|
||||
cache_features = None
|
||||
interval_seq = list(range(0, num_inference_steps, cache_interval))
|
||||
interval_seq = sorted(interval_seq)
|
||||
|
||||
# 8. Denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
self._num_timesteps = len(timesteps)
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
# Concatenate image_latents over channels dimention
|
||||
latent_model_input = torch.cat([latent_model_input, image_latents], dim=2)
|
||||
|
||||
if i in interval_seq:
|
||||
cache_features = None
|
||||
|
||||
# predict the noise residual
|
||||
noise_pred, cache_features = self.unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=image_embeddings,
|
||||
added_time_ids=added_time_ids,
|
||||
cache_features=cache_features,
|
||||
cache_branch=cache_branch,
|
||||
return_dict=False,
|
||||
)
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(noise_pred, t, latents).prev_sample
|
||||
|
||||
if callback_on_step_end is not None:
|
||||
callback_kwargs = {}
|
||||
for k in callback_on_step_end_tensor_inputs:
|
||||
callback_kwargs[k] = locals()[k]
|
||||
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
||||
|
||||
latents = callback_outputs.pop("latents", latents)
|
||||
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
|
||||
if not output_type == "latent":
|
||||
# cast back to fp16 if needed
|
||||
if needs_upcasting:
|
||||
self.vae.to(dtype=torch.float16)
|
||||
frames = self.decode_latents(latents, num_frames, decode_chunk_size)
|
||||
frames = tensor2vid(frames, self.image_processor, output_type=output_type)
|
||||
else:
|
||||
frames = latents
|
||||
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if not return_dict:
|
||||
return frames
|
||||
|
||||
return StableVideoDiffusionPipelineOutput(frames=frames)
|
||||
|
||||
|
||||
# resizing utils
|
||||
# TODO: clean up later
|
||||
def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True):
|
||||
h, w = input.shape[-2:]
|
||||
factors = (h / size[0], w / size[1])
|
||||
|
||||
# First, we have to determine sigma
|
||||
# Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171
|
||||
sigmas = (
|
||||
max((factors[0] - 1.0) / 2.0, 0.001),
|
||||
max((factors[1] - 1.0) / 2.0, 0.001),
|
||||
)
|
||||
|
||||
# Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma
|
||||
# https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206
|
||||
# But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now
|
||||
ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))
|
||||
|
||||
# Make sure it is odd
|
||||
if (ks[0] % 2) == 0:
|
||||
ks = ks[0] + 1, ks[1]
|
||||
|
||||
if (ks[1] % 2) == 0:
|
||||
ks = ks[0], ks[1] + 1
|
||||
|
||||
input = _gaussian_blur2d(input, ks, sigmas)
|
||||
|
||||
output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
|
||||
return output
|
||||
|
||||
|
||||
def _compute_padding(kernel_size):
|
||||
"""Compute padding tuple."""
|
||||
# 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom)
|
||||
# https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
|
||||
if len(kernel_size) < 2:
|
||||
raise AssertionError(kernel_size)
|
||||
computed = [k - 1 for k in kernel_size]
|
||||
|
||||
# for even kernels we need to do asymmetric padding :(
|
||||
out_padding = 2 * len(kernel_size) * [0]
|
||||
|
||||
for i in range(len(kernel_size)):
|
||||
computed_tmp = computed[-(i + 1)]
|
||||
|
||||
pad_front = computed_tmp // 2
|
||||
pad_rear = computed_tmp - pad_front
|
||||
|
||||
out_padding[2 * i + 0] = pad_front
|
||||
out_padding[2 * i + 1] = pad_rear
|
||||
|
||||
return out_padding
|
||||
|
||||
|
||||
def _filter2d(input, kernel):
|
||||
# prepare kernel
|
||||
b, c, h, w = input.shape
|
||||
tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype)
|
||||
|
||||
tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
|
||||
|
||||
height, width = tmp_kernel.shape[-2:]
|
||||
|
||||
padding_shape: list[int] = _compute_padding([height, width])
|
||||
input = torch.nn.functional.pad(input, padding_shape, mode="reflect")
|
||||
|
||||
# kernel and input tensor reshape to align element-wise or batch-wise params
|
||||
tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
|
||||
input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))
|
||||
|
||||
# convolve the tensor with the kernel.
|
||||
output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
|
||||
|
||||
out = output.view(b, c, h, w)
|
||||
return out
|
||||
|
||||
|
||||
def _gaussian(window_size: int, sigma):
|
||||
if isinstance(sigma, float):
|
||||
sigma = torch.tensor([[sigma]])
|
||||
|
||||
batch_size = sigma.shape[0]
|
||||
|
||||
x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)
|
||||
|
||||
if window_size % 2 == 0:
|
||||
x = x + 0.5
|
||||
|
||||
gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
|
||||
|
||||
return gauss / gauss.sum(-1, keepdim=True)
|
||||
|
||||
|
||||
def _gaussian_blur2d(input, kernel_size, sigma):
|
||||
if isinstance(sigma, tuple):
|
||||
sigma = torch.tensor([sigma], dtype=input.dtype)
|
||||
else:
|
||||
sigma = sigma.to(dtype=input.dtype)
|
||||
|
||||
ky, kx = int(kernel_size[0]), int(kernel_size[1])
|
||||
bs = sigma.shape[0]
|
||||
kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
|
||||
kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
|
||||
out_x = _filter2d(input, kernel_x[..., None, :])
|
||||
out = _filter2d(out_x, kernel_y[..., None])
|
||||
|
||||
return out
|
||||
2108
ixformer_sdk/contrib/DeepCache/svd/pipeline_utils.py
Normal file
2108
ixformer_sdk/contrib/DeepCache/svd/pipeline_utils.py
Normal file
File diff suppressed because it is too large
Load Diff
2412
ixformer_sdk/contrib/DeepCache/svd/unet_3d_blocks.py
Normal file
2412
ixformer_sdk/contrib/DeepCache/svd/unet_3d_blocks.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,566 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.loaders import UNet2DConditionLoadersMixin
|
||||
from diffusers.utils import BaseOutput, logging
|
||||
from diffusers.models.attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
|
||||
from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class UNetSpatioTemporalConditionOutput(BaseOutput):
|
||||
"""
|
||||
The output of [`UNetSpatioTemporalConditionModel`].
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
|
||||
The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor = None
|
||||
|
||||
|
||||
class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
|
||||
r"""
|
||||
A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and returns a sample
|
||||
shaped output.
|
||||
|
||||
This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
|
||||
for all models (such as downloading or saving).
|
||||
|
||||
Parameters:
|
||||
sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
|
||||
Height and width of input/output sample.
|
||||
in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample.
|
||||
out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
|
||||
down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
|
||||
The tuple of downsample blocks to use.
|
||||
up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
|
||||
The tuple of upsample blocks to use.
|
||||
block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
|
||||
The tuple of output channels for each block.
|
||||
addition_time_embed_dim: (`int`, defaults to 256):
|
||||
Dimension to to encode the additional time ids.
|
||||
projection_class_embeddings_input_dim (`int`, defaults to 768):
|
||||
The dimension of the projection of encoded `added_time_ids`.
|
||||
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
|
||||
cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
|
||||
The dimension of the cross attention features.
|
||||
transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
|
||||
The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
|
||||
[`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
|
||||
[`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
|
||||
num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
|
||||
The number of attention heads.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
sample_size: Optional[int] = None,
|
||||
in_channels: int = 8,
|
||||
out_channels: int = 4,
|
||||
down_block_types: Tuple[str] = (
|
||||
"CrossAttnDownBlockSpatioTemporal",
|
||||
"CrossAttnDownBlockSpatioTemporal",
|
||||
"CrossAttnDownBlockSpatioTemporal",
|
||||
"DownBlockSpatioTemporal",
|
||||
),
|
||||
up_block_types: Tuple[str] = (
|
||||
"UpBlockSpatioTemporal",
|
||||
"CrossAttnUpBlockSpatioTemporal",
|
||||
"CrossAttnUpBlockSpatioTemporal",
|
||||
"CrossAttnUpBlockSpatioTemporal",
|
||||
),
|
||||
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
||||
addition_time_embed_dim: int = 256,
|
||||
projection_class_embeddings_input_dim: int = 768,
|
||||
layers_per_block: Union[int, Tuple[int]] = 2,
|
||||
cross_attention_dim: Union[int, Tuple[int]] = 1024,
|
||||
transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
|
||||
num_attention_heads: Union[int, Tuple[int]] = (5, 10, 10, 20),
|
||||
num_frames: int = 25,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_size = sample_size
|
||||
|
||||
# Check inputs
|
||||
if len(down_block_types) != len(up_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
|
||||
)
|
||||
|
||||
if len(block_out_channels) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
# input
|
||||
self.conv_in = nn.Conv2d(
|
||||
in_channels,
|
||||
block_out_channels[0],
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
)
|
||||
|
||||
# time
|
||||
time_embed_dim = block_out_channels[0] * 4
|
||||
|
||||
self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0)
|
||||
timestep_input_dim = block_out_channels[0]
|
||||
|
||||
self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||||
|
||||
self.add_time_proj = Timesteps(addition_time_embed_dim, True, downscale_freq_shift=0)
|
||||
self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
|
||||
if isinstance(num_attention_heads, int):
|
||||
num_attention_heads = (num_attention_heads,) * len(down_block_types)
|
||||
|
||||
if isinstance(cross_attention_dim, int):
|
||||
cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
|
||||
|
||||
if isinstance(layers_per_block, int):
|
||||
layers_per_block = [layers_per_block] * len(down_block_types)
|
||||
|
||||
if isinstance(transformer_layers_per_block, int):
|
||||
transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
|
||||
|
||||
blocks_time_embed_dim = time_embed_dim
|
||||
|
||||
# down
|
||||
output_channel = block_out_channels[0]
|
||||
for i, down_block_type in enumerate(down_block_types):
|
||||
input_channel = output_channel
|
||||
output_channel = block_out_channels[i]
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
down_block = get_down_block(
|
||||
down_block_type,
|
||||
num_layers=layers_per_block[i],
|
||||
transformer_layers_per_block=transformer_layers_per_block[i],
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
add_downsample=not is_final_block,
|
||||
resnet_eps=1e-5,
|
||||
cross_attention_dim=cross_attention_dim[i],
|
||||
num_attention_heads=num_attention_heads[i],
|
||||
resnet_act_fn="silu",
|
||||
)
|
||||
self.down_blocks.append(down_block)
|
||||
|
||||
# mid
|
||||
self.mid_block = UNetMidBlockSpatioTemporal(
|
||||
block_out_channels[-1],
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
transformer_layers_per_block=transformer_layers_per_block[-1],
|
||||
cross_attention_dim=cross_attention_dim[-1],
|
||||
num_attention_heads=num_attention_heads[-1],
|
||||
)
|
||||
|
||||
# count how many layers upsample the images
|
||||
self.num_upsamplers = 0
|
||||
|
||||
# up
|
||||
reversed_block_out_channels = list(reversed(block_out_channels))
|
||||
reversed_num_attention_heads = list(reversed(num_attention_heads))
|
||||
reversed_layers_per_block = list(reversed(layers_per_block))
|
||||
reversed_cross_attention_dim = list(reversed(cross_attention_dim))
|
||||
reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
|
||||
|
||||
output_channel = reversed_block_out_channels[0]
|
||||
for i, up_block_type in enumerate(up_block_types):
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
prev_output_channel = output_channel
|
||||
output_channel = reversed_block_out_channels[i]
|
||||
input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
|
||||
|
||||
# add upsample block for all BUT final layer
|
||||
if not is_final_block:
|
||||
add_upsample = True
|
||||
self.num_upsamplers += 1
|
||||
else:
|
||||
add_upsample = False
|
||||
|
||||
up_block = get_up_block(
|
||||
up_block_type,
|
||||
num_layers=reversed_layers_per_block[i] + 1,
|
||||
transformer_layers_per_block=reversed_transformer_layers_per_block[i],
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=1e-5,
|
||||
resolution_idx=i,
|
||||
cross_attention_dim=reversed_cross_attention_dim[i],
|
||||
num_attention_heads=reversed_num_attention_heads[i],
|
||||
resnet_act_fn="silu",
|
||||
)
|
||||
self.up_blocks.append(up_block)
|
||||
prev_output_channel = output_channel
|
||||
|
||||
# out
|
||||
self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5)
|
||||
self.conv_act = nn.SiLU()
|
||||
|
||||
self.conv_out = nn.Conv2d(
|
||||
block_out_channels[0],
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
)
|
||||
|
||||
@property
|
||||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||||
r"""
|
||||
Returns:
|
||||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||||
indexed by its weight name.
|
||||
"""
|
||||
# set recursively
|
||||
processors = {}
|
||||
|
||||
def fn_recursive_add_processors(
|
||||
name: str,
|
||||
module: torch.nn.Module,
|
||||
processors: Dict[str, AttentionProcessor],
|
||||
):
|
||||
if hasattr(module, "get_processor"):
|
||||
processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
|
||||
r"""
|
||||
Sets the attention processor to use to compute attention.
|
||||
|
||||
Parameters:
|
||||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
for **all** `Attention` layers.
|
||||
|
||||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||||
processor. This is strongly recommended when setting trainable attention processors.
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
def set_default_attn_processor(self):
|
||||
"""
|
||||
Disables custom attention processors and sets the default attention implementation.
|
||||
"""
|
||||
if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
|
||||
processor = AttnProcessor()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
|
||||
)
|
||||
|
||||
self.set_attn_processor(processor)
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
# Copied from diffusers.models.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
|
||||
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
|
||||
"""
|
||||
Sets the attention processor to use [feed forward
|
||||
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
|
||||
|
||||
Parameters:
|
||||
chunk_size (`int`, *optional*):
|
||||
The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
|
||||
over each tensor of dim=`dim`.
|
||||
dim (`int`, *optional*, defaults to `0`):
|
||||
The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
|
||||
or dim=1 (sequence length).
|
||||
"""
|
||||
if dim not in [0, 1]:
|
||||
raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
|
||||
|
||||
# By default chunk size is 1
|
||||
chunk_size = chunk_size or 1
|
||||
|
||||
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
|
||||
if hasattr(module, "set_chunk_feed_forward"):
|
||||
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_feed_forward(child, chunk_size, dim)
|
||||
|
||||
for module in self.children():
|
||||
fn_recursive_feed_forward(module, chunk_size, dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
added_time_ids: torch.Tensor,
|
||||
cache_features: Optional[torch.Tensor] = None,
|
||||
cache_branch: Optional[int] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
|
||||
r"""
|
||||
The [`UNetSpatioTemporalConditionModel`] forward method.
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor`):
|
||||
The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`.
|
||||
timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
|
||||
encoder_hidden_states (`torch.FloatTensor`):
|
||||
The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`.
|
||||
added_time_ids: (`torch.FloatTensor`):
|
||||
The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal
|
||||
embeddings and added to the time embeddings.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead of a plain
|
||||
tuple.
|
||||
Returns:
|
||||
[`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`:
|
||||
If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is returned, otherwise
|
||||
a `tuple` is returned where the first element is the sample tensor.
|
||||
"""
|
||||
# 1. time
|
||||
timesteps = timestep
|
||||
if not torch.is_tensor(timesteps):
|
||||
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
||||
# This would be a good case for the `match` statement (Python 3.10+)
|
||||
is_mps = sample.device.type == "mps"
|
||||
if isinstance(timestep, float):
|
||||
dtype = torch.float32 if is_mps else torch.float64
|
||||
else:
|
||||
dtype = torch.int32 if is_mps else torch.int64
|
||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
||||
elif len(timesteps.shape) == 0:
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
batch_size, num_frames = sample.shape[:2]
|
||||
timesteps = timesteps.expand(batch_size)
|
||||
|
||||
t_emb = self.time_proj(timesteps)
|
||||
|
||||
# `Timesteps` does not contain any weights and will always return f32 tensors
|
||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
t_emb = t_emb.to(dtype=sample.dtype)
|
||||
|
||||
emb = self.time_embedding(t_emb)
|
||||
|
||||
time_embeds = self.add_time_proj(added_time_ids.flatten())
|
||||
time_embeds = time_embeds.reshape((batch_size, -1))
|
||||
time_embeds = time_embeds.to(emb.dtype)
|
||||
aug_emb = self.add_embedding(time_embeds)
|
||||
emb = emb + aug_emb
|
||||
|
||||
# Flatten the batch and frames dimensions
|
||||
# sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width]
|
||||
sample = sample.flatten(0, 1)
|
||||
# Repeat the embeddings num_video_frames times
|
||||
# emb: [batch, channels] -> [batch * frames, channels]
|
||||
emb = emb.repeat_interleave(num_frames, dim=0)
|
||||
# encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels]
|
||||
encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
|
||||
|
||||
# 2. pre-process
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device)
|
||||
|
||||
# Branch: 4 down_blocks, each with 3 skip connections. Here we ignore the first skip branch, whose computations only has up_blocks but without down_blocks.
|
||||
if cache_branch is not None:
|
||||
each_module_num = len(self.down_blocks[0].resnets) + 1
|
||||
down_cache_block_idx = cache_branch // each_module_num
|
||||
down_cache_module_idx = cache_branch % each_module_num
|
||||
|
||||
up_cache_block_idx = len(self.up_blocks) - 1 - down_cache_block_idx
|
||||
up_cache_module_idx = 1 - down_cache_module_idx
|
||||
if down_cache_module_idx == each_module_num - 1:
|
||||
up_cache_block_idx -= 1
|
||||
up_cache_module_idx = 2
|
||||
|
||||
if cache_features is not None:
|
||||
# 3. down
|
||||
down_block_res_samples = (sample,)
|
||||
for block_id, downsample_block in enumerate(self.down_blocks):
|
||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
image_only_indicator=image_only_indicator,
|
||||
exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
image_only_indicator=image_only_indicator,
|
||||
exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None
|
||||
)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
if down_cache_block_idx == block_id:
|
||||
break
|
||||
|
||||
# 4. no mid
|
||||
sample = cache_features
|
||||
|
||||
# 5. up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
if i < up_cache_block_idx:
|
||||
continue
|
||||
|
||||
if i == up_cache_block_idx:
|
||||
trunc_res_samples_len = len(upsample_block.resnets) - up_cache_module_idx
|
||||
else:
|
||||
trunc_res_samples_len = len(upsample_block.resnets)
|
||||
|
||||
res_samples = down_block_res_samples[-trunc_res_samples_len :]
|
||||
down_block_res_samples = down_block_res_samples[: -trunc_res_samples_len]
|
||||
|
||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
||||
sample, _ = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
image_only_indicator=image_only_indicator,
|
||||
enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None
|
||||
)
|
||||
else:
|
||||
sample, _ = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
image_only_indicator=image_only_indicator,
|
||||
enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None
|
||||
)
|
||||
else:
|
||||
# 3. down
|
||||
down_block_res_samples = (sample,)
|
||||
for downsample_block in self.down_blocks:
|
||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
# 4. mid
|
||||
sample = self.mid_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
|
||||
# 5. up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
||||
|
||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
||||
sample, current_record_f = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
else:
|
||||
sample, current_record_f = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
|
||||
if cache_branch is not None and i == up_cache_block_idx:
|
||||
cache_features = current_record_f[up_cache_module_idx]
|
||||
|
||||
# 6. post-process
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
# 7. Reshape back to original shape
|
||||
sample = sample.reshape(batch_size, num_frames, *sample.shape[1:])
|
||||
|
||||
if not return_dict:
|
||||
return (sample, cache_features)
|
||||
|
||||
return UNetSpatioTemporalConditionOutput(sample=sample)
|
||||
0
ixformer_sdk/contrib/__init__.py
Normal file
0
ixformer_sdk/contrib/__init__.py
Normal file
0
ixformer_sdk/contrib/comfy/__init__.py
Normal file
0
ixformer_sdk/contrib/comfy/__init__.py
Normal file
496
ixformer_sdk/contrib/comfy/unet_model_wrapper.py
Normal file
496
ixformer_sdk/contrib/comfy/unet_model_wrapper.py
Normal file
@@ -0,0 +1,496 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import ixformer.functions as ixf_F
|
||||
|
||||
|
||||
def time_embed(t_emb, weight1, bias1, weight2, bias2):
|
||||
# unet time_emd
|
||||
# linear + silu + linear
|
||||
emb = ixf_F.act_bias_mm(
|
||||
t_emb, weight1, act_type="silu", bias=bias1, scale=1, trans_format="TN"
|
||||
)
|
||||
emb = ixf_F.act_bias_mm(
|
||||
emb, weight2, act_type="none", bias=bias2, scale=1, trans_format="TN"
|
||||
)
|
||||
return emb
|
||||
|
||||
|
||||
def ixf_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
|
||||
return ixf_F.layernorm(input, weight, bias, normalized_shape)
|
||||
|
||||
|
||||
def ixf_pt_scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
|
||||
):
|
||||
if (
|
||||
not query.is_contiguous()
|
||||
and query.transpose(1, 2).is_contiguous()
|
||||
and key.transpose(1, 2).is_contiguous()
|
||||
and value.transpose(1, 2).is_contiguous()
|
||||
and attn_mask is None
|
||||
):
|
||||
|
||||
batch_size, head_num, seq_len_q, head_dim = query.shape
|
||||
_, _, seq_len_k, _ = key.shape
|
||||
|
||||
query = query.transpose(1, 2).view(batch_size * seq_len_q, head_num, head_dim)
|
||||
key = key.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim)
|
||||
value = value.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim)
|
||||
|
||||
cu_seqlens_q = torch.arange(
|
||||
0,
|
||||
seq_len_q * (batch_size + 1),
|
||||
seq_len_q,
|
||||
dtype=torch.int32,
|
||||
device=query.device,
|
||||
)
|
||||
if seq_len_q == seq_len_k:
|
||||
cu_seqlens_k = cu_seqlens_q
|
||||
else:
|
||||
cu_seqlens_k = torch.arange(
|
||||
0,
|
||||
seq_len_k * (batch_size + 1),
|
||||
seq_len_k,
|
||||
dtype=torch.int32,
|
||||
device=query.device,
|
||||
)
|
||||
|
||||
res = ixf_F.flash_attn_varlen_func(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seqlens_q.int(),
|
||||
cu_seqlens_k.int(),
|
||||
seq_len_q,
|
||||
seq_len_k,
|
||||
)
|
||||
res = res.view(batch_size, seq_len_q, head_num, head_dim).transpose(1, 2)
|
||||
return res
|
||||
|
||||
if not query.is_contiguous():
|
||||
query = query.contiguous()
|
||||
if not key.is_contiguous():
|
||||
key = key.contiguous()
|
||||
if not value.is_contiguous():
|
||||
value = value.contiguous()
|
||||
return ixf_F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attn_mask, is_causal=is_causal
|
||||
)
|
||||
|
||||
|
||||
class UnetIxformerFunction:
|
||||
def __init__(self) -> None:
|
||||
self.ixf_linear = ixf_F.linear
|
||||
self.pt_linear = F.linear
|
||||
self.pt_layer_norm = F.layer_norm
|
||||
self.pt_scaled_dot_product_attention = F.scaled_dot_product_attention
|
||||
|
||||
def __enter__(self):
|
||||
F.linear = self.ixf_linear
|
||||
F.layer_norm = ixf_layer_norm
|
||||
F.scaled_dot_product_attention = ixf_pt_scaled_dot_product_attention
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
F.linear = self.pt_linear
|
||||
F.layer_norm = self.pt_layer_norm
|
||||
F.scaled_dot_product_attention = self.pt_scaled_dot_product_attention
|
||||
if exc_tb is not None:
|
||||
print(f"{exc_type} {exc_val}")
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def ForwardWrapper(fun):
|
||||
def wrap(*args, **kwargs):
|
||||
with UnetIxformerFunction() as w:
|
||||
return fun(*args, **kwargs)
|
||||
|
||||
return wrap
|
||||
|
||||
|
||||
class IxformerComfyWrapper(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.is_ixf_wrapper = True
|
||||
|
||||
|
||||
class Conv2dNhwcWrapper(IxformerComfyWrapper):
|
||||
def __init__(self, module):
|
||||
super().__init__()
|
||||
module.weight.data = module.weight.permute(0, 2, 3, 1).contiguous()
|
||||
module.bias.data = module.bias.float()
|
||||
self.weight = module.weight.data
|
||||
self.bias = module.bias.data
|
||||
self.stride = module.stride
|
||||
self.padding = module.padding
|
||||
self.dilation = module.dilation
|
||||
self.groups = module.groups
|
||||
|
||||
def forward(self, x):
|
||||
h2 = ixf_F.conv2d(
|
||||
x,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
return h2
|
||||
|
||||
|
||||
class ResBlockNhwcWrapper(IxformerComfyWrapper):
|
||||
def __init__(self, module) -> None:
|
||||
super().__init__()
|
||||
assert not module.updown
|
||||
assert not module.use_scale_shift_norm
|
||||
assert not module.skip_t_emb
|
||||
assert not module.exchange_temb_dims
|
||||
|
||||
if isinstance(module.skip_connection, nn.Identity):
|
||||
self.skip_connection = module.skip_connection
|
||||
elif get_class_name(module.skip_connection) == "Conv2d":
|
||||
self.skip_connection = Conv2dNhwcWrapper(module.skip_connection)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"ResBlockNhwcWrapper support Conv2d or nn.Identity, but got {module.skip_connection}"
|
||||
)
|
||||
|
||||
self.in_layers = module.in_layers
|
||||
self.out_layers = module.out_layers
|
||||
self.emb_layers = module.emb_layers
|
||||
self.in_layers_conv = Conv2dNhwcWrapper(module.in_layers[2])
|
||||
self.out_layers_conv = Conv2dNhwcWrapper(module.out_layers[3])
|
||||
|
||||
def forward(self, x, emb):
|
||||
# x: nhwc
|
||||
x1 = x
|
||||
# print(x1.shape)
|
||||
# fused group_norm silu
|
||||
h = ixf_F.group_norm(
|
||||
x1, # nchw->nhwc
|
||||
self.in_layers[0].num_groups,
|
||||
self.in_layers[0].weight,
|
||||
self.in_layers[0].bias,
|
||||
format=False,
|
||||
act_type=1,
|
||||
)
|
||||
h = self.in_layers_conv(h)
|
||||
|
||||
emb_out = self.emb_layers(emb)
|
||||
while len(emb_out.shape) < len(h.shape):
|
||||
emb_out = emb_out[..., None]
|
||||
|
||||
h = h + emb_out.permute(0, 2, 3, 1)
|
||||
# print(h.shape)
|
||||
h = ixf_F.group_norm(
|
||||
h,
|
||||
self.out_layers[0].num_groups,
|
||||
self.out_layers[0].weight,
|
||||
self.out_layers[0].bias,
|
||||
format=False,
|
||||
act_type=1,
|
||||
)
|
||||
|
||||
h = self.out_layers[2](h)
|
||||
h = self.out_layers_conv(h)
|
||||
# TODO: support other skip_connection
|
||||
return self.skip_connection(x) + h
|
||||
|
||||
|
||||
class DownsampleNhwcWrapper(IxformerComfyWrapper):
|
||||
def __init__(self, module) -> None:
|
||||
# TODO: support avg_pool_nd
|
||||
super().__init__()
|
||||
assert module.use_conv
|
||||
self.channels = module.channels
|
||||
self.op = Conv2dNhwcWrapper(module.op)
|
||||
|
||||
def forward(self, x):
|
||||
assert x.shape[-1] == self.channels
|
||||
return self.op(x)
|
||||
|
||||
|
||||
def ffn_forward(self, x):
|
||||
if get_class_name(self.net[0]) == "GEGLU":
|
||||
net = self.net[1:]
|
||||
geglu_net = self.net[0]
|
||||
x = geglu_net.proj(x)
|
||||
x = ixf_F.gelu_and_mul(x)
|
||||
return net(x)
|
||||
else:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
# ComfyUI/comfy/ldm/modules/attention.py `class BasicTransformerBlock(nn.Module)`
|
||||
def transformer_block_forward(self, x, context=None, transformer_options={}):
|
||||
extra_options = {}
|
||||
block = transformer_options.get("block", None)
|
||||
block_index = transformer_options.get("block_index", 0)
|
||||
transformer_patches = {}
|
||||
transformer_patches_replace = {}
|
||||
|
||||
for k in transformer_options:
|
||||
if k == "patches":
|
||||
transformer_patches = transformer_options[k]
|
||||
elif k == "patches_replace":
|
||||
transformer_patches_replace = transformer_options[k]
|
||||
else:
|
||||
extra_options[k] = transformer_options[k]
|
||||
|
||||
extra_options["n_heads"] = self.n_heads
|
||||
extra_options["dim_head"] = self.d_head
|
||||
|
||||
if self.ff_in:
|
||||
x_skip = x
|
||||
x = self.ff_in(self.norm_in(x))
|
||||
if self.is_res:
|
||||
x += x_skip
|
||||
|
||||
n = self.norm1(x)
|
||||
if self.disable_self_attn:
|
||||
context_attn1 = context
|
||||
else:
|
||||
context_attn1 = None
|
||||
value_attn1 = None
|
||||
|
||||
if "attn1_patch" in transformer_patches:
|
||||
patch = transformer_patches["attn1_patch"]
|
||||
if context_attn1 is None:
|
||||
context_attn1 = n
|
||||
value_attn1 = context_attn1
|
||||
for p in patch:
|
||||
n, context_attn1, value_attn1 = p(
|
||||
n, context_attn1, value_attn1, extra_options
|
||||
)
|
||||
|
||||
if block is not None:
|
||||
transformer_block = (block[0], block[1], block_index)
|
||||
else:
|
||||
transformer_block = None
|
||||
|
||||
attn1_replace_patch = transformer_patches_replace.get("attn1", {})
|
||||
block_attn1 = transformer_block
|
||||
if block_attn1 not in attn1_replace_patch:
|
||||
block_attn1 = block
|
||||
|
||||
if block_attn1 in attn1_replace_patch:
|
||||
if context_attn1 is None:
|
||||
context_attn1 = n
|
||||
value_attn1 = n
|
||||
n = self.attn1.to_q(n)
|
||||
context_attn1 = self.attn1.to_k(context_attn1)
|
||||
value_attn1 = self.attn1.to_v(value_attn1)
|
||||
n = attn1_replace_patch[block_attn1](
|
||||
n, context_attn1, value_attn1, extra_options
|
||||
)
|
||||
n = self.attn1.to_out(n)
|
||||
else:
|
||||
n = self.attn1(n, context=context_attn1, value=value_attn1)
|
||||
|
||||
if "attn1_output_patch" in transformer_patches:
|
||||
patch = transformer_patches["attn1_output_patch"]
|
||||
for p in patch:
|
||||
n = p(n, extra_options)
|
||||
|
||||
x += n
|
||||
if "middle_patch" in transformer_patches:
|
||||
patch = transformer_patches["middle_patch"]
|
||||
for p in patch:
|
||||
x = p(x, extra_options)
|
||||
|
||||
if self.attn2 is not None:
|
||||
n = self.norm2(x)
|
||||
if self.switch_temporal_ca_to_sa:
|
||||
context_attn2 = n
|
||||
else:
|
||||
context_attn2 = context
|
||||
value_attn2 = None
|
||||
if "attn2_patch" in transformer_patches:
|
||||
patch = transformer_patches["attn2_patch"]
|
||||
value_attn2 = context_attn2
|
||||
for p in patch:
|
||||
n, context_attn2, value_attn2 = p(
|
||||
n, context_attn2, value_attn2, extra_options
|
||||
)
|
||||
|
||||
attn2_replace_patch = transformer_patches_replace.get("attn2", {})
|
||||
block_attn2 = transformer_block
|
||||
if block_attn2 not in attn2_replace_patch:
|
||||
block_attn2 = block
|
||||
|
||||
if block_attn2 in attn2_replace_patch:
|
||||
if value_attn2 is None:
|
||||
value_attn2 = context_attn2
|
||||
n = self.attn2.to_q(n)
|
||||
context_attn2 = self.attn2.to_k(context_attn2)
|
||||
value_attn2 = self.attn2.to_v(value_attn2)
|
||||
n = attn2_replace_patch[block_attn2](
|
||||
n, context_attn2, value_attn2, extra_options
|
||||
)
|
||||
n = self.attn2.to_out(n)
|
||||
else:
|
||||
n = self.attn2(n, context=context_attn2, value=value_attn2)
|
||||
|
||||
if "attn2_output_patch" in transformer_patches:
|
||||
patch = transformer_patches["attn2_output_patch"]
|
||||
for p in patch:
|
||||
n = p(n, extra_options)
|
||||
|
||||
# x += n
|
||||
# if self.is_res:
|
||||
# x_skip = x
|
||||
# x = self.ff(self.norm3(x))
|
||||
|
||||
x, x_skip = ixf_F.residual_layer_norm(
|
||||
n,
|
||||
self.norm3.normalized_shape,
|
||||
self.norm3.weight,
|
||||
self.norm3.bias,
|
||||
x,
|
||||
eps=self.norm3.eps,
|
||||
is_post_ln=False,
|
||||
)
|
||||
x = ffn_forward(self.ff, x)
|
||||
|
||||
# x = ffn_forward(self.ff, self.norm3(x))
|
||||
if self.is_res:
|
||||
x += x_skip
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class SpatialTransformerNhwcWrapper(IxformerComfyWrapper):
|
||||
def __init__(self, module):
|
||||
super().__init__()
|
||||
self.use_linear = module.use_linear
|
||||
self.transformer_blocks = module.transformer_blocks
|
||||
self.norm = module.norm
|
||||
if not self.use_linear:
|
||||
self.proj_in = Conv2dNhwcWrapper(module.proj_in)
|
||||
self.proj_out = Conv2dNhwcWrapper(module.proj_out)
|
||||
else:
|
||||
self.proj_in = module.proj_in
|
||||
self.proj_out = module.proj_out
|
||||
|
||||
@ForwardWrapper
|
||||
def forward(self, x, context=None, transformer_options={}):
|
||||
# note: if no context is given, cross-attention defaults to self-attention
|
||||
if not isinstance(context, list):
|
||||
context = [context] * len(self.transformer_blocks)
|
||||
|
||||
b, h, w, c = x.shape
|
||||
x_in = x
|
||||
|
||||
# group_norm
|
||||
x = ixf_F.group_norm(
|
||||
x,
|
||||
self.norm.num_groups,
|
||||
self.norm.weight,
|
||||
self.norm.bias,
|
||||
format=False,
|
||||
)
|
||||
# conv2d
|
||||
if not self.use_linear:
|
||||
x = self.proj_in(x)
|
||||
# n,(hw),c
|
||||
x = x.view(x.shape[0], -1, x.shape[-1])
|
||||
if self.use_linear:
|
||||
x = self.proj_in(x)
|
||||
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
transformer_options["block_index"] = i
|
||||
# x = block(x, context=context[i], transformer_options=transformer_options)
|
||||
x = transformer_block_forward(
|
||||
block, x, context=context[i], transformer_options=transformer_options
|
||||
)
|
||||
|
||||
if self.use_linear:
|
||||
x = self.proj_out(x)
|
||||
x = x.view(b, h, w, c)
|
||||
if not self.use_linear:
|
||||
x = self.proj_out(x)
|
||||
return x + x_in
|
||||
|
||||
|
||||
class UpsampleNhwcWrapper(IxformerComfyWrapper):
|
||||
def __init__(self, module) -> None:
|
||||
# TODO: support mhwc interpolate
|
||||
super().__init__()
|
||||
self.dims = module.dims
|
||||
self.use_conv = module.use_conv
|
||||
self.channels = module.channels
|
||||
if self.use_conv:
|
||||
self.conv = Conv2dNhwcWrapper(module.conv)
|
||||
|
||||
def forward(self, x, output_shape=None):
|
||||
# print("================== Upsample is running ==================")
|
||||
assert x.shape[-1] == self.channels
|
||||
assert len(x.shape) == 4
|
||||
|
||||
# nhwc -> nchw
|
||||
if output_shape is not None:
|
||||
assert len(output_shape) == 4
|
||||
output_shape = [
|
||||
output_shape[0],
|
||||
output_shape[3],
|
||||
output_shape[1],
|
||||
output_shape[2],
|
||||
]
|
||||
x = x.permute(0, 3, 1, 2).contiguous()
|
||||
if self.dims == 3:
|
||||
shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
|
||||
if output_shape is not None:
|
||||
shape[1] = output_shape[3]
|
||||
shape[2] = output_shape[4]
|
||||
else:
|
||||
shape = [x.shape[2] * 2, x.shape[3] * 2]
|
||||
if output_shape is not None:
|
||||
shape[0] = output_shape[2]
|
||||
shape[1] = output_shape[3]
|
||||
# TODO: interpolate 支持 nhwc， 去掉前后转置
|
||||
x = F.interpolate(x, size=shape, mode="nearest")
|
||||
# nchw -> nhwc
|
||||
x = x.permute(0, 2, 3, 1).contiguous()
|
||||
if self.use_conv:
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
unet_wrappers = {
|
||||
"Conv2d": Conv2dNhwcWrapper,
|
||||
"ResBlock": ResBlockNhwcWrapper,
|
||||
"Downsample": DownsampleNhwcWrapper,
|
||||
"SpatialTransformer": SpatialTransformerNhwcWrapper,
|
||||
"Upsample": UpsampleNhwcWrapper,
|
||||
}
|
||||
|
||||
|
||||
def get_class_name(module):
|
||||
return module.__class__.__name__
|
||||
|
||||
|
||||
def module_wrapper(module):
|
||||
# 将原始的 module 封装为 nhwc 模式
|
||||
module_name = get_class_name(module)
|
||||
assert (
|
||||
module_name == "TimestepEmbedSequential"
|
||||
), f"ixformer unet_model_wrapper only support 'TimestepEmbedSequential' now, but got {module_name}"
|
||||
|
||||
num_sequential = len(module)
|
||||
for idx_seq in range(num_sequential):
|
||||
sub_module = module[idx_seq]
|
||||
sub_module_name = get_class_name(sub_module)
|
||||
# 判断模块是否已经封装
|
||||
if not getattr(sub_module, "is_ixf_wrapper", False):
|
||||
if sub_module_name in unet_wrappers:
|
||||
module[idx_seq].forward = unet_wrappers[sub_module_name](
|
||||
sub_module
|
||||
).forward
|
||||
module[idx_seq].is_ixf_wrapper = True
|
||||
else:
|
||||
raise NotImplementedError(f"{sub_module_name} not support")
|
||||
return module
|
||||
17
ixformer_sdk/contrib/flashinfer/__init__.py
Normal file
17
ixformer_sdk/contrib/flashinfer/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
from .decode import BatchDecodeWithPagedKVCacheWrapper
|
||||
from .prefill import (
|
||||
BatchPrefillWithPagedKVCacheWrapper,
|
||||
BatchPrefillWithRaggedKVCacheWrapper,
|
||||
)
|
||||
|
||||
|
||||
def bmm_fp8():
|
||||
pass
|
||||
|
||||
|
||||
def SegmentGEMMWrapper():
|
||||
pass
|
||||
|
||||
|
||||
def bmm_fp8():
|
||||
pass
|
||||
29
ixformer_sdk/contrib/flashinfer/activation.py
Normal file
29
ixformer_sdk/contrib/flashinfer/activation.py
Normal file
@@ -0,0 +1,29 @@
|
||||
import ixformer.inference.functions as ops
|
||||
import torch
|
||||
|
||||
|
||||
def gelu_and_mul():
|
||||
pass
|
||||
|
||||
|
||||
def gelu_tanh_and_mul():
|
||||
pass
|
||||
|
||||
|
||||
def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
|
||||
r"""Fused SiLU and Mul operation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input: torch.Tensor
|
||||
Input tensor, shape (..., 2 * hidden_size).
|
||||
|
||||
out: Optional[torch.Tensor]
|
||||
The the output tensor, if specified, the kernel will update this tensor inplace.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output: torch.Tensor
|
||||
Output tensor, shape (..., hidden_size).
|
||||
"""
|
||||
return ops.silu_and_mul(input=input, output=out)
|
||||
2
ixformer_sdk/contrib/flashinfer/cascade.py
Normal file
2
ixformer_sdk/contrib/flashinfer/cascade.py
Normal file
@@ -0,0 +1,2 @@
|
||||
def merge_state():
|
||||
pass
|
||||
101
ixformer_sdk/contrib/flashinfer/decode.py
Normal file
101
ixformer_sdk/contrib/flashinfer/decode.py
Normal file
@@ -0,0 +1,101 @@
|
||||
import math
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import ixformer.inference.functions as ops
|
||||
import torch
|
||||
|
||||
|
||||
def _grouped_size_compiled_for_decode_kernels(
|
||||
num_qo_heads: int, num_kv_heads: int
|
||||
) -> bool:
|
||||
return (num_qo_heads // num_kv_heads) in [1, 2, 4, 8]
|
||||
|
||||
|
||||
class BatchDecodeWithPagedKVCacheWrapper:
|
||||
def __init__(
|
||||
self,
|
||||
float_workspace_buffer: torch.Tensor,
|
||||
kv_layout: str = "NHD",
|
||||
use_cuda_graph: bool = False,
|
||||
use_tensor_cores: bool = False,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def plan(
|
||||
self,
|
||||
indptr: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
last_page_len: torch.Tensor,
|
||||
num_qo_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_dim: int,
|
||||
page_size: int,
|
||||
# pos_encoding_mode: str = "NONE",
|
||||
# window_left: int = -1,
|
||||
# logits_soft_cap: Optional[float] = None,
|
||||
data_type: Union[str, torch.dtype] = "float16",
|
||||
q_data_type: Optional[Union[str, torch.dtype]] = None,
|
||||
sm_scale: Optional[float] = None,
|
||||
# rope_scale: Optional[float] = None,
|
||||
# rope_theta: Optional[float] = None,
|
||||
max_seqlen_q: int = None,
|
||||
max_seqlen_k: int = None,
|
||||
) -> None:
|
||||
self.indptr = indptr
|
||||
self.indices = indices
|
||||
self.last_page_len = last_page_len
|
||||
self.num_qo_heads = num_qo_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.head_dim = head_dim
|
||||
|
||||
assert page_size == 1
|
||||
|
||||
self.cu_seqlens_q = torch.ones_like(indptr)
|
||||
self.cu_seqlens_q[0] = 0
|
||||
self.cu_seqlens_q = torch.cumsum(self.cu_seqlens_q, dim=0).int()
|
||||
|
||||
self.cu_seqlens_k = indptr
|
||||
if sm_scale is None:
|
||||
sm_scale = 1.0 / math.sqrt(head_dim)
|
||||
|
||||
self.sm_scale = sm_scale
|
||||
self.max_seqlen_q = max_seqlen_q
|
||||
self.max_seqlen_k = max_seqlen_k
|
||||
|
||||
begin_forward = plan
|
||||
|
||||
def forward(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
|
||||
pos_encoding_mode: str = "NONE",
|
||||
q_scale: Optional[float] = None,
|
||||
k_scale: Optional[float] = None,
|
||||
v_scale: Optional[float] = None,
|
||||
window_left: int = -1,
|
||||
logits_soft_cap: Optional[float] = None,
|
||||
sm_scale: Optional[float] = None,
|
||||
rope_scale: Optional[float] = None,
|
||||
rope_theta: Optional[float] = None,
|
||||
) -> torch.Tensor:
|
||||
k_cache, v_cache = paged_kv_cache
|
||||
|
||||
out = torch.empty_like(q)
|
||||
|
||||
ops.paged_attention_flashinfer(
|
||||
output=out,
|
||||
query=q,
|
||||
paged_kv_data=(k_cache.unsqueeze(1), v_cache.unsqueeze(1)),
|
||||
paged_kv_indptr=self.indptr,
|
||||
paged_kv_indices=self.indices,
|
||||
paged_kv_last_page_len=self.last_page_len,
|
||||
scale=self.sm_scale,
|
||||
max_seq_len=self.max_seqlen_k,
|
||||
kv_cache_format="NHD",
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
def end_forward(self) -> None:
|
||||
r"""Warning: this function is deprecated and has no effect."""
|
||||
pass
|
||||
61
ixformer_sdk/contrib/flashinfer/norm.py
Normal file
61
ixformer_sdk/contrib/flashinfer/norm.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import ixformer.inference.functions as ops
|
||||
import torch
|
||||
|
||||
|
||||
def fused_add_rmsnorm(
|
||||
input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
|
||||
):
|
||||
r"""Fused add root mean square normalization.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input: torch.Tensor
|
||||
Input tensor, shape (batch_size, hidden_size).
|
||||
residual: torch.Tensor
|
||||
Residual tensor, shape (batch_size, hidden_size).
|
||||
weight: torch.Tensor
|
||||
Weight tensor, shape (hidden_size,).
|
||||
eps: float
|
||||
Epsilon for numerical stability.
|
||||
"""
|
||||
return ops.residual_rms_norm(
|
||||
input=input,
|
||||
residual=residual,
|
||||
weight=weight,
|
||||
eps=eps,
|
||||
)
|
||||
|
||||
|
||||
def gemma_fused_add_rmsnorm():
|
||||
pass
|
||||
|
||||
|
||||
def gemma_rmsnorm():
|
||||
pass
|
||||
|
||||
|
||||
def rmsnorm(
|
||||
input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
|
||||
) -> torch.Tensor:
|
||||
r"""Root mean square normalization.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input: torch.Tensor
|
||||
Input tensor, shape (batch_size, hidden_size).
|
||||
weight: torch.Tensor
|
||||
Weight tensor, shape (hidden_size,).
|
||||
eps: float
|
||||
Epsilon for numerical stability.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output: torch.Tensor
|
||||
Normalized tensor, shape (batch_size, hidden_size).
|
||||
"""
|
||||
|
||||
return ops.rms_norm(
|
||||
input=input,
|
||||
weight=weight,
|
||||
eps=eps,
|
||||
)
|
||||
113
ixformer_sdk/contrib/flashinfer/prefill.py
Normal file
113
ixformer_sdk/contrib/flashinfer/prefill.py
Normal file
@@ -0,0 +1,113 @@
|
||||
import math
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
|
||||
class BatchPrefillWithRaggedKVCacheWrapper:
|
||||
def __init__(
|
||||
self,
|
||||
float_workspace_buffer: torch.Tensor,
|
||||
kv_layout: str = "NHD",
|
||||
):
|
||||
pass
|
||||
|
||||
def plan(
|
||||
self,
|
||||
qo_indptr: torch.Tensor,
|
||||
kv_indptr: torch.Tensor,
|
||||
num_qo_heads: int,
|
||||
num_kv_heads: int,
|
||||
head_dim: int,
|
||||
max_seqlen_q: int,
|
||||
max_seqlen_k: int,
|
||||
# custom_mask: Optional[torch.Tensor] = None,
|
||||
# packed_custom_mask: Optional[torch.Tensor] = None,
|
||||
causal: bool = True,
|
||||
# pos_encoding_mode: str = "NONE",
|
||||
# allow_fp16_qk_reduction: bool = False,
|
||||
# window_left: int = -1,
|
||||
# logits_soft_cap: Optional[float] = None,
|
||||
sm_scale: Optional[float] = None,
|
||||
# rope_scale: Optional[float] = None,
|
||||
# rope_theta: Optional[float] = None,
|
||||
# q_data_type: str = "float16",
|
||||
) -> None:
|
||||
batch_size = len(qo_indptr) - 1
|
||||
if len(kv_indptr) != batch_size + 1:
|
||||
raise ValueError(
|
||||
"The kv_indptr length should be equal to qk_indptr length."
|
||||
)
|
||||
self._causal = causal
|
||||
self._sm_scale = sm_scale
|
||||
if sm_scale is None:
|
||||
sm_scale = 1.0 / math.sqrt(head_dim)
|
||||
|
||||
self.cu_seqlens_q = qo_indptr
|
||||
self.cu_seqlens_k = kv_indptr
|
||||
self.num_qo_heads = num_qo_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.head_dim = head_dim
|
||||
self.max_seqlen_q = max_seqlen_q
|
||||
self.max_seqlen_k = max_seqlen_k
|
||||
|
||||
begin_forward = plan
|
||||
|
||||
def forward(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
causal: bool = True,
|
||||
# pos_encoding_mode: str = "NONE",
|
||||
# allow_fp16_qk_reduction: bool = False,
|
||||
# window_left: int = -1,
|
||||
logits_soft_cap: Optional[float] = None,
|
||||
sm_scale: Optional[float] = None,
|
||||
# rope_scale: Optional[float] = None,
|
||||
# rope_theta: Optional[float] = None,
|
||||
) -> torch.Tensor:
|
||||
r"""Warning: This function is deprecated, please use :meth:`run` instead."""
|
||||
|
||||
q = q.view(-1, self.num_qo_heads, self.head_dim)
|
||||
k = k.view(-1, self.num_kv_heads, self.head_dim)
|
||||
v = v.view(-1, self.num_kv_heads, self.head_dim)
|
||||
|
||||
out = torch.empty_like(q)
|
||||
|
||||
assert causal
|
||||
assert (
|
||||
logits_soft_cap is None or logits_soft_cap == 0
|
||||
), f"logits_soft_cap not supported, but got logits_soft_cap={logits_soft_cap}"
|
||||
|
||||
ops.infer.ixinfer_flash_attn_unpad(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
out,
|
||||
self.cu_seqlens_q,
|
||||
self.cu_seqlens_k,
|
||||
self.max_seqlen_q,
|
||||
self.max_seqlen_k,
|
||||
causal,
|
||||
False, # need_lse =False
|
||||
sm_scale,
|
||||
False,
|
||||
None,
|
||||
)
|
||||
return out
|
||||
|
||||
def end_forward(self) -> None:
|
||||
r"""Warning: this function is deprecated and has no effect."""
|
||||
pass
|
||||
|
||||
|
||||
class BatchPrefillWithPagedKVCacheWrapper:
|
||||
def __init__(
|
||||
self,
|
||||
float_workspace_buffer: torch.Tensor,
|
||||
kv_layout: str = "NHD",
|
||||
use_cuda_graph: bool = False,
|
||||
) -> None:
|
||||
pass
|
||||
14
ixformer_sdk/contrib/flashinfer/sampling.py
Normal file
14
ixformer_sdk/contrib/flashinfer/sampling.py
Normal file
@@ -0,0 +1,14 @@
|
||||
def min_p_sampling_from_probs():
|
||||
pass
|
||||
|
||||
|
||||
def top_k_renorm_prob():
|
||||
pass
|
||||
|
||||
|
||||
def top_k_top_p_sampling_from_probs():
|
||||
pass
|
||||
|
||||
|
||||
def top_p_renorm_prob():
|
||||
pass
|
||||
1
ixformer_sdk/contrib/tgi/__init__.py
Normal file
1
ixformer_sdk/contrib/tgi/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from .fused_moe import fused_moe
|
||||
430
ixformer_sdk/contrib/tgi/fused_moe.py
Normal file
430
ixformer_sdk/contrib/tgi/fused_moe.py
Normal file
@@ -0,0 +1,430 @@
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
from loguru import logger
|
||||
import torch
|
||||
import ixformer.inference.functions as ops
|
||||
|
||||
CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "65536"))
|
||||
|
||||
def fused_topk(
|
||||
hidden_states: torch.Tensor,
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
):
|
||||
assert hidden_states.shape[0] == gating_output.shape[0], (
|
||||
"Number of tokens mismatch")
|
||||
|
||||
M, _ = hidden_states.shape
|
||||
|
||||
topk_weights = torch.empty(M,
|
||||
topk,
|
||||
dtype=torch.float32,
|
||||
device=hidden_states.device)
|
||||
topk_ids = torch.empty(M,
|
||||
topk,
|
||||
dtype=torch.int32,
|
||||
device=hidden_states.device)
|
||||
token_expert_indicies = torch.empty(M,
|
||||
topk,
|
||||
dtype=torch.int32,
|
||||
device=hidden_states.device)
|
||||
ops.vllm_moe_topk_softmax(
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
token_expert_indicies,
|
||||
gating_output.float(), # TODO(woosuk): Optimize this.
|
||||
)
|
||||
del token_expert_indicies # Not used. Will be used in the future.
|
||||
|
||||
if renormalize:
|
||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
return topk_weights, topk_ids
|
||||
|
||||
# This is used by the Deepseek-V2 model
|
||||
def grouped_topk(hidden_states: torch.Tensor,
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
num_expert_group: int = 0,
|
||||
topk_group: int = 0):
|
||||
|
||||
assert hidden_states.shape[0] == gating_output.shape[0], (
|
||||
"Number of tokens mismatch")
|
||||
|
||||
scores = torch.softmax(gating_output, dim=-1)
|
||||
num_token = scores.shape[0]
|
||||
group_scores = scores.view(num_token, num_expert_group,
|
||||
-1).max(dim=-1).values # [n, n_group]
|
||||
group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
|
||||
sorted=False)[1] # [n, top_k_group]
|
||||
group_mask = torch.zeros_like(group_scores) # [n, n_group]
|
||||
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
|
||||
score_mask = group_mask.unsqueeze(-1).expand(
|
||||
num_token, num_expert_group,
|
||||
scores.shape[-1] // num_expert_group).reshape(num_token, -1) # [n, e]
|
||||
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
|
||||
topk_weights, topk_ids = torch.topk(tmp_scores,
|
||||
k=topk,
|
||||
dim=-1,
|
||||
sorted=False)
|
||||
|
||||
if renormalize:
|
||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
return topk_weights, topk_ids
|
||||
|
||||
def get_config_file_name(E: int, N: int, dtype: Optional[str]) -> str:
|
||||
device_name = torch.cuda.get_device_name().replace(" ", "_")
|
||||
dtype_selector = "" if not dtype else f",dtype={dtype}"
|
||||
return f"E={E},N={N},device_name={device_name}{dtype_selector}.json"
|
||||
|
||||
@functools.lru_cache
|
||||
def get_moe_configs(E: int, N: int,
|
||||
dtype: Optional[str]) -> Optional[Dict[int, Any]]:
|
||||
"""
|
||||
Return optimized configurations for the fused MoE kernel.
|
||||
|
||||
The return value will be a dictionary that maps an irregular grid of
|
||||
batch sizes to configurations of the fused_moe kernel. To evaluate the
|
||||
kernel on a given batch size bs, the closest batch size in the grid should
|
||||
be picked and the associated configuration chosen to invoke the kernel.
|
||||
"""
|
||||
|
||||
# First look up if an optimized configuration is available in the configs
|
||||
# directory
|
||||
json_file_name = get_config_file_name(E, N, dtype)
|
||||
|
||||
config_file_path = os.path.join(
|
||||
os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
|
||||
if os.path.exists(config_file_path):
|
||||
with open(config_file_path) as f:
|
||||
logger.info("Using configuration from %s for MoE layer.",
|
||||
config_file_path)
|
||||
# If a configuration has been found, return it
|
||||
return {int(key): val for key, val in json.load(f).items()}
|
||||
|
||||
# If no optimized configuration is available, we will use the default
|
||||
# configuration
|
||||
return None
|
||||
|
||||
def get_default_config(
|
||||
M: int,
|
||||
E: int,
|
||||
N: int,
|
||||
K: int,
|
||||
topk: int,
|
||||
dtype: Optional[str],
|
||||
) -> Dict[str, int]:
|
||||
config = {
|
||||
'BLOCK_SIZE_M': 64,
|
||||
'BLOCK_SIZE_N': 64,
|
||||
'BLOCK_SIZE_K': 32,
|
||||
'GROUP_SIZE_M': 8
|
||||
}
|
||||
if M <= E:
|
||||
config = {
|
||||
'BLOCK_SIZE_M': 16,
|
||||
'BLOCK_SIZE_N': 32,
|
||||
'BLOCK_SIZE_K': 64,
|
||||
'GROUP_SIZE_M': 1
|
||||
}
|
||||
numel = M * topk
|
||||
if numel <= 64:
|
||||
config['BLOCK_SIZE_M'] = 32
|
||||
elif numel <= 1024:
|
||||
config['BLOCK_SIZE_M'] = 64
|
||||
else:
|
||||
config['BLOCK_SIZE_M'] = 256
|
||||
return config
|
||||
|
||||
def try_get_optimal_moe_config(
|
||||
w1_shape: Tuple[int, ...],
|
||||
w2_shape: Tuple[int, ...],
|
||||
top_k: int,
|
||||
dtype: Optional[str],
|
||||
M: int,
|
||||
override_config: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
if override_config:
|
||||
config = override_config
|
||||
else:
|
||||
# First try to load optimal config from the file
|
||||
E, _, N = w2_shape
|
||||
configs = get_moe_configs(E, N, dtype)
|
||||
|
||||
if configs:
|
||||
# If an optimal configuration map has been found, look up the
|
||||
# optimal config
|
||||
config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
|
||||
else:
|
||||
# Else use the default config
|
||||
config = get_default_config(M, E, N, w1_shape[2], top_k, dtype)
|
||||
return config
|
||||
|
||||
def moe_align_block_size(
|
||||
topk_ids: torch.Tensor, block_size: int,
|
||||
num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Aligns the token distribution across experts to be compatible with block
|
||||
size for matrix multiplication.
|
||||
|
||||
Parameters:
|
||||
- topk_ids: A tensor of shape [total_tokens, top_k] representing the
|
||||
top-k expert indices for each token.
|
||||
- block_size: The block size used in block matrix multiplication.
|
||||
- num_experts: The total number of experts.
|
||||
|
||||
Returns:
|
||||
- sorted_token_ids: A tensor containing the sorted token indices according
|
||||
to their allocated expert.
|
||||
- expert_ids: A tensor indicating the assigned expert index for each block.
|
||||
- num_tokens_post_padded: The total number of tokens after padding,
|
||||
ensuring divisibility by block_size.
|
||||
|
||||
This function pads the number of tokens that each expert needs to process
|
||||
so that it is divisible by block_size.
|
||||
Padding ensures that during block matrix multiplication, the dimensions
|
||||
align correctly.
|
||||
|
||||
Example:
|
||||
Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
|
||||
block_size = 4, and num_experts = 4:
|
||||
- We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
|
||||
with each expert needing to process 3 tokens.
|
||||
- As block_size is 4, we pad 1 token for each expert.
|
||||
- First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
|
||||
- Then append padding tokens [12, 12, 12, 12] for each block.
|
||||
- After sorting by expert index, we obtain token_ids
|
||||
[3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
|
||||
Tokens 12 are non-existent (padding) and are ignored in
|
||||
the subsequent matrix multiplication.
|
||||
- The padding ensures that the total number of tokens is now divisible
|
||||
by block_size for proper block matrix operations.
|
||||
"""
|
||||
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
|
||||
sorted_ids = torch.empty((max_num_tokens_padded, ),
|
||||
dtype=torch.int32,
|
||||
device=topk_ids.device)
|
||||
sorted_ids.fill_(topk_ids.numel())
|
||||
# max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
|
||||
max_num_m_blocks = topk_ids.numel() + num_experts
|
||||
expert_ids = torch.empty((max_num_m_blocks, ),
|
||||
dtype=torch.int32,
|
||||
device=topk_ids.device)
|
||||
num_tokens_post_pad = torch.empty((1),
|
||||
dtype=torch.int32,
|
||||
device=topk_ids.device)
|
||||
ops.vllm_moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
|
||||
expert_ids, num_tokens_post_pad)
|
||||
return sorted_ids, expert_ids, num_tokens_post_pad
|
||||
|
||||
def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
|
||||
A_scale: Optional[torch.Tensor],
|
||||
B_scale: Optional[torch.Tensor],
|
||||
topk_weights: torch.Tensor, topk_ids: torch.Tensor,
|
||||
sorted_token_ids: torch.Tensor,
|
||||
expert_ids: torch.Tensor,
|
||||
num_tokens_post_padded: torch.Tensor,
|
||||
mul_routed_weight: bool, top_k: int,
|
||||
config: Dict[str, Any], compute_type: torch.dtype,
|
||||
use_fp8: bool) -> None:
|
||||
ops.vllm_invoke_fused_moe_kernel(A, B, C, topk_weights, topk_ids,
|
||||
sorted_token_ids,expert_ids, num_tokens_post_padded,
|
||||
mul_routed_weight, top_k, config['BLOCK_SIZE_M'])
|
||||
|
||||
def fused_experts(hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
inplace: bool = False,
|
||||
override_config: Optional[Dict[str, Any]] = None,
|
||||
use_fp8: bool = False,
|
||||
w1_scale: Optional[torch.Tensor] = None,
|
||||
w2_scale: Optional[torch.Tensor] = None,
|
||||
a1_scale: Optional[torch.Tensor] = None,
|
||||
a2_scale: Optional[torch.Tensor] = None):
|
||||
# Check constraints.
|
||||
assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
|
||||
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
|
||||
assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
|
||||
assert w1.is_contiguous(), "Expert weights1 must be contiguous"
|
||||
assert w2.is_contiguous(), "Expert weights2 must be contiguous"
|
||||
assert hidden_states.dtype in [
|
||||
torch.float32, torch.float16, torch.bfloat16
|
||||
]
|
||||
|
||||
num_tokens, _ = hidden_states.shape
|
||||
E, N, _ = w1.shape
|
||||
# We execute the fused_moe kernel in chunks to circumvent this issue:
|
||||
# https://github.com/vllm-project/vllm/issues/5938
|
||||
M = min(num_tokens, CHUNK_SIZE)
|
||||
|
||||
get_config_func = functools.partial(
|
||||
try_get_optimal_moe_config,
|
||||
w1.shape,
|
||||
w2.shape,
|
||||
topk_ids.shape[1],
|
||||
"float8" if use_fp8 else None,
|
||||
override_config=override_config,
|
||||
)
|
||||
|
||||
config = get_config_func(M)
|
||||
|
||||
intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
|
||||
device=hidden_states.device,
|
||||
dtype=hidden_states.dtype)
|
||||
intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2),
|
||||
device=hidden_states.device,
|
||||
dtype=hidden_states.dtype)
|
||||
intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
|
||||
device=hidden_states.device,
|
||||
dtype=hidden_states.dtype)
|
||||
|
||||
compute_type = (torch.bfloat16
|
||||
if hidden_states.dtype == torch.bfloat16 else torch.float16)
|
||||
|
||||
if inplace:
|
||||
out_hidden_states = hidden_states
|
||||
else:
|
||||
out_hidden_states = torch.empty_like(hidden_states)
|
||||
|
||||
for chunk in range((num_tokens // CHUNK_SIZE) + 1):
|
||||
begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
|
||||
min((chunk + 1) * CHUNK_SIZE,
|
||||
num_tokens))
|
||||
curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
|
||||
tokens_in_chunk, _ = curr_hidden_states.shape
|
||||
|
||||
if tokens_in_chunk == 0:
|
||||
break
|
||||
|
||||
if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
|
||||
# Adjust the intermediate cache size and config for the last
|
||||
# chunk. Note that in most cases we only have one chunk
|
||||
# so the cache size and config are already set correctly and
|
||||
# do not need to be adjusted.
|
||||
intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
|
||||
intermediate_cache2 = intermediate_cache2[:tokens_in_chunk]
|
||||
intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
|
||||
config = get_config_func(tokens_in_chunk)
|
||||
|
||||
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
|
||||
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
|
||||
|
||||
sorted_token_ids, expert_ids, num_tokens_post_padded = (
|
||||
moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E))
|
||||
|
||||
invoke_fused_moe_kernel(curr_hidden_states,
|
||||
w1,
|
||||
intermediate_cache1,
|
||||
a1_scale,
|
||||
w1_scale,
|
||||
curr_topk_weights,
|
||||
curr_topk_ids,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
False,
|
||||
topk_ids.shape[1],
|
||||
config,
|
||||
compute_type=compute_type,
|
||||
use_fp8=use_fp8)
|
||||
|
||||
ops.silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
|
||||
|
||||
invoke_fused_moe_kernel(intermediate_cache2,
|
||||
w2,
|
||||
intermediate_cache3,
|
||||
a2_scale,
|
||||
w2_scale,
|
||||
curr_topk_weights,
|
||||
curr_topk_ids,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
True,
|
||||
1,
|
||||
config,
|
||||
compute_type=compute_type,
|
||||
use_fp8=use_fp8)
|
||||
|
||||
torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
|
||||
dim=1,
|
||||
out=out_hidden_states[begin_chunk_idx:end_chunk_idx])
|
||||
return out_hidden_states
|
||||
|
||||
def fused_moe(
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
gating_output: torch.Tensor,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
inplace: bool = False,
|
||||
override_config: Optional[Dict[str, Any]] = None,
|
||||
use_grouped_topk: bool = False,
|
||||
num_expert_group: Optional[int] = None,
|
||||
topk_group: Optional[int] = None,
|
||||
use_fp8: bool = False,
|
||||
w1_scale: Optional[torch.Tensor] = None,
|
||||
w2_scale: Optional[torch.Tensor] = None,
|
||||
a1_scale: Optional[torch.Tensor] = None,
|
||||
a2_scale: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
This function computes a Mixture of Experts (MoE) layer using two sets of
|
||||
weights, w1 and w2, and top-k gating mechanism.
|
||||
|
||||
Parameters:
|
||||
- hidden_states (torch.Tensor): The input tensor to the MoE layer.
|
||||
- w1 (torch.Tensor): The first set of expert weights.
|
||||
- w2 (torch.Tensor): The second set of expert weights.
|
||||
- gating_output (torch.Tensor): The output of the gating operation
|
||||
(before softmax).
|
||||
- topk (int): The number of top-k experts to select.
|
||||
- renormalize (bool): If True, renormalize the top-k weights to sum to 1.
|
||||
- inplace (bool): If True, perform the operation in-place.
|
||||
Defaults to False.
|
||||
- override_config (Optional[Dict[str, Any]]): Optional override
|
||||
for the kernel configuration.
|
||||
- num_expert_group: Optional[int]: additional parameter for grouped_topk
|
||||
- topk_group: Optional[int]: additional parameter for grouped_topk
|
||||
- use_grouped_topk: If True, use grouped_topk instead of fused_topk
|
||||
note: Deepseekv2 model uses grouped_topk
|
||||
- use_fp8 (bool): If True, use fp8 arithmetic to compute the inner
|
||||
products for w1 and w2. Defaults to False.
|
||||
- w1_scale (Optional[torch.Tensor]): Optional scale to be used for
|
||||
w1.
|
||||
- w2_scale (Optional[torch.Tensor]): Optional scale to be used for
|
||||
w2.
|
||||
|
||||
Returns:
|
||||
- torch.Tensor: The output tensor after applying the MoE layer.
|
||||
"""
|
||||
# Check constraints.
|
||||
assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
|
||||
|
||||
if use_grouped_topk:
|
||||
assert num_expert_group is not None and topk_group is not None
|
||||
topk_weights, topk_ids = grouped_topk(hidden_states, gating_output,
|
||||
topk, renormalize,
|
||||
num_expert_group, topk_group)
|
||||
else:
|
||||
topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
|
||||
renormalize)
|
||||
|
||||
return fused_experts(hidden_states,
|
||||
w1,
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
inplace=inplace,
|
||||
override_config=override_config,
|
||||
use_fp8=use_fp8,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale)
|
||||
2
ixformer_sdk/contrib/transformers/__init__.py
Normal file
2
ixformer_sdk/contrib/transformers/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from .models.bert.modeling_bert import BertForQuestionAnswering
|
||||
from .models.t5.modeling_t5 import T5ForConditionalGeneration
|
||||
@@ -0,0 +1,150 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
|
||||
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# 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.
|
||||
"""BERT model configuration"""
|
||||
|
||||
from collections import OrderedDict
|
||||
from typing import Mapping
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.onnx import OnnxConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class BertConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`BertModel`] or a [`TFBertModel`]. It is used to
|
||||
instantiate a BERT model according to the specified arguments, defining the model architecture. Instantiating a
|
||||
configuration with the defaults will yield a similar configuration to that of the BERT
|
||||
[google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased) architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 30522):
|
||||
Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`BertModel`] or [`TFBertModel`].
|
||||
hidden_size (`int`, *optional*, defaults to 768):
|
||||
Dimensionality of the encoder layers and the pooler layer.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 12):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 12):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
intermediate_size (`int`, *optional*, defaults to 3072):
|
||||
Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
|
||||
hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
|
||||
The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
|
||||
`"relu"`, `"silu"` and `"gelu_new"` are supported.
|
||||
hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
|
||||
The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
|
||||
attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
|
||||
The dropout ratio for the attention probabilities.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 512):
|
||||
The maximum sequence length that this model might ever be used with. Typically set this to something large
|
||||
just in case (e.g., 512 or 1024 or 2048).
|
||||
type_vocab_size (`int`, *optional*, defaults to 2):
|
||||
The vocabulary size of the `token_type_ids` passed when calling [`BertModel`] or [`TFBertModel`].
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
layer_norm_eps (`float`, *optional*, defaults to 1e-12):
|
||||
The epsilon used by the layer normalization layers.
|
||||
position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
|
||||
Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
|
||||
positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
|
||||
[Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
|
||||
For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
|
||||
with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
|
||||
is_decoder (`bool`, *optional*, defaults to `False`):
|
||||
Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
classifier_dropout (`float`, *optional*):
|
||||
The dropout ratio for the classification head.
|
||||
|
||||
Examples:
|
||||
|
||||
```python
|
||||
>>> from transformers import BertConfig, BertModel
|
||||
|
||||
>>> # Initializing a BERT google-bert/bert-base-uncased style configuration
|
||||
>>> configuration = BertConfig()
|
||||
|
||||
>>> # Initializing a model (with random weights) from the google-bert/bert-base-uncased style configuration
|
||||
>>> model = BertModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "bert"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=30522,
|
||||
hidden_size=768,
|
||||
num_hidden_layers=12,
|
||||
num_attention_heads=12,
|
||||
intermediate_size=3072,
|
||||
hidden_act="gelu",
|
||||
hidden_dropout_prob=0.1,
|
||||
attention_probs_dropout_prob=0.1,
|
||||
max_position_embeddings=512,
|
||||
type_vocab_size=2,
|
||||
initializer_range=0.02,
|
||||
layer_norm_eps=1e-12,
|
||||
pad_token_id=0,
|
||||
position_embedding_type="absolute",
|
||||
use_cache=True,
|
||||
classifier_dropout=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(pad_token_id=pad_token_id, **kwargs)
|
||||
|
||||
self.vocab_size = vocab_size
|
||||
self.hidden_size = hidden_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.intermediate_size = intermediate_size
|
||||
self.hidden_dropout_prob = hidden_dropout_prob
|
||||
self.attention_probs_dropout_prob = attention_probs_dropout_prob
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.type_vocab_size = type_vocab_size
|
||||
self.initializer_range = initializer_range
|
||||
self.layer_norm_eps = layer_norm_eps
|
||||
self.position_embedding_type = position_embedding_type
|
||||
self.use_cache = use_cache
|
||||
self.classifier_dropout = classifier_dropout
|
||||
|
||||
|
||||
class BertOnnxConfig(OnnxConfig):
|
||||
@property
|
||||
def inputs(self) -> Mapping[str, Mapping[int, str]]:
|
||||
if self.task == "multiple-choice":
|
||||
dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
|
||||
else:
|
||||
dynamic_axis = {0: "batch", 1: "sequence"}
|
||||
return OrderedDict(
|
||||
[
|
||||
("input_ids", dynamic_axis),
|
||||
("attention_mask", dynamic_axis),
|
||||
("token_type_ids", dynamic_axis),
|
||||
]
|
||||
)
|
||||
2145
ixformer_sdk/contrib/transformers/models/bert/modeling_bert.py
Normal file
2145
ixformer_sdk/contrib/transformers/models/bert/modeling_bert.py
Normal file
File diff suppressed because it is too large
Load Diff
174
ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py
Normal file
174
ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py
Normal file
@@ -0,0 +1,174 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2020, The T5 Authors and HuggingFace Inc.
|
||||
#
|
||||
# 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.
|
||||
""" T5 model configuration"""
|
||||
from typing import Mapping
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.onnx import OnnxSeq2SeqConfigWithPast
|
||||
from transformers.utils import logging
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
T5_PRETRAINED_CONFIG_ARCHIVE_MAP = {
|
||||
"t5-small": "https://huggingface.co/t5-small/resolve/main/config.json",
|
||||
"t5-base": "https://huggingface.co/t5-base/resolve/main/config.json",
|
||||
"t5-large": "https://huggingface.co/t5-large/resolve/main/config.json",
|
||||
"t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json",
|
||||
"t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json",
|
||||
}
|
||||
|
||||
|
||||
class T5Config(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`T5Model`] or a [`TFT5Model`]. It is used to
|
||||
instantiate a T5 model according to the specified arguments, defining the model architecture. Instantiating a
|
||||
configuration with the defaults will yield a similar configuration to that of the T5
|
||||
[t5-small](https://huggingface.co/t5-small) architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
Arguments:
|
||||
vocab_size (`int`, *optional*, defaults to 32128):
|
||||
Vocabulary size of the T5 model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`T5Model`] or [`TFT5Model`].
|
||||
d_model (`int`, *optional*, defaults to 512):
|
||||
Size of the encoder layers and the pooler layer.
|
||||
d_kv (`int`, *optional*, defaults to 64):
|
||||
Size of the key, query, value projections per attention head. The `inner_dim` of the projection layer will
|
||||
be defined as `num_heads * d_kv`.
|
||||
d_ff (`int`, *optional*, defaults to 2048):
|
||||
Size of the intermediate feed forward layer in each `T5Block`.
|
||||
num_layers (`int`, *optional*, defaults to 6):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
num_decoder_layers (`int`, *optional*):
|
||||
Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set.
|
||||
num_heads (`int`, *optional*, defaults to 8):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
relative_attention_num_buckets (`int`, *optional*, defaults to 32):
|
||||
The number of buckets to use for each attention layer.
|
||||
relative_attention_max_distance (`int`, *optional*, defaults to 128):
|
||||
The maximum distance of the longer sequences for the bucket separation.
|
||||
dropout_rate (`float`, *optional*, defaults to 0.1):
|
||||
The ratio for all dropout layers.
|
||||
layer_norm_eps (`float`, *optional*, defaults to 1e-6):
|
||||
The epsilon used by the layer normalization layers.
|
||||
initializer_factor (`float`, *optional*, defaults to 1):
|
||||
A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
|
||||
testing).
|
||||
feed_forward_proj (`string`, *optional*, defaults to `"relu"`):
|
||||
Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. T5v1.1 uses the
|
||||
`"gated-gelu"` feed forward projection. Original T5 uses `"relu"`.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models).
|
||||
"""
|
||||
model_type = "t5"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
attribute_map = {
|
||||
"hidden_size": "d_model",
|
||||
"num_attention_heads": "num_heads",
|
||||
"num_hidden_layers": "num_layers",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=32128,
|
||||
d_model=512,
|
||||
d_kv=64,
|
||||
d_ff=2048,
|
||||
num_layers=6,
|
||||
num_decoder_layers=None,
|
||||
num_heads=8,
|
||||
relative_attention_num_buckets=32,
|
||||
relative_attention_max_distance=128,
|
||||
dropout_rate=0.1,
|
||||
layer_norm_epsilon=1e-6,
|
||||
initializer_factor=1.0,
|
||||
feed_forward_proj="relu",
|
||||
is_encoder_decoder=True,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
eos_token_id=1,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.d_model = d_model
|
||||
self.d_kv = d_kv
|
||||
self.d_ff = d_ff
|
||||
self.num_layers = num_layers
|
||||
self.num_decoder_layers = (
|
||||
num_decoder_layers if num_decoder_layers is not None else self.num_layers
|
||||
) # default = symmetry
|
||||
self.num_heads = num_heads
|
||||
self.relative_attention_num_buckets = relative_attention_num_buckets
|
||||
self.relative_attention_max_distance = relative_attention_max_distance
|
||||
self.dropout_rate = dropout_rate
|
||||
self.layer_norm_epsilon = layer_norm_epsilon
|
||||
self.initializer_factor = initializer_factor
|
||||
self.feed_forward_proj = feed_forward_proj
|
||||
self.use_cache = use_cache
|
||||
|
||||
act_info = self.feed_forward_proj.split("-")
|
||||
self.dense_act_fn = act_info[-1]
|
||||
self.is_gated_act = act_info[0] == "gated"
|
||||
|
||||
if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
|
||||
raise ValueError(
|
||||
f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
|
||||
"Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
|
||||
"'gated-gelu' or 'relu'"
|
||||
)
|
||||
|
||||
# for backwards compatibility
|
||||
if feed_forward_proj == "gated-gelu":
|
||||
self.dense_act_fn = "gelu_new"
|
||||
|
||||
super().__init__(
|
||||
pad_token_id=pad_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
is_encoder_decoder=is_encoder_decoder,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class T5OnnxConfig(OnnxSeq2SeqConfigWithPast):
|
||||
@property
|
||||
def inputs(self) -> Mapping[str, Mapping[int, str]]:
|
||||
common_inputs = {
|
||||
"input_ids": {0: "batch", 1: "encoder_sequence"},
|
||||
"attention_mask": {0: "batch", 1: "encoder_sequence"},
|
||||
}
|
||||
if self.use_past:
|
||||
common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence"
|
||||
common_inputs["decoder_input_ids"] = {0: "batch"}
|
||||
common_inputs["decoder_attention_mask"] = {
|
||||
0: "batch",
|
||||
1: "past_decoder_sequence + sequence",
|
||||
}
|
||||
else:
|
||||
common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"}
|
||||
common_inputs["decoder_attention_mask"] = {
|
||||
0: "batch",
|
||||
1: "decoder_sequence",
|
||||
}
|
||||
|
||||
if self.use_past:
|
||||
self.fill_with_past_key_values_(common_inputs, direction="inputs")
|
||||
|
||||
return common_inputs
|
||||
|
||||
@property
|
||||
def default_onnx_opset(self) -> int:
|
||||
return 13
|
||||
@@ -0,0 +1,315 @@
|
||||
import torch
|
||||
from transformers.activations import NewGELUActivation
|
||||
|
||||
import ixformer
|
||||
|
||||
|
||||
def self_attention_forward(
|
||||
self,
|
||||
hidden_states,
|
||||
attention_mask=None,
|
||||
position_bias=None,
|
||||
layer_head_mask=None,
|
||||
past_key_value=None,
|
||||
use_cache=False,
|
||||
output_attentions=False,
|
||||
):
|
||||
assert output_attentions is False
|
||||
assert layer_head_mask is None
|
||||
|
||||
normed_hidden_states = self.layer_norm(hidden_states)
|
||||
|
||||
if not hasattr(self, "qkv_weight"):
|
||||
self.qkv_weight = torch.cat(
|
||||
[
|
||||
self.SelfAttention.q.weight,
|
||||
self.SelfAttention.k.weight,
|
||||
self.SelfAttention.v.weight,
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
self.qkv_bias = None
|
||||
|
||||
del self.SelfAttention.q.weight
|
||||
del self.SelfAttention.k.weight
|
||||
del self.SelfAttention.v.weight
|
||||
|
||||
batch_size, seq_length = hidden_states.shape[:2]
|
||||
real_seq_length = seq_length
|
||||
if past_key_value is not None:
|
||||
if len(past_key_value) != 2:
|
||||
raise ValueError(
|
||||
f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
|
||||
)
|
||||
real_seq_length += past_key_value[0].shape[2]
|
||||
key_length = real_seq_length
|
||||
|
||||
def unshape(states):
|
||||
"""reshape"""
|
||||
return (
|
||||
states.transpose(1, 2)
|
||||
.contiguous()
|
||||
.view(batch_size, -1, self.SelfAttention.inner_dim)
|
||||
)
|
||||
|
||||
qkv = ixformer.functions.linear(
|
||||
normed_hidden_states, self.qkv_weight, self.qkv_bias
|
||||
)
|
||||
|
||||
if past_key_value is not None:
|
||||
pask_key, past_value = past_key_value
|
||||
(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
) = ixformer.functions.t5_split_qkv_update_kv_cache(
|
||||
qkv,
|
||||
pask_key,
|
||||
past_value,
|
||||
self.SelfAttention.n_heads,
|
||||
self.SelfAttention.key_value_proj_dim,
|
||||
)
|
||||
else:
|
||||
query_states, key_states, value_states = ixformer.functions.t5_split_qkv(
|
||||
qkv, self.SelfAttention.n_heads, self.SelfAttention.key_value_proj_dim
|
||||
)
|
||||
|
||||
if position_bias is None:
|
||||
if not self.SelfAttention.has_relative_attention_bias:
|
||||
position_bias = torch.zeros(
|
||||
(1, self.SelfAttention.n_heads, real_seq_length, key_length),
|
||||
device=query_states.device,
|
||||
dtype=query_states.dtype,
|
||||
)
|
||||
else:
|
||||
position_bias = self.SelfAttention.compute_bias(
|
||||
real_seq_length, key_length, device=query_states.device
|
||||
)
|
||||
|
||||
# if key and values are already calculated
|
||||
# we want only the last query position bias
|
||||
if past_key_value is not None:
|
||||
position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
|
||||
|
||||
if attention_mask is not None:
|
||||
# (batch_size, n_heads, seq_length, key_length)
|
||||
position_bias = position_bias + attention_mask
|
||||
|
||||
if self.SelfAttention.pruned_heads:
|
||||
mask = torch.ones(position_bias.shape[1])
|
||||
mask[list(self.pruned_heads)] = 0
|
||||
position_bias_masked = position_bias[:, mask.bool()]
|
||||
else:
|
||||
position_bias_masked = position_bias
|
||||
|
||||
attn_output = ixformer.functions.ixinfer_flash_attn_pad(
|
||||
query_states.contiguous(),
|
||||
key_states.contiguous(),
|
||||
value_states.contiguous(),
|
||||
mask=position_bias_masked.float().contiguous(),
|
||||
atten_scale=1,
|
||||
)
|
||||
attn_output = unshape(attn_output)
|
||||
|
||||
attn_output = self.SelfAttention.o(attn_output)
|
||||
|
||||
present_key_value_state = (
|
||||
(key_states, value_states)
|
||||
if (self.SelfAttention.is_decoder and use_cache)
|
||||
else None
|
||||
)
|
||||
outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
|
||||
|
||||
if output_attentions:
|
||||
outputs = outputs + (None,)
|
||||
hidden_states = attn_output + hidden_states
|
||||
outputs = (hidden_states,) + outputs[1:]
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def cross_attention_forward(
|
||||
self,
|
||||
hidden_states,
|
||||
key_value_states,
|
||||
attention_mask=None,
|
||||
position_bias=None,
|
||||
layer_head_mask=None,
|
||||
past_key_value=None,
|
||||
use_cache=False,
|
||||
query_length=None,
|
||||
output_attentions=False,
|
||||
):
|
||||
|
||||
assert output_attentions is False
|
||||
assert layer_head_mask is None
|
||||
|
||||
def unshape(states):
|
||||
"""reshape"""
|
||||
return (
|
||||
states.transpose(1, 2)
|
||||
.contiguous()
|
||||
.view(batch_size, -1, self.EncDecAttention.inner_dim)
|
||||
)
|
||||
|
||||
normed_hidden_states = self.layer_norm(hidden_states)
|
||||
|
||||
# cross attn need key_value_states
|
||||
assert key_value_states is not None
|
||||
batch_size, seq_length = hidden_states.shape[:2]
|
||||
real_seq_length = seq_length
|
||||
|
||||
if past_key_value is not None:
|
||||
if len(past_key_value) != 2:
|
||||
raise ValueError(
|
||||
f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
|
||||
)
|
||||
real_seq_length += (
|
||||
past_key_value[0].shape[2] if query_length is None else query_length
|
||||
)
|
||||
|
||||
key_length = (
|
||||
real_seq_length if key_value_states is None else key_value_states.shape[1]
|
||||
)
|
||||
head_num, head_dim = (
|
||||
self.EncDecAttention.n_heads,
|
||||
self.EncDecAttention.key_value_proj_dim,
|
||||
)
|
||||
|
||||
query_states = (
|
||||
self.EncDecAttention.q(normed_hidden_states)
|
||||
.view(batch_size, seq_length, head_num, head_dim)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
if past_key_value is not None:
|
||||
if past_key_value[0].shape[2] != key_value_states.shape[1]:
|
||||
# checking that the `sequence_length` of the `past_key_value` is the same as
|
||||
# the provided `key_value_states` to support prefix tuning
|
||||
# cross-attn
|
||||
# (batch_size, n_heads, seq_length, dim_per_head)
|
||||
key_states = (
|
||||
self.EncDecAttention.k(key_value_states)
|
||||
.view(batch_size, key_length, head_num, head_dim)
|
||||
.transpose(1, 2)
|
||||
)
|
||||
value_states = (
|
||||
self.EncDecAttention.v(key_value_states)
|
||||
.view(batch_size, key_length, head_num, head_dim)
|
||||
.transpose(1, 2)
|
||||
)
|
||||
else:
|
||||
# cross-attn
|
||||
key_states = past_key_value[0]
|
||||
value_states = past_key_value[1]
|
||||
else:
|
||||
key_states = (
|
||||
self.EncDecAttention.k(key_value_states)
|
||||
.view(batch_size, key_length, head_num, head_dim)
|
||||
.transpose(1, 2)
|
||||
)
|
||||
value_states = (
|
||||
self.EncDecAttention.v(key_value_states)
|
||||
.view(batch_size, key_length, head_num, head_dim)
|
||||
.transpose(1, 2)
|
||||
)
|
||||
|
||||
if not query_states.is_contiguous():
|
||||
query_states = query_states.contiguous()
|
||||
|
||||
# TODO: fix this bug
|
||||
if not value_states.is_contiguous():
|
||||
new_value_states = query_states.new_empty(value_states.shape)
|
||||
new_value_states.copy_(value_states)
|
||||
value_states = new_value_states
|
||||
if not key_states.is_contiguous():
|
||||
numel = torch.numel(key_states)
|
||||
new_key_states = query_states.new_empty([numel * 2])[:numel].view(
|
||||
*list(key_states.shape)
|
||||
)
|
||||
new_key_states.copy_(key_states)
|
||||
key_states = new_key_states
|
||||
|
||||
if position_bias is None:
|
||||
if not self.EncDecAttention.has_relative_attention_bias:
|
||||
position_bias = torch.zeros(
|
||||
(1, self.EncDecAttention.n_heads, real_seq_length, key_length),
|
||||
device=query_states.device,
|
||||
dtype=query_states.dtype,
|
||||
)
|
||||
else:
|
||||
position_bias = self.EncDecAttention.compute_bias(
|
||||
real_seq_length, key_length, device=query_states.device
|
||||
)
|
||||
|
||||
# if key and values are already calculated
|
||||
# we want only the last query position bias
|
||||
if past_key_value is not None:
|
||||
position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
|
||||
|
||||
if attention_mask is not None:
|
||||
# (batch_size, n_heads, seq_length, key_length)
|
||||
position_bias = position_bias + attention_mask
|
||||
|
||||
if self.EncDecAttention.pruned_heads:
|
||||
mask = torch.ones(position_bias.shape[1])
|
||||
mask[list(self.pruned_heads)] = 0
|
||||
position_bias_masked = position_bias[:, mask.bool()]
|
||||
else:
|
||||
position_bias_masked = position_bias
|
||||
|
||||
attn_output = ixformer.functions.ixinfer_flash_attn_pad(
|
||||
query_states,
|
||||
key_states.contiguous(),
|
||||
value_states.contiguous(),
|
||||
mask=position_bias_masked.float().contiguous(),
|
||||
atten_scale=1,
|
||||
)
|
||||
attn_output = unshape(attn_output)
|
||||
|
||||
attn_output = self.EncDecAttention.o(attn_output)
|
||||
|
||||
present_key_value_state = (
|
||||
(key_states, value_states)
|
||||
if (self.EncDecAttention.is_decoder and use_cache)
|
||||
else None
|
||||
)
|
||||
outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
|
||||
|
||||
if output_attentions:
|
||||
outputs = outputs + (None,)
|
||||
hidden_states = attn_output + hidden_states
|
||||
outputs = (hidden_states,) + outputs[1:]
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def dense_gated_act_dense_forward(self, hidden_states):
|
||||
if isinstance(self.act, NewGELUActivation):
|
||||
if not hasattr(self, "wi"):
|
||||
self.wi = torch.cat([self.wi_1.weight, self.wi_0.weight], dim=0)
|
||||
del self.wi_1
|
||||
del self.wi_0
|
||||
hidden_states = ixformer.functions.linear(hidden_states, self.wi, None)
|
||||
hidden_states = ixformer.functions.gelu_and_mul(hidden_states)
|
||||
hidden_states = ixformer.functions.linear(hidden_states, self.wo.weight, None)
|
||||
else:
|
||||
hidden_gelu = self.act(self.wi_0(hidden_states))
|
||||
hidden_linear = self.wi_1(hidden_states)
|
||||
hidden_states = hidden_gelu * hidden_linear
|
||||
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
# To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32.
|
||||
# See https://github.com/huggingface/transformers/issues/20287
|
||||
# we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None``
|
||||
if (
|
||||
isinstance(self.wo.weight, torch.Tensor)
|
||||
and hidden_states.dtype != self.wo.weight.dtype
|
||||
and self.wo.weight.dtype != torch.int8
|
||||
):
|
||||
hidden_states = hidden_states.to(self.wo.weight.dtype)
|
||||
|
||||
hidden_states = self.wo(hidden_states)
|
||||
return hidden_states
|
||||
2644
ixformer_sdk/contrib/transformers/models/t5/modeling_t5.py
Normal file
2644
ixformer_sdk/contrib/transformers/models/t5/modeling_t5.py
Normal file
File diff suppressed because it is too large
Load Diff
0
ixformer_sdk/contrib/vllm/__init__.py
Normal file
0
ixformer_sdk/contrib/vllm/__init__.py
Normal file
30
ixformer_sdk/contrib/vllm/layers/__init__.py
Normal file
30
ixformer_sdk/contrib/vllm/layers/__init__.py
Normal file
@@ -0,0 +1,30 @@
|
||||
from .llama import forward_smoothquant
|
||||
from .mixtral import mixtral_decoder_layer_forward
|
||||
|
||||
SUPPORT_REPLACE_METHOD = {
|
||||
"llama": forward_smoothquant,
|
||||
}
|
||||
|
||||
SUPPORT_REPLACE_LAYER = {
|
||||
"llama": None,
|
||||
}
|
||||
|
||||
|
||||
def get_replace_forward(name: str):
|
||||
try:
|
||||
method = SUPPORT_REPLACE_METHOD[name]
|
||||
except:
|
||||
raise ValueError(
|
||||
f"Only support replace names: {SUPPORT_REPLACE_METHOD.keys()}, but got {name}"
|
||||
)
|
||||
return method
|
||||
|
||||
|
||||
def get_replace_layer(name: str):
|
||||
try:
|
||||
layer = SUPPORT_REPLACE_LAYER[name]
|
||||
except:
|
||||
raise ValueError(
|
||||
f"Only support replace names: {SUPPORT_REPLACE_LAYER.keys()}, but got {name}"
|
||||
)
|
||||
return layer
|
||||
100
ixformer_sdk/contrib/vllm/layers/llama.py
Normal file
100
ixformer_sdk/contrib/vllm/layers/llama.py
Normal file
@@ -0,0 +1,100 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from transformers import LlamaConfig
|
||||
|
||||
from vllm.attention import AttentionMetadata
|
||||
from vllm.config import CacheConfig
|
||||
from vllm.distributed import tensor_model_parallel_all_reduce
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
from vllm.model_executor.models.llama import LlamaDecoderLayer as VllmLlamaDecoderLayer
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
# from ..overlap_comm import DecoderLayerOverlapComm, get_overlap_linear_method
|
||||
|
||||
|
||||
# This method is needed for support smoothquant no overlap forward
|
||||
def forward_smoothquant(
|
||||
input_ids: Optional[torch.Tensor],
|
||||
positions: torch.Tensor,
|
||||
kv_caches: List[torch.Tensor],
|
||||
attn_metadata: AttentionMetadata,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
self = None, # will be set by partial
|
||||
) -> torch.Tensor:
|
||||
dtype = self.dtype
|
||||
|
||||
def forward_smoothquant_mlp(self,x,scales):
|
||||
# gate_up_proj
|
||||
# Int8 Matrix multiply.
|
||||
bias = self.gate_up_proj.bias if not self.gate_up_proj.skip_bias_add else None
|
||||
gate_up = ops.w8a8(x, self.gate_up_proj.weight, scales, self.gate_up_proj.weight_scales, dtype)
|
||||
if bias:
|
||||
gate_up += bias
|
||||
|
||||
# act_fun
|
||||
x, scales = ops.silu_and_mul_smoothquant(gate_up, self.down_proj.smooth_scales)
|
||||
|
||||
# down_proj
|
||||
output_parallel = ops.w8a8(x, self.down_proj.weight, scales, self.down_proj.weight_scales, dtype)
|
||||
if self.down_proj.reduce_results and self.down_proj.tp_size > 1:
|
||||
output = tensor_model_parallel_all_reduce(output_parallel)
|
||||
else:
|
||||
output = output_parallel
|
||||
|
||||
if not self.down_proj.skip_bias_add:
|
||||
output = output + self.down_proj.bias if self.down_proj.bias is not None else output
|
||||
|
||||
return output
|
||||
|
||||
def forward_smoothquant_attn(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
scales: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# qkv proj
|
||||
bias = self.qkv_proj.bias if not self.qkv_proj.skip_bias_add else None
|
||||
|
||||
qkv = ops.w8a8(hidden_states, self.qkv_proj.weight, scales, self.qkv_proj.weight_scales, dtype)
|
||||
if bias:
|
||||
qkv += bias
|
||||
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
|
||||
output, _ = self.o_proj(attn_output) # TODO
|
||||
return output
|
||||
|
||||
if inputs_embeds is not None:
|
||||
hidden_states = inputs_embeds
|
||||
else:
|
||||
hidden_states = self.get_input_embeddings(input_ids)
|
||||
residual = None
|
||||
for i in range(len(self.layers)):
|
||||
layer = self.layers[i]
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
hidden_states, scales = ops.rms_norm_smoothquant(hidden_states,layer.input_layernorm.weight,layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
|
||||
else:
|
||||
hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.input_layernorm.weight, layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
|
||||
|
||||
hidden_states = forward_smoothquant_attn(
|
||||
layer.self_attn,
|
||||
positions=positions,
|
||||
hidden_states=hidden_states,
|
||||
kv_cache=kv_caches[i],
|
||||
attn_metadata=attn_metadata,
|
||||
scales=scales,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.post_attention_layernorm.weight, layer.post_attention_layernorm.variance_epsilon, layer.mlp.gate_up_proj.smooth_scales)
|
||||
|
||||
hidden_states = forward_smoothquant_mlp(layer.mlp, hidden_states, scales)
|
||||
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
return hidden_states
|
||||
|
||||
331
ixformer_sdk/contrib/vllm/layers/mixtral.py
Normal file
331
ixformer_sdk/contrib/vllm/layers/mixtral.py
Normal file
@@ -0,0 +1,331 @@
|
||||
import functools
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import ixformer.inference.functions as ixf
|
||||
import torch
|
||||
|
||||
|
||||
def mixtral_decoder_layer_forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata,
|
||||
residual: Optional[torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
if self.use_int_w8a8:
|
||||
return w8a8_forward(
|
||||
self, positions, hidden_states, kv_cache, attn_metadata, residual
|
||||
)
|
||||
else:
|
||||
return original_forward(
|
||||
self, positions, hidden_states, kv_cache, attn_metadata, residual
|
||||
)
|
||||
|
||||
|
||||
def original_forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata,
|
||||
residual: Optional[torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
# Self Attention
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
else:
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
hidden_states = self.self_attn(
|
||||
positions=positions,
|
||||
hidden_states=hidden_states,
|
||||
kv_cache=kv_cache,
|
||||
attn_metadata=attn_metadata,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||
hidden_states = self.block_sparse_moe(hidden_states)
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
def dynamic_scaled_int8_quant(x):
|
||||
m, k = x.shape
|
||||
i8_x = x.new_empty([m, k], dtype=torch.int8, device="cuda")
|
||||
i8_scales = torch.empty([m], dtype=torch.float32, device="cuda")
|
||||
ixf.dynamic_scaled_int8_quant(i8_x, x, i8_scales)
|
||||
return i8_x, i8_scales
|
||||
|
||||
|
||||
def dynamic_w8a8(x, i8_weight, weight_scale):
|
||||
i8_x, i8_scale = dynamic_scaled_int8_quant(x)
|
||||
m, k = x.shape
|
||||
k, n = i8_weight.shape
|
||||
output = x.new_empty([m, n], dtype=x.dtype, device="cuda")
|
||||
ixf.w8a8(
|
||||
i8_x,
|
||||
i8_weight.transpose(0, 1),
|
||||
i8_scale,
|
||||
weight_scale,
|
||||
output=output,
|
||||
out_dtype=x.dtype,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def fused_rms_norm_quant_linear(
|
||||
self,
|
||||
hidden_states,
|
||||
ln_weight,
|
||||
eps,
|
||||
linear_weight,
|
||||
linear_weight_scale,
|
||||
residual=None,
|
||||
):
|
||||
# lower rouge
|
||||
# if residual is None:
|
||||
# residual = hidden_states
|
||||
# i8_hidden_states, _, i8_scales = ixf.residual_rms_norm_dynamic_int8(
|
||||
# input=hidden_states,
|
||||
# weight=ln_weight,
|
||||
# residual=None,
|
||||
# eps=eps,
|
||||
# )
|
||||
# else:
|
||||
# i8_hidden_states, residual, i8_scales = ixf.residual_rms_norm_dynamic_int8(
|
||||
# input=hidden_states,
|
||||
# weight=ln_weight,
|
||||
# residual=residual,
|
||||
# eps=eps,
|
||||
# )
|
||||
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
else:
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
i8_hidden_states, i8_scales = dynamic_scaled_int8_quant(hidden_states)
|
||||
|
||||
qkv = hidden_states.new_empty(hidden_states.shape[0], linear_weight.shape[1])
|
||||
ixf.w8a8(
|
||||
i8_hidden_states,
|
||||
linear_weight.transpose(0, 1),
|
||||
i8_scales,
|
||||
linear_weight_scale,
|
||||
output=qkv,
|
||||
out_dtype=hidden_states.dtype,
|
||||
)
|
||||
return qkv, residual
|
||||
|
||||
|
||||
def attention(qkv, positions, kv_cache, attn_metadata, self_attn):
|
||||
q, k, v = qkv.split(
|
||||
[self_attn.q_size, self_attn.kv_size, self_attn.kv_size], dim=-1
|
||||
)
|
||||
q, k = self_attn.rotary_emb(positions, q, k)
|
||||
attn_output = self_attn.attn(q, k, v, kv_cache, attn_metadata)
|
||||
return attn_output
|
||||
|
||||
|
||||
def fused_rms_norm_attention(
|
||||
self,
|
||||
hidden_states,
|
||||
ln_weight,
|
||||
eps,
|
||||
positions,
|
||||
kv_cache,
|
||||
attn_metadata,
|
||||
self_attn,
|
||||
residual=None,
|
||||
):
|
||||
hidden_states, residual = fused_rms_norm_quant_linear(
|
||||
self,
|
||||
hidden_states,
|
||||
ln_weight,
|
||||
eps,
|
||||
self_attn.qkv_proj.weight,
|
||||
self_attn.qkv_proj.weight_scale,
|
||||
residual,
|
||||
)
|
||||
hidden_states = attention(
|
||||
hidden_states, positions, kv_cache, attn_metadata, self_attn
|
||||
)
|
||||
|
||||
hidden_states = dynamic_w8a8(
|
||||
hidden_states, self_attn.o_proj.weight, self_attn.o_proj.weight_scale
|
||||
)
|
||||
# hidden_states,_ = self_attn.o_proj(hidden_states) # quant+linear+allreduce
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
def w8a8_forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata,
|
||||
residual: Optional[torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
|
||||
# qkv,_ = self.self_attn.qkv_proj(hidden_states)
|
||||
hidden_states, residual = fused_rms_norm_attention(
|
||||
self,
|
||||
hidden_states,
|
||||
self.input_layernorm.weight,
|
||||
self.input_layernorm.variance_epsilon,
|
||||
positions,
|
||||
kv_cache,
|
||||
attn_metadata,
|
||||
self.self_attn,
|
||||
residual,
|
||||
)
|
||||
|
||||
# allreduce
|
||||
tp_size = self.block_sparse_moe.experts.tp_size
|
||||
if tp_size > 1:
|
||||
from vllm.distributed import tensor_model_parallel_all_reduce
|
||||
|
||||
hidden_states = tensor_model_parallel_all_reduce(hidden_states)
|
||||
|
||||
# rms norm
|
||||
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||
# moe
|
||||
hidden_states = fused_moe(
|
||||
hidden_states,
|
||||
self.block_sparse_moe.gate.weight,
|
||||
top_k=self.block_sparse_moe.experts.top_k,
|
||||
w1=self.block_sparse_moe.experts.w13_weight,
|
||||
w2=self.block_sparse_moe.experts.w2_weight,
|
||||
w1_scale=self.block_sparse_moe.experts.w13_weight_scale,
|
||||
w2_scale=self.block_sparse_moe.experts.w2_weight_scale,
|
||||
)
|
||||
|
||||
# allreduce
|
||||
if tp_size > 1:
|
||||
from vllm.distributed import tensor_model_parallel_all_reduce
|
||||
|
||||
hidden_states = tensor_model_parallel_all_reduce(hidden_states)
|
||||
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
def fused_experts(hidden_states, router_logits, top_k, w1, w2, w1_scale, w2_scale):
|
||||
|
||||
"""
|
||||
Args:
|
||||
hidden_states: (num_tokens, k) dtype
|
||||
router_logits: (num_tokens, num_experts) torch.float32
|
||||
top_k int
|
||||
w1: (num_experts, 2n, k) torch.int8
|
||||
w2: (num_experts, k, n) torch.int8
|
||||
w1_scale: (num_experts, 2n) torch.float32
|
||||
w2_scale: (num_experts, k) torch.float32
|
||||
Returns
|
||||
final_hidden_states: (num_tokens, k) dtype
|
||||
"""
|
||||
|
||||
# topk_weight: (num_tokens, top_k) torch.float32
|
||||
# topk_ids: (num_tokens, top_k) torch.int32
|
||||
topk_weight, topk_ids = ixf.moe_topk_softmax(
|
||||
gating_output=router_logits,
|
||||
topk=top_k,
|
||||
renormalize=True,
|
||||
)
|
||||
|
||||
dtype = hidden_states.dtype
|
||||
num_tokens, num_experts = router_logits.shape
|
||||
expand_tokens = num_tokens * top_k
|
||||
|
||||
(
|
||||
src_to_dst,
|
||||
sorted_token_ids,
|
||||
expert_sizes_gpu,
|
||||
expert_sizes_cpu,
|
||||
) = ixf.moe_compute_token_index(
|
||||
topk_ids=topk_ids,
|
||||
num_experts=num_experts,
|
||||
)
|
||||
expert_sizes_cpu = expert_sizes_gpu.cpu()
|
||||
|
||||
# expand + reorder + quant
|
||||
# i8_hidden_states: (expand_tokens, k) torch.int8
|
||||
i8_hidden_states, a_scale = ixf.moe_expand_input_dynamic_scaled_int8(
|
||||
hidden_states=hidden_states,
|
||||
dst_to_src=sorted_token_ids,
|
||||
dst_tokens=expand_tokens,
|
||||
topk=top_k,
|
||||
src_to_dst=src_to_dst,
|
||||
topk_ids=None, # use smooth quant
|
||||
smooth_scales=None, # use smooth quant
|
||||
)
|
||||
|
||||
# w8a8 group gemm 1
|
||||
# pt_output_1: (expand_tokens, 2n) dtype
|
||||
pt_output_1 = ixf.moe_w8a8_group_gemm(
|
||||
input=i8_hidden_states,
|
||||
weight=w1,
|
||||
i_scales=a_scale,
|
||||
w_scales=w1_scale,
|
||||
output_dtype=dtype,
|
||||
tokens_per_experts=expert_sizes_cpu,
|
||||
dst_to_src=None,
|
||||
format="TN",
|
||||
)
|
||||
|
||||
# act + quant
|
||||
# pt_output_2: (expand_tokens, n) torch.int8
|
||||
pt_output_2, a2_scale = ixf.activation_dynamic_scaled_int8(
|
||||
input=pt_output_1,
|
||||
bias=None, # add gemm bias
|
||||
smooth_scales=None, # use smooth quant
|
||||
dst_to_src=sorted_token_ids,
|
||||
topk_ids=None, # add gemm bias or use smooth quant
|
||||
act_type="swiglu",
|
||||
)
|
||||
|
||||
# w8a8 group gemm 2 + reorder
|
||||
# pt_output_3: (expand_tokens, k) dtype
|
||||
pt_output_3 = ixf.moe_w8a8_group_gemm(
|
||||
input=pt_output_2,
|
||||
weight=w2,
|
||||
i_scales=a2_scale,
|
||||
w_scales=w2_scale,
|
||||
output_dtype=dtype,
|
||||
tokens_per_experts=expert_sizes_cpu,
|
||||
dst_to_src=sorted_token_ids,
|
||||
format="TN",
|
||||
)
|
||||
|
||||
# mul + reduce_sum
|
||||
# final_hidden_states: (num_tokens, k)
|
||||
final_hidden_states = ixf.moe_output_reduce_sum(
|
||||
input=pt_output_3.view(num_tokens, top_k, -1),
|
||||
topk_weight=topk_weight,
|
||||
)
|
||||
|
||||
return final_hidden_states
|
||||
|
||||
|
||||
def fused_moe(hidden_states, gate_weight, top_k, w1, w2, w1_scale, w2_scale):
|
||||
orig_shape = hidden_states.shape
|
||||
hidden_size = hidden_states.shape[-1]
|
||||
|
||||
hidden_states = hidden_states.view(-1, hidden_size)
|
||||
|
||||
# router_logits: (num_tokens, n_experts)
|
||||
# gate_weight: fp16
|
||||
router_logits = ixf.linear(hidden_states, gate_weight)
|
||||
router_logits = router_logits.to(torch.float32)
|
||||
|
||||
final_hidden_states = fused_experts(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
top_k,
|
||||
w1,
|
||||
w2,
|
||||
w1_scale,
|
||||
w2_scale,
|
||||
)
|
||||
|
||||
return final_hidden_states.view(orig_shape)
|
||||
14
ixformer_sdk/contrib/vllm/quantize/__init__.py
Normal file
14
ixformer_sdk/contrib/vllm/quantize/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
from .smoothquant import smoothquant_prepare_quantize,smoothquant_export_quantized_weights
|
||||
from .w8a16 import w8a16_prepare_quantize,w8a16_export_quantized_weights
|
||||
|
||||
SUPPORT_METHOD = {
|
||||
"smoothquant": [smoothquant_prepare_quantize,smoothquant_export_quantized_weights],
|
||||
"w8a16": [w8a16_prepare_quantize,w8a16_export_quantized_weights],
|
||||
}
|
||||
|
||||
def get_quantize_method(method_name:str):
|
||||
try:
|
||||
method = SUPPORT_METHOD[method_name]
|
||||
except:
|
||||
raise ValueError(f"Only support quantization methods: {SUPPORT_METHOD.keys()}, but got {method_name}")
|
||||
return method
|
||||
407
ixformer_sdk/contrib/vllm/quantize/smoothquant.py
Normal file
407
ixformer_sdk/contrib/vllm/quantize/smoothquant.py
Normal file
@@ -0,0 +1,407 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def smoothquant_prepare_quantize(self, quant_params={}):
|
||||
model = self.model_runner.model
|
||||
|
||||
def update_act_scales(act_scales, x):
|
||||
# 动态统计每次输入的最大值
|
||||
hidden_dim = x.shape[-1]
|
||||
x = x.view(-1, hidden_dim).abs().detach()
|
||||
# [k]
|
||||
comming_max = torch.max(x, dim=0, keepdim=True)[0].float()
|
||||
|
||||
if act_scales is None:
|
||||
act_scales = comming_max
|
||||
else:
|
||||
act_scales = torch.max(act_scales, comming_max)
|
||||
return act_scales
|
||||
|
||||
from functools import partial
|
||||
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
def new_forward(input_, m, raw_forward):
|
||||
if not hasattr(m, "act_scales"):
|
||||
m.act_scales = None
|
||||
m.act_scales = update_act_scales(m.act_scales, input_)
|
||||
return raw_forward(input_)
|
||||
|
||||
for name, m in model.named_modules():
|
||||
if (
|
||||
isinstance(m, QKVParallelLinear)
|
||||
or isinstance(m, RowParallelLinear)
|
||||
or isinstance(m, MergedColumnParallelLinear)
|
||||
or isinstance(m, ColumnParallelLinear)
|
||||
):
|
||||
m.forward = partial(new_forward, m=m, raw_forward=m.forward)
|
||||
|
||||
def smoothquant_export_quantized_weights(self, save_path, quant_params={}):
|
||||
gb_per_file = quant_params.get("filesize_limit", None)
|
||||
smooth_alpha = quant_params.get("smooth_alpha", 0.5)
|
||||
dynamic_quant_type = quant_params.get("dynamic_quant_type", "gpu")
|
||||
assert dynamic_quant_type in ["gpu","cpu","kernel"]
|
||||
if self.rank == 0:
|
||||
print(f"set smooth_alpha={smooth_alpha}")
|
||||
print(f"use quantize weight type: {dynamic_quant_type}")
|
||||
|
||||
import ixformer._C as ops
|
||||
def per_token_quant_8bit(weight):
|
||||
# weight: [m,k]
|
||||
dtype = weight.dtype
|
||||
i8_weight = weight
|
||||
scale = i8_weight.abs().max(dim=-1, keepdim=True)[0] / 127
|
||||
i8_weight = i8_weight / scale.to(dtype)
|
||||
i8_weight = torch.clamp(torch.round(i8_weight), -128, 127).to(torch.int8)
|
||||
return i8_weight, scale.float()
|
||||
|
||||
def smooth_quant_weight_gpu_cpu(weight, act_scale, alpha=0.5, device="cpu"):
|
||||
device = torch.device("cpu") if device == "cpu" else weight.device
|
||||
ori_dtype = weight.dtype
|
||||
# [1, k]
|
||||
act_scale = act_scale.float().to(device).view(1, -1)
|
||||
weight = weight.to(device)
|
||||
# [1, k]
|
||||
weight_scale = weight.abs().max(dim=0, keepdim=True)[0].float()
|
||||
if alpha == -1:
|
||||
smooth_scales = torch.ones_like(act_scale)
|
||||
else:
|
||||
smooth_scales = act_scale.pow(alpha) / weight_scale.pow(1 - alpha).clamp(
|
||||
min=1e-5
|
||||
)
|
||||
weight = weight * smooth_scales.to(ori_dtype)
|
||||
i8_weight, weight_scales = per_token_quant_8bit(weight)
|
||||
# 为了可以使用 input * smooth_scales
|
||||
if alpha == -1:
|
||||
smooth_scales = torch.ones_like(act_scale)
|
||||
else:
|
||||
smooth_scales = weight_scale.pow(1 - alpha) / act_scale.pow(alpha).clamp(
|
||||
min=1e-5
|
||||
)
|
||||
return i8_weight, weight_scales, smooth_scales.to(ori_dtype)
|
||||
|
||||
def smooth_quant_weight_kernel(weight, act_scale, alpha=0.5):
|
||||
output = torch.zeros_like(weight,dtype=torch.int8)
|
||||
weight_scales = torch.zeros(weight.shape[:-1],dtype=torch.float, device=weight.device)
|
||||
weight_max = torch.zeros(weight.shape[-1],dtype=torch.float, device=weight.device)
|
||||
smooth_scales = torch.zeros(weight.shape[-1],dtype=weight.dtype, device=weight.device)
|
||||
ops.infer.weight_quant_smoothquant(
|
||||
weight, act_scale, alpha, output, weight_scales, smooth_scales, weight_max
|
||||
)
|
||||
return output, weight_scales.view(-1,1), smooth_scales.view(1,-1)
|
||||
|
||||
def smooth_quant_weight(weight, act_scale, alpha=0.5):
|
||||
if dynamic_quant_type == "kernel":
|
||||
return smooth_quant_weight_kernel(weight,act_scale,alpha)
|
||||
else:
|
||||
return smooth_quant_weight_gpu_cpu(weight,act_scale,alpha,dynamic_quant_type)
|
||||
|
||||
model = self.model_runner.model
|
||||
|
||||
from vllm.distributed import (
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_all_reduce,
|
||||
get_tensor_model_parallel_world_size
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.model_executor.models.falcon import FalconForCausalLM
|
||||
|
||||
for name, m in model.named_modules():
|
||||
if isinstance(m, VocabParallelEmbedding):
|
||||
# weight shape: [vocab_size // tp, embedding_dim]
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
||||
weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
|
||||
print(f"merged: {name}, shape={m.weight.shape}")
|
||||
|
||||
elif isinstance(m, ParallelLMHead):
|
||||
# weight shape: [vocab_size // tp, embedding_dim]
|
||||
# bias shape: [vocab_size // tp]
|
||||
if m.bias is not None:
|
||||
bias = tensor_model_parallel_all_gather(m.bias, dim=0)
|
||||
bias = bias[:m.org_vocab_size].contiguous()
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
|
||||
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
||||
weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
|
||||
print(f"merged: {name}, shape={m.weight.shape}")
|
||||
|
||||
elif isinstance(m, QKVParallelLinear):
|
||||
# weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
|
||||
# bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
|
||||
if self.parallel_config.world_size > 1:
|
||||
total_q_hidden_size = m.total_num_heads * m.head_size
|
||||
partial_q_hidden_size = m.num_heads * m.head_size
|
||||
total_kv_hidden_size = m.total_num_kv_heads * m.head_size
|
||||
partial_kv_hidden_size = m.num_kv_heads * m.head_size
|
||||
|
||||
if m.bias is not None:
|
||||
# TODO do not support padding..
|
||||
bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
|
||||
q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
|
||||
k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
|
||||
[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
||||
v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
|
||||
[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
||||
|
||||
q_bias[:] = m.bias[:partial_q_hidden_size]
|
||||
k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
|
||||
v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
|
||||
|
||||
bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
|
||||
|
||||
q_tensor = m.weight.new_zeros(m.total_num_heads * m.head_size, m.weight.shape[1])
|
||||
k_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
|
||||
v_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
|
||||
|
||||
q_in_weight = m.weight[:-m.num_kv_heads * m.head_size * 2]
|
||||
k_in_weight = m.weight[-m.num_kv_heads * m.head_size * 2:-m.num_kv_heads * m.head_size]
|
||||
v_in_weight = m.weight[-m.num_kv_heads * m.head_size:]
|
||||
|
||||
if getattr(m,"start_idx",None) is not None:
|
||||
start_idx = getattr(m,"start_idx")
|
||||
weight_end_idx = m.num_heads * m.head_size if not getattr(m,"is_padding") else (m.num_heads - 1) * m.head_size
|
||||
end_idx = start_idx + weight_end_idx
|
||||
else:
|
||||
start_idx = self.rank * m.num_heads * m.head_size
|
||||
weight_end_idx = m.num_heads * m.head_size
|
||||
end_idx = start_idx + weight_end_idx
|
||||
assert q_tensor[start_idx:end_idx,:].shape == q_in_weight[:weight_end_idx, :].shape
|
||||
q_tensor[start_idx:end_idx,:] = q_in_weight[:weight_end_idx, :]
|
||||
|
||||
if m.num_kv_head_replicas > 1:
|
||||
if self.rank % m.num_kv_head_replicas == 0:
|
||||
rank = self.rank // m.num_kv_head_replicas
|
||||
k_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
|
||||
v_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
|
||||
else:
|
||||
k_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
|
||||
v_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
|
||||
|
||||
q_tensor = tensor_model_parallel_all_reduce(q_tensor)
|
||||
k_tensor = tensor_model_parallel_all_reduce(k_tensor)
|
||||
v_tensor = tensor_model_parallel_all_reduce(v_tensor)
|
||||
|
||||
if isinstance(model, FalconForCausalLM):
|
||||
num_query_heads_per_kv_head = (
|
||||
m.total_num_heads // m.total_num_kv_heads
|
||||
)
|
||||
q_tensor = q_tensor.view(
|
||||
m.total_num_kv_heads,
|
||||
num_query_heads_per_kv_head,
|
||||
m.head_size,
|
||||
-1,
|
||||
)
|
||||
k_tensor = k_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
|
||||
v_tensor = v_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
|
||||
weight_tensor = torch.cat(
|
||||
[q_tensor, k_tensor, v_tensor], dim=1
|
||||
).view(-1, m.hidden_size)
|
||||
else:
|
||||
weight_tensor = torch.cat([q_tensor, k_tensor, v_tensor])
|
||||
assert (
|
||||
weight_tensor.shape[0]
|
||||
== total_q_hidden_size + total_kv_hidden_size * 2
|
||||
)
|
||||
assert weight_tensor.shape[1] == m.hidden_size
|
||||
|
||||
else:
|
||||
weight_tensor = m.weight
|
||||
if m.bias is not None and self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
||||
|
||||
if self.is_driver_worker:
|
||||
i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
|
||||
weight_tensor, m.act_scales, smooth_alpha
|
||||
)
|
||||
m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
|
||||
m.weight_scales = torch.nn.Parameter(
|
||||
weight_scales.cpu(), requires_grad=False
|
||||
)
|
||||
m.smooth_scales = torch.nn.Parameter(
|
||||
smooth_scales.cpu(), requires_grad=False
|
||||
)
|
||||
print(f"Quantized: {name}")
|
||||
|
||||
elif isinstance(m, MergedColumnParallelLinear):
|
||||
if self.parallel_config.world_size > 1:
|
||||
# weight shape: [intermediate_size // tp * 2, hidden_size]
|
||||
# bias shape: [intermediate_size // tp * 2]
|
||||
output_sizes = m.output_sizes
|
||||
output_size = sum(output_sizes)
|
||||
partial_output_sizes = [
|
||||
i // self.parallel_config.world_size for i in output_sizes
|
||||
]
|
||||
|
||||
if m.bias is not None:
|
||||
index_start = 0
|
||||
partial_index_start = 0
|
||||
bias_tenosr = m.bias.new_zeros(output_size)
|
||||
for i in range(len(output_sizes)):
|
||||
index_out = index_start + output_sizes[i]
|
||||
sub_bias_tensor = bias_tenosr[index_start:index_out]
|
||||
partial_size = partial_output_sizes[i]
|
||||
sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
|
||||
|
||||
index_start += output_sizes[i]
|
||||
partial_index_start += partial_size
|
||||
bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
|
||||
|
||||
weight_tensor = m.weight.new_zeros(output_size, m.input_size)
|
||||
|
||||
idx_out_start = 0
|
||||
idx_partial_satrt = 0
|
||||
for i in range(len(output_sizes)):
|
||||
idx_out_end = idx_out_start + output_sizes[i]
|
||||
sub_weight_tensor = weight_tensor[idx_out_start:idx_out_end]
|
||||
partial_size = partial_output_sizes[i]
|
||||
sub_weight_tensor[
|
||||
self.rank * partial_size : (self.rank + 1) * partial_size
|
||||
] = m.weight[idx_partial_satrt : idx_partial_satrt + partial_size]
|
||||
|
||||
idx_out_start += output_sizes[i]
|
||||
idx_partial_satrt += partial_size
|
||||
weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
|
||||
else:
|
||||
weight_tensor = m.weight
|
||||
if m.bias is not None and self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
||||
|
||||
if self.is_driver_worker:
|
||||
i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
|
||||
weight_tensor, m.act_scales, smooth_alpha
|
||||
)
|
||||
m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
|
||||
m.weight_scales = torch.nn.Parameter(
|
||||
weight_scales.cpu(), requires_grad=False
|
||||
)
|
||||
m.smooth_scales = torch.nn.Parameter(
|
||||
smooth_scales.cpu(), requires_grad=False
|
||||
)
|
||||
print(f"Quantized: {name}")
|
||||
|
||||
elif isinstance(m, ColumnParallelLinear):
|
||||
# weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
|
||||
# bias shape: [some_dim // tp]
|
||||
if m.bias is not None:
|
||||
bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
|
||||
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
||||
|
||||
if self.is_driver_worker:
|
||||
i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
|
||||
weight_tensor, m.act_scales, smooth_alpha
|
||||
)
|
||||
m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
|
||||
m.weight_scales = torch.nn.Parameter(
|
||||
weight_scales.cpu(), requires_grad=False
|
||||
)
|
||||
m.smooth_scales = torch.nn.Parameter(
|
||||
smooth_scales.cpu(), requires_grad=False
|
||||
)
|
||||
print(f"Quantized: {name}")
|
||||
|
||||
elif isinstance(m, RowParallelLinear):
|
||||
# weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
|
||||
# bias shape: [hidden_size]
|
||||
if m.bias is not None:
|
||||
bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
||||
|
||||
if getattr(m,"start_idx", None) is not None:
|
||||
start_idx = getattr(m,"start_idx")
|
||||
end_idx = start_idx + (m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size))
|
||||
weight_end_idx = m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size)
|
||||
else:
|
||||
start_idx = m.input_size_per_partition * self.rank
|
||||
end_idx = start_idx + m.input_size_per_partition
|
||||
weight_end_idx = m.input_size_per_partition
|
||||
|
||||
act_scales = m.act_scales.new_zeros(m.input_size)
|
||||
assert act_scales[start_idx:end_idx].shape == m.act_scales.view(-1)[:weight_end_idx].shape
|
||||
act_scales[start_idx:end_idx] = m.act_scales.view(-1)[:weight_end_idx]
|
||||
act_scales = tensor_model_parallel_all_reduce(act_scales)
|
||||
m.act_scales = act_scales
|
||||
|
||||
weight_tensor = m.weight.new_zeros(m.weight.shape[0],m.input_size)
|
||||
assert weight_tensor[:,start_idx:end_idx].shape == m.weight[:,:weight_end_idx].shape
|
||||
weight_tensor[:,start_idx:end_idx] = m.weight[:,:weight_end_idx]
|
||||
weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
|
||||
|
||||
if self.is_driver_worker:
|
||||
i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
|
||||
weight_tensor, m.act_scales, smooth_alpha
|
||||
)
|
||||
smooth_scales = smooth_scales.view(1,-1)
|
||||
m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
|
||||
m.weight_scales = torch.nn.Parameter(
|
||||
weight_scales.cpu(), requires_grad=False
|
||||
)
|
||||
m.smooth_scales = torch.nn.Parameter(
|
||||
smooth_scales.cpu(), requires_grad=False
|
||||
)
|
||||
print(f"Quantized: {name}")
|
||||
else:
|
||||
pass
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# save weights
|
||||
if self.is_driver_worker:
|
||||
from safetensors.torch import save_file
|
||||
|
||||
tensors = {}
|
||||
saved = False
|
||||
count = 0
|
||||
size_in_bytes = 0
|
||||
|
||||
tensors = {}
|
||||
for name, weight in model.named_parameters():
|
||||
if "act_scales" in name:
|
||||
continue
|
||||
# skip lm_head_weight if needed..
|
||||
if "lm_head" in name and model.config.tie_word_embeddings:
|
||||
continue
|
||||
tensors[name] = weight
|
||||
|
||||
saved = False
|
||||
if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
|
||||
weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
|
||||
save_file(tensors, weight_path)
|
||||
print(f"The quantified weights were successfully saved in {weight_path}.")
|
||||
tensors.clear()
|
||||
saved = True
|
||||
count += 1
|
||||
size_in_bytes = 0
|
||||
|
||||
if not saved:
|
||||
weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
|
||||
save_file(tensors, weight_path)
|
||||
print(f"The quantified weights were successfully saved in {weight_path}.")
|
||||
233
ixformer_sdk/contrib/vllm/quantize/w8a16.py
Normal file
233
ixformer_sdk/contrib/vllm/quantize/w8a16.py
Normal file
@@ -0,0 +1,233 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
def w8a16_prepare_quantize(self, quant_params={}):
|
||||
# We need do nothing in here
|
||||
pass
|
||||
|
||||
|
||||
def w8a16_export_quantized_weights(self, save_path, quant_params={}):
|
||||
gb_per_file = quant_params.get("filesize_limit", None)
|
||||
int8_min = -127
|
||||
|
||||
def w8a16_quantization(weight):
|
||||
# all weights should be [output,input], otherwise, we may get an wrong weight and scale...
|
||||
scale = torch.abs(weight).max(dim=-1)[0] / 127.0
|
||||
int8_weight = torch.clamp(weight / scale.view(-1,1),min=int8_min,max=127).to(torch.int8).contiguous()
|
||||
scale = scale.view(1,-1).contiguous()
|
||||
return int8_weight, scale
|
||||
|
||||
|
||||
model = self.model_runner.model
|
||||
|
||||
from vllm.distributed.communication_op import (
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
|
||||
for name, m in model.named_modules():
|
||||
if isinstance(m, VocabParallelEmbedding):
|
||||
# weight shape: [vocab_size // tp, embedding_dim]
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
||||
weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
|
||||
print(f"merged: {name}, shape={m.weight.shape}")
|
||||
|
||||
elif isinstance(m, ParallelLMHead):
|
||||
# weight shape: [vocab_size // tp, embedding_dim]
|
||||
# bias shape: [vocab_size // tp]
|
||||
if m.bias is not None:
|
||||
bias = tensor_model_parallel_all_gather(m.bias, dim=0)
|
||||
bias = bias[:m.org_vocab_size].contiguous()
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
|
||||
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
||||
weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
|
||||
print(f"merged: {name}, shape={m.weight.shape}")
|
||||
|
||||
elif isinstance(m, QKVParallelLinear):
|
||||
# weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
|
||||
# bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
|
||||
if self.parallel_config.world_size > 1:
|
||||
total_q_hidden_size = m.total_num_heads * m.head_size
|
||||
partial_q_hidden_size = m.num_heads * m.head_size
|
||||
total_kv_hidden_size = m.total_num_kv_heads * m.head_size
|
||||
partial_kv_hidden_size = m.num_kv_heads * m.head_size
|
||||
|
||||
if m.bias is not None:
|
||||
bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
|
||||
q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
|
||||
k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
|
||||
[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
||||
v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
|
||||
[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
||||
|
||||
q_bias[:] = m.bias[:partial_q_hidden_size]
|
||||
k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
|
||||
v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
|
||||
|
||||
bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
|
||||
|
||||
weight_tensor = m.weight.new_zeros(
|
||||
total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size
|
||||
)
|
||||
|
||||
q_tensor = weight_tensor[:total_q_hidden_size, :]
|
||||
q_tensor = q_tensor[self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
|
||||
|
||||
k_tensor = weight_tensor[total_q_hidden_size : total_q_hidden_size + total_kv_hidden_size]
|
||||
k_tensor = k_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
||||
|
||||
v_tensor = weight_tensor[total_q_hidden_size + total_kv_hidden_size :]
|
||||
v_tensor = v_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
||||
|
||||
q_tensor[:, :] = m.weight[: partial_q_hidden_size, :]
|
||||
k_tensor[:, :] = m.weight[partial_q_hidden_size : partial_q_hidden_size + partial_kv_hidden_size, :]
|
||||
v_tensor[:, :] = m.weight[partial_q_hidden_size + partial_kv_hidden_size : , :]
|
||||
|
||||
weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
|
||||
else:
|
||||
weight_tensor = m.weight
|
||||
if m.bias is not None and self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
||||
|
||||
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
||||
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
||||
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
||||
print(f"Quantized: {name}")
|
||||
|
||||
elif isinstance(m, MergedColumnParallelLinear):
|
||||
# weight shape: [intermediate_size // tp * 2, hidden_size]
|
||||
# bias shape: [intermediate_size // tp * 2]
|
||||
if self.parallel_config.world_size > 1:
|
||||
output_sizes = m.output_sizes
|
||||
output_size = sum(output_sizes)
|
||||
partial_output_sizes = [
|
||||
i // self.parallel_config.world_size for i in output_sizes
|
||||
]
|
||||
|
||||
if m.bias is not None:
|
||||
index_start = 0
|
||||
partial_index_start = 0
|
||||
bias_tenosr = m.bias.new_zeros(output_size)
|
||||
for i in range(len(output_sizes)):
|
||||
index_out = index_start + output_sizes[i]
|
||||
sub_bias_tensor = bias_tenosr[index_start:index_out]
|
||||
partial_size = partial_output_sizes[i]
|
||||
sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
|
||||
|
||||
index_start += output_sizes[i]
|
||||
partial_index_start += partial_size
|
||||
bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
|
||||
|
||||
weight_tensor = m.weight.new_zeros(output_size, m.input_size)
|
||||
|
||||
index_start = 0
|
||||
partial_index_start = 0
|
||||
for i in range(len(output_sizes)):
|
||||
index_out = index_start + output_sizes[i]
|
||||
sub_weight_tensor = weight_tensor[index_start:index_out]
|
||||
partial_size = partial_output_sizes[i]
|
||||
sub_weight_tensor[self.rank * partial_size : (self.rank + 1) * partial_size] = m.weight[partial_index_start : partial_index_start + partial_size]
|
||||
|
||||
index_start += output_sizes[i]
|
||||
partial_index_start += partial_size
|
||||
weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
|
||||
else:
|
||||
weight_tensor = m.weight
|
||||
if m.bias is not None and self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
||||
|
||||
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
||||
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
||||
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
||||
print(f"Quantized: {name}")
|
||||
|
||||
elif isinstance(m, ColumnParallelLinear):
|
||||
# weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
|
||||
# bias shape: [some_dim // tp]
|
||||
if m.bias is not None:
|
||||
bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
|
||||
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
||||
|
||||
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
||||
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
||||
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
||||
print(f"Quantized: {name}")
|
||||
|
||||
elif isinstance(m, RowParallelLinear):
|
||||
# weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
|
||||
# bias shape: [hidden_size]
|
||||
if m.bias is not None:
|
||||
bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
|
||||
if self.is_driver_worker:
|
||||
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
||||
|
||||
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=-1)
|
||||
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
||||
|
||||
if self.is_driver_worker:
|
||||
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
||||
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
||||
print(f"Quantized: {name}")
|
||||
else:
|
||||
pass
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# save weights
|
||||
if self.is_driver_worker:
|
||||
from safetensors.torch import save_file
|
||||
|
||||
tensors = {}
|
||||
saved = False
|
||||
count = 0
|
||||
size_in_bytes = 0
|
||||
|
||||
for name, weight in model.named_parameters():
|
||||
if "lm_head" in name and model.config.tie_word_embeddings:
|
||||
continue
|
||||
size_in_bytes += weight.numel() * weight.element_size()
|
||||
tensors[name] = weight
|
||||
saved = False
|
||||
if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
|
||||
weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
|
||||
save_file(tensors, weight_path)
|
||||
print(f"The quantified weights were successfully saved in {weight_path}.")
|
||||
tensors.clear()
|
||||
saved = True
|
||||
count += 1
|
||||
size_in_bytes = 0
|
||||
|
||||
if not saved:
|
||||
weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
|
||||
save_file(tensors, weight_path)
|
||||
print(f"The quantified weights were successfully saved in {weight_path}.")
|
||||
3
ixformer_sdk/contrib/vllm_flash_attn/__init__.py
Normal file
3
ixformer_sdk/contrib/vllm_flash_attn/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
__version__ = "2.6.1"
|
||||
|
||||
from .flash_attn_interface import *
|
||||
1018
ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py
Normal file
1018
ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user