draft kwai kolors

modules/model_kolors.py

@@ -0,0 +1,28 @@
+import torch
+import transformers
+import diffusers
+
+
+repo_id = 'Kwai-Kolors/Kolors'
+encoder_id = 'THUDM/chatglm3-6b'
+
+
+def load_kolors(_checkpoint_info, diffusers_load_config={}):
+    from modules import shared, devices, modelloader
+    modelloader.hf_login()
+    diffusers_load_config['variant'] = "fp16"
+    if 'torch_dtype' not in diffusers_load_config:
+        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)
+    # 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(
+        repo_id,
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        cache_dir = shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    devices.torch_gc()
+    return pipe

modules/sd_models.py

@@ -613,6 +613,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 guess = 'PixArt-Sigma'
             if 'lumina-next' in f.lower():
                 guess = 'Lumina-Next'
+            if 'kolors' in f.lower():
+                guess = 'Kolors'
             # switch for specific variant
             if guess == 'Stable Diffusion' and 'inpaint' in f.lower():
                 guess = 'Stable Diffusion Inpaint'
@@ -1003,6 +1005,15 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     if debug_load:
                         errors.display(e, 'Load')
                     return
+            elif model_type in ['Kolors']: # forced pipeline
+                try:
+                    from modules.model_kolors import load_kolors
+                    sd_model = load_kolors(checkpoint_info, diffusers_load_config)
+                except Exception as e:
+                    shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}')
+                    if debug_load:
+                        errors.display(e, 'Load')
+                    return
             elif model_type in ['Stable Diffusion 3']:
                 try:
                     from modules.model_sd3 import load_sd3

modules/shared_items.py

@@ -71,6 +71,7 @@ def get_pipelines():
         'Kandinsky 3': getattr(diffusers, 'Kandinsky3Pipeline', None),
         'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
         'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None),
+        'Kolors': getattr(diffusers, 'StableDiffusionXLPipeline', None),
         'InstaFlow': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
         'SegMoE': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
     }