reorg imports for faster startup

This commit is contained in:
Vladimir Mandic
2024-04-07 16:29:33 -04:00
parent 01b9c4ddf0
commit 19db61a861
3 changed files with 17 additions and 7 deletions
+7 -3
View File
@@ -37,13 +37,17 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__:
torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)
timer.startup.record("torch")
import transformers # pylint: disable=W0611,C0411
timer.startup.record("transformers")
import onnxruntime
onnxruntime.set_default_logger_severity(3)
timer.startup.record("onnx")
from modules.onnx_impl import initialize_olive # pylint: disable=ungrouped-imports
initialize_olive()
timer.startup.record("olive")
# moved to webui.py:initialize()
# from modules.onnx_impl import initialize_olive # pylint: disable=ungrouped-imports
# initialize_olive()
# timer.startup.record("olive")
from fastapi import FastAPI # pylint: disable=W0611,C0411
import gradio # pylint: disable=W0611,C0411
+4 -4
View File
@@ -1,14 +1,10 @@
import torch
import torch.nn as nn
import kornia
from torch.utils.checkpoint import checkpoint
from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel
import open_clip
from ldm.util import default, count_params, autocast
class AbstractEncoder(nn.Module):
def __init__(self):
super().__init__()
@@ -62,6 +58,7 @@ class FrozenT5Embedder(AbstractEncoder):
def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77,
freeze=True): # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
super().__init__()
from transformers import T5Tokenizer, T5EncoderModel
self.tokenizer = T5Tokenizer.from_pretrained(version)
self.transformer = T5EncoderModel.from_pretrained(version)
self.device = device
@@ -99,6 +96,7 @@ class FrozenCLIPEmbedder(AbstractEncoder):
def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,
freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32
super().__init__()
from transformers import CLIPTokenizer, CLIPTextModel
assert layer in self.LAYERS
self.tokenizer = CLIPTokenizer.from_pretrained(version)
self.transformer = CLIPTextModel.from_pretrained(version)
@@ -155,6 +153,7 @@ class ClipImageEmbedder(nn.Module):
self.ucg_rate = ucg_rate
def preprocess(self, x):
import kornia
# normalize to [0,1]
x = kornia.geometry.resize(x, (224, 224),
interpolation='bicubic', align_corners=True,
@@ -265,6 +264,7 @@ class FrozenOpenCLIPImageEmbedder(AbstractEncoder):
self.ucg_rate = ucg_rate
def preprocess(self, x):
import kornia
# normalize to [0,1]
x = kornia.geometry.resize(x, (224, 224),
interpolation='bicubic', align_corners=True,
+6
View File
@@ -81,6 +81,12 @@ def initialize():
log.debug('Initializing')
check_rollback_vae()
if shared.opts.cuda_compile_backend == "olive-ai":
from modules.onnx_impl import initialize_olive, install_olive
install_olive()
initialize_olive()
timer.startup.record("olive")
modules.sd_samplers.list_samplers()
timer.startup.record("samplers")