prep kolors

This commit is contained in:
Vladimir Mandic
2024-07-10 10:28:55 -04:00
parent 03a6e25d48
commit 8e170f54a6
+7 -8
View File
@@ -1,10 +1,8 @@
import torch
import transformers
import diffusers
repo_id = 'Kwai-Kolors/Kolors'
encoder_id = 'THUDM/chatglm3-6b'
def load_kolors(_checkpoint_info, diffusers_load_config={}):
@@ -12,15 +10,16 @@ def load_kolors(_checkpoint_info, diffusers_load_config={}):
modelloader.hf_login()
diffusers_load_config['variant'] = "fp16"
if 'torch_dtype' not in diffusers_load_config:
diffusers_load_config['torch_dtype'] = 'torch.float16'
diffusers_load_config['torch_dtype'] = torch.float16
text_encoder = transformers.AutoModel.from_pretrained(encoder_id, torch_dtype=torch.float16, trust_remote_code=True, cache_dir=shared.opts.diffusers_dir)
# 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.StableDiffusionXLPipeline.from_pretrained(
# tokenizer = transformers.AutoTokenizer.from_pretrained(encoder_id, trust_remote_code=True, cache_dir=shared.opts.diffusers_dir)
pipe = diffusers.KolorsPipeline.from_pretrained(
repo_id,
tokenizer=tokenizer,
text_encoder=text_encoder,
cache_dir = shared.opts.diffusers_dir,
**diffusers_load_config,
)