fix chroma

Signed-off-by: Vladimir Mandic <mandic00@live.com>

cli/test-all-models.py

@@ -1,10 +1,14 @@
 #!/usr/bin/env python
 """
+Warning:
 - fal/AuraFlow-v0.3: layer_class_name=Linear layer_weight_shape=torch.Size([3072, 2, 1024]) weights_dtype=int8 unsupported
+- Kwai-Kolors/Kolors-diffusers: `set_input_embeddings` not auto‑handled for ChatGLMModel
+Error:
 - nvidia/Cosmos-Predict2-2B-Text2Image: mat1 and mat2 shapes cannot be multiplied (512x4096 and 1024x2048)
 - nvidia/Cosmos-Predict2-14B-Text2Image: mat1 and mat2 shapes cannot be multiplied (512x4096 and 1024x5120)
+- Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers: CUDA error: device-side assert triggered
+Other:
 - HiDream-ai/HiDream-I1-Full: 30+s/it
-- Kwai-Kolors/Kolors-diffusers: `set_input_embeddings` not auto‑handled for ChatGLMModel
 """
 
 import io
@@ -25,49 +29,53 @@ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 
 output_folder = 'outputs/compare'
-models = [
-    "sdxl-base-v10-vaefix",
-    "tempest-by-vlad-0.1",
-    "icbinpXL_v6",
-    "briaai/BRIA-3.2",
-    "Freepik/F-Lite",
-    "Freepik/F-Lite-Texture",
-    "ostris/Flex.2-preview",
-    "stabilityai/stable-diffusion-3.5-medium",
-    "stabilityai/stable-diffusion-3.5-large",
-    "fal/AuraFlow-v0.3",
-    "fal/AuraFlow-v0.2",
-    "zai-org/CogView4-6B",
-    "zai-org/CogView3-Plus-3B",
-    # "nvidia/Cosmos-Predict2-2B-Text2Image",
-    # "nvidia/Cosmos-Predict2-14B-Text2Image",
-    "Qwen/Qwen-Image",
-    "vladmandic/Qwen-Lightning",
-    "Shitao/OmniGen-v1-diffusers",
-    "OmniGen2/OmniGen2",
-    # "HiDream-ai/HiDream-I1-Full",
-    "Kwai-Kolors/Kolors-diffusers",
-    "lodestones/Chroma1-HD",
-    "vladmandic/chroma-unlocked-v50-annealed",
-    "vladmandic/chroma-unlocked-v48",
-    "vladmandic/chroma-unlocked-v48-detail-calibrated",
-    "Alpha-VLLM/Lumina-Next-SFT-diffusers",
-    "Alpha-VLLM/Lumina-Image-2.0",
-    "MeissonFlow/Meissonic",
-    "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers",
-    "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers",
-    "PixArt-alpha/PixArt-XL-2-1024-MS",
-    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-    "stabilityai/stable-cascade",
-]
+models = {
+    "sdxl-base-v10-vaefix": {},
+    "tempest-by-vlad-0.1": {},
+    "icbinpXL_v6": {},
+    "briaai/BRIA-3.2": {},
+    "Freepik/F-Lite": {},
+    "Freepik/F-Lite-Texture": {},
+    "ostris/Flex.2-preview": {},
+    "playgroundai/playground-v2-1024px-aesthetic": {},
+    "playground-v2.5-1024px-aesthetic.fp16": { "sampler_name": "DPM++ 2M EDM" },
+    "stabilityai/stable-diffusion-3.5-medium": {},
+    "stabilityai/stable-diffusion-3.5-large": {},
+    "fal/AuraFlow-v0.3": {},
+    "fal/AuraFlow-v0.2": {},
+    "zai-org/CogView4-6B": {},
+    "zai-org/CogView3-Plus-3B": {},
+    # "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers": {},
+    # "nvidia/Cosmos-Predict2-2B-Text2Image": {},
+    # "nvidia/Cosmos-Predict2-14B-Text2Image": {},
+    "Qwen/Qwen-Image": {},
+    "vladmandic/Qwen-Lightning": {},
+    "Shitao/OmniGen-v1-diffusers": {},
+    "OmniGen2/OmniGen2": {},
+    # "HiDream-ai/HiDream-I1-Full": {},
+    "Kwai-Kolors/Kolors-diffusers": {},
+    "lodestones/Chroma1-HD": {},
+    "vladmandic/chroma-unlocked-v50-annealed": {},
+    "vladmandic/chroma-unlocked-v48": {},
+    "vladmandic/chroma-unlocked-v48-detail-calibrated": {},
+    "Alpha-VLLM/Lumina-Next-SFT-diffusers": {},
+    "Alpha-VLLM/Lumina-Image-2.0": {},
+    "MeissonFlow/Meissonic": {},
+    "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers": {},
+    "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers": {},
+    "PixArt-alpha/PixArt-XL-2-1024-MS": {},
+    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS": {},
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": {},
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers": {},
+    "stabilityai/stable-cascade": {},
+}
 models_tbd = [
     "black-forest-labs/FLUX.1-dev",
     "black-forest-labs/FLUX.1-Kontext-dev",
     "black-forest-labs/FLUX.1-Krea-dev",
-    "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers", # TODO
     "kandinsky-community/kandinsky-3", # TODO
+    "kandinsky-community/kandinsky-2-2-decoder",
+    "kandinsky-community/kandinsky-2-1",
 ]
 styles = [
     'Fixed Astronaut',
@@ -103,7 +111,7 @@ def read_history():
     log.info(f'history: file="{fn}" records={len(history)}')
 
 
-def write_history(model:str, style:str, image:str='', size:tuple=(0,0), duration:float=0, info:str='', error:str=''):
+def write_history(model:str, style:str, image:str='', size:tuple=(0,0), duration:float=0, info:str=''):
     fn = os.path.join(output_folder, 'history.json')
     history.append({
         'model': model,
@@ -113,7 +121,6 @@ def write_history(model:str, style:str, image:str='', size:tuple=(0,0), duration
         'size': size,
         'time': duration,
         'info': info,
-        'error': error,
     })
     with open(fn, "w", encoding='utf8') as file:
         data = json.dumps(history) # pylint: disable=no-member
@@ -137,9 +144,11 @@ def request(endpoint: str, dct: dict = None, method: str = 'POST'):
 
 
 def generate(): # pylint: disable=redefined-outer-name
-    for m, model in enumerate(models):
+    idx = 0
+    for model, args in models.items():
+        idx += 1
         model_name = pathvalidate.sanitize_filename(model, replacement_text='_')
-        log.info(f'model: name="{model}" n={m+1}/{len(models)}')
+        log.info(f'model: n={idx+1}/{len(models)} name="{model}"')
         for s, style in enumerate(styles):
             try:
                 model_name = pathvalidate.sanitize_filename(model, replacement_text='_')
@@ -152,9 +161,12 @@ def generate(): # pylint: disable=redefined-outer-name
                 if not loaded or not (model in loaded.get('checkpoint') or model in loaded.get('title') or model in loaded.get('name')):
                     log.error(f' model: error="{model}"')
                     continue
-                log.info(f' style: name="{style}" n={s+1}/{len(styles)} fn="{fn}"')
                 t0 = time.time()
-                data = request('/sdapi/v1/txt2img', { 'styles': [style] })
+                params = { 'styles': [style] }
+                for k, v in args.items():
+                    params[k] = v
+                log.info(f' style: n={s+1}/{len(styles)} name="{style}" args={params} fn="{fn}"')
+                data = request('/sdapi/v1/txt2img', params)
                 t1 = time.time()
                 if 'images' in data and len(data['images']) > 0:
                     b64 = data['images'][0].split(',',1)[0]
@@ -164,13 +176,13 @@ def generate(): # pylint: disable=redefined-outer-name
                     image.save(fn)
                     write_history(model=model, style=style, image=fn, size=image.size, duration=round(t1-t0, 3), info=info)
                 else:
-                    # write_history(model=model, style=style, duration=round(t1-t0, 3), error='no image')
+                    # write_history(model=model, style=style, duration=round(t1-t0, 3), info='no image')
                     log.error(f' model: error="{model}" style="{style}" no image')
             except Exception as e:
                 if 'Connection refused' in str(e) or 'RemoteDisconnected' in str(e):
                     log.error('server offline')
                     os._exit(1)
-                # write_history(model=model, style=style, duration=round(t1-t0, 3), error=str(e))
+                # write_history(model=model, style=style, duration=round(t1-t0, 3), info=str(e))
                 log.error(f' model: error="{model}" style="{style}" exception="{e}"')
 
 

installer.py

@@ -593,7 +593,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git:
         return
-    sha = '7b10e4ae65cc5830c581fba58638f5afb6e587cf' # diffusers commit hash
+    sha = '4a9dbd56f68214f0c949b8036a58c9ac3607f54e' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else -1)
     cur = opts.get('diffusers_version', '') if minor > -1 else ''

modules/sd_models.py

@@ -323,7 +323,8 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
             allow_post_quant = False
         elif model_type in ['Chroma']:
             from pipelines.model_chroma import load_chroma
-            sd_model, allow_post_quant = load_chroma(checkpoint_info, diffusers_load_config)
+            sd_model = load_chroma(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Lumina 2']:
             from pipelines.model_lumina import load_lumina2
             sd_model = load_lumina2(checkpoint_info, diffusers_load_config)
@@ -376,6 +377,10 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
             from pipelines.model_qwen import load_qwen
             sd_model = load_qwen(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
+        elif model_type in ['HunyuanDiT']:
+            from pipelines.model_hunyuandit import load_hunyuandit
+            sd_model = load_hunyuandit(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
     except Exception as e:
         shared.log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
         if debug_load:

pipelines/generic.py

@@ -8,7 +8,7 @@ from modules import shared, devices, sd_models, model_quant
 debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
-def load_transformer(repo_id, cls_name, load_config={}, subfolder="transformer", allow_quant=True):
+def load_transformer(repo_id, cls_name, load_config={}, subfolder="transformer", allow_quant=True, variant=None):
     load_args, quant_args = model_quant.get_dit_args(load_config, module='Model', device_map=True, allow_quant=allow_quant)
     quant_type = model_quant.get_quant_type(quant_args)
 
@@ -45,6 +45,8 @@ def load_transformer(repo_id, cls_name, load_config={}, subfolder="transformer",
         shared.log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" args={load_args}')
         if subfolder is not None:
             load_args['subfolder'] = subfolder
+        if variant is not None:
+            load_args['variant'] = variant
         transformer = cls_name.from_pretrained(
             repo_id,
             cache_dir=shared.opts.hfcache_dir,
@@ -56,7 +58,7 @@ def load_transformer(repo_id, cls_name, load_config={}, subfolder="transformer",
     return transformer
 
 
-def load_text_encoder(repo_id, cls_name, load_config={}, subfolder="text_encoder", allow_quant=True, allow_shared=True):
+def load_text_encoder(repo_id, cls_name, load_config={}, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None):
     load_args, quant_args = model_quant.get_dit_args(load_config, module='TE', device_map=True, allow_quant=allow_quant)
     quant_type = model_quant.get_quant_type(quant_args)
     text_encoder = None
@@ -120,6 +122,8 @@ def load_text_encoder(repo_id, cls_name, load_config={}, subfolder="text_encoder
         shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}')
         if subfolder is not None:
             load_args['subfolder'] = subfolder
+        if variant is not None:
+            load_args['variant'] = variant
         text_encoder = cls_name.from_pretrained(
             repo_id,
             cache_dir=shared.opts.hfcache_dir,

pipelines/model_chroma.py

@@ -30,5 +30,6 @@ def load_chroma(checkpoint_info, diffusers_load_config={}):
     diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["chroma"] = diffusers.ChromaImg2ImgPipeline
     del text_encoder
     del transformer
+
     devices.torch_gc(force=True, reason='load')
     return pipe

pipelines/model_hunyuandit.py

@@ -0,0 +1,29 @@
+import transformers
+import diffusers
+from modules import shared, sd_models, devices, model_quant, sd_hijack_te
+from pipelines import generic
+
+
+def load_hunyuandit(checkpoint_info, diffusers_load_config={}):
+    repo_id = sd_models.path_to_repo(checkpoint_info)
+    sd_models.hf_auth_check(checkpoint_info)
+
+    load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config)
+    shared.log.debug(f'Load model: type=HunyuanDiT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+    transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config)
+    text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2")
+
+    pipe = diffusers.HunyuanDiTPipeline.from_pretrained(
+        repo_id,
+        transformer=transformer,
+        text_encoder_2=text_encoder_2,
+        cache_dir=shared.opts.diffusers_dir,
+        **load_args,
+    )
+
+    del text_encoder_2
+    del transformer
+    sd_hijack_te.init_hijack(pipe)
+    devices.torch_gc(force=True, reason='load')
+    return pipe