mirror of https://github.com/vladmandic/automatic
28 lines
1.1 KiB
Python
28 lines
1.1 KiB
Python
import torch
|
|
import diffusers
|
|
|
|
|
|
repo_id = 'Kwai-Kolors/Kolors-diffusers'
|
|
|
|
|
|
def load_kolors(_checkpoint_info, diffusers_load_config={}):
|
|
from modules import shared, devices
|
|
diffusers_load_config['variant'] = "fp16"
|
|
if 'torch_dtype' not in diffusers_load_config:
|
|
diffusers_load_config['torch_dtype'] = torch.float16
|
|
|
|
# import torch
|
|
# import transformers
|
|
# encoder_id = 'THUDM/chatglm3-6b'
|
|
# text_encoder = transformers.AutoModel.from_pretrained(encoder_id, torch_dtype=torch.float16, trust_remote_code=True, cache_dir=shared.opts.diffusers_dir)
|
|
# text_encoder = transformers.AutoModel.from_pretrained("THUDM/chatglm3-6b", torch_dtype=torch.float16, trust_remote_code=True).quantize(4).cuda()
|
|
# tokenizer = transformers.AutoTokenizer.from_pretrained(encoder_id, trust_remote_code=True, cache_dir=shared.opts.diffusers_dir)
|
|
pipe = diffusers.KolorsPipeline.from_pretrained(
|
|
repo_id,
|
|
cache_dir = shared.opts.diffusers_dir,
|
|
**diffusers_load_config,
|
|
)
|
|
pipe.vae.config.force_upcast = True
|
|
devices.torch_gc()
|
|
return pipe
|