diff --git a/modules/model_kolors.py b/modules/model_kolors.py index dcf7c1f26..a89af487b 100644 --- a/modules/model_kolors.py +++ b/modules/model_kolors.py @@ -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, )