141 lines
5.3 KiB
Python
141 lines
5.3 KiB
Python
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# 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.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import torch
|
|
from diffusers.models import AutoencoderKL
|
|
from termcolor import colored
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, T5EncoderModel, T5Tokenizer
|
|
from transformers import logging as transformers_logging
|
|
from .dc_ae.efficientvit.ae_model_zoo import DCAE_HF
|
|
from .utils import set_fp32_attention, set_grad_checkpoint
|
|
|
|
try:
|
|
from mmengine.registry import Registry
|
|
except ImportError as e:
|
|
from mmcv import Registry
|
|
|
|
MODELS = Registry("models")
|
|
|
|
transformers_logging.set_verbosity_error()
|
|
|
|
|
|
def build_model(cfg, use_grad_checkpoint=False, use_fp32_attention=False, gc_step=1, **kwargs):
|
|
if isinstance(cfg, str):
|
|
cfg = dict(type=cfg)
|
|
model = MODELS.build(cfg, default_args=kwargs)
|
|
|
|
if use_grad_checkpoint:
|
|
set_grad_checkpoint(model, gc_step=gc_step)
|
|
if use_fp32_attention:
|
|
set_fp32_attention(model)
|
|
return model
|
|
|
|
|
|
def get_tokenizer_and_text_encoder(name="T5", device="cuda", text_encoder_dict_custom={}):
|
|
text_encoder_dict = {
|
|
"T5": "DeepFloyd/t5-v1_1-xxl",
|
|
"T5-small": "google/t5-v1_1-small",
|
|
"T5-base": "google/t5-v1_1-base",
|
|
"T5-large": "google/t5-v1_1-large",
|
|
"T5-xl": "google/t5-v1_1-xl",
|
|
"T5-xxl": "google/t5-v1_1-xxl",
|
|
"gemma-2b": "google/gemma-2b",
|
|
"gemma-2b-it": "google/gemma-2b-it",
|
|
"gemma-2-2b": "google/gemma-2-2b",
|
|
"gemma-2-2b-it": "google/gemma-2-2b-it",
|
|
"gemma-2-9b": "google/gemma-2-9b",
|
|
"gemma-2-9b-it": "google/gemma-2-9b-it",
|
|
"Qwen2-0.5B-Instruct": "Qwen/Qwen2-0.5B-Instruct",
|
|
"Qwen2-1.5B-Instruct": "Qwen/Qwen2-1.5B-Instruct",
|
|
}
|
|
assert name in list(text_encoder_dict_custom.keys()), f"not support this text encoder: {name}"
|
|
if "T5" in name:
|
|
tokenizer = T5Tokenizer.from_pretrained(text_encoder_dict_custom[name])
|
|
text_encoder = T5EncoderModel.from_pretrained(text_encoder_dict_custom[name], torch_dtype=torch.float16).to(device)
|
|
elif "gemma" in name or "Qwen" in name:
|
|
print(f"\033[93m{text_encoder_dict_custom[name]}\033[0m")
|
|
tokenizer = AutoTokenizer.from_pretrained(text_encoder_dict_custom[name])
|
|
tokenizer.padding_side = "right"
|
|
# text_encoder = (
|
|
# AutoModelForCausalLM.from_pretrained(text_encoder_dict_custom[name], torch_dtype=torch.bfloat16)
|
|
# .get_decoder()
|
|
# .to(device)
|
|
# )
|
|
text_encoder = AutoModelForCausalLM.from_pretrained(text_encoder_dict_custom[name], torch_dtype=torch.bfloat16).get_decoder()
|
|
# elif "gemma" in name or "Qwen" in name:
|
|
# tokenizer = AutoTokenizer.from_pretrained(r'C:\Users\pc\Desktop\New_Folder\SANA\models--google--gemma-2-2b-it')
|
|
# tokenizer.padding_side = "right"
|
|
# text_encoder = (
|
|
# AutoModelForCausalLM.from_pretrained(r'C:\Users\pc\Desktop\New_Folder\SANA\models--google--gemma-2-2b-it', torch_dtype=torch.bfloat16)
|
|
# .get_decoder()
|
|
# .to(device)
|
|
# )
|
|
else:
|
|
print("error load text encoder")
|
|
exit()
|
|
|
|
return tokenizer, text_encoder
|
|
|
|
|
|
def get_vae(name, model_path, device="cuda"):
|
|
if name == "sdxl" or name == "sd3":
|
|
vae = AutoencoderKL.from_pretrained(model_path).to(device).to(torch.float16)
|
|
if name == "sdxl":
|
|
vae.config.shift_factor = 0
|
|
return vae
|
|
# elif "dc-ae" in name:
|
|
# print(colored(f"[DC-AE] Loading model from {model_path}", attrs=["bold"]))
|
|
# dc_ae = DCAE_HF.from_pretrained(model_path).to(device).eval()
|
|
# return dc_ae
|
|
elif "dc-ae" in name:
|
|
print(colored(f"[DC-AE] Loading model from {model_path}", attrs=["bold"]))
|
|
dc_ae = DCAE_HF.from_pretrained(model_path).to(device).eval()
|
|
return dc_ae
|
|
else:
|
|
print("error load vae")
|
|
exit()
|
|
|
|
|
|
def vae_encode(name, vae, images, sample_posterior, device):
|
|
if name == "sdxl" or name == "sd3":
|
|
posterior = vae.encode(images.to(device)).latent_dist
|
|
if sample_posterior:
|
|
z = posterior.sample()
|
|
else:
|
|
z = posterior.mode()
|
|
z = (z - vae.config.shift_factor) * vae.config.scaling_factor
|
|
elif "dc-ae" in name:
|
|
ae = vae
|
|
z = ae.encode(images.to(device))
|
|
z = z * ae.cfg.scaling_factor
|
|
else:
|
|
print("error load vae")
|
|
exit()
|
|
return z
|
|
|
|
|
|
def vae_decode(name, vae, latent):
|
|
if name == "sdxl" or name == "sd3":
|
|
latent = (latent.detach() / vae.config.scaling_factor) + vae.config.shift_factor
|
|
samples = vae.decode(latent).sample
|
|
elif "dc-ae" in name:
|
|
ae = vae
|
|
samples = ae.decode(latent.detach() / ae.cfg.scaling_factor)
|
|
else:
|
|
print("error load vae")
|
|
exit()
|
|
return samples
|