Merge pull request #4424 from CalamitousFelicitousness/z-image

Add Z-Image model support

modules/sd_detect.py

@@ -135,6 +135,8 @@ def guess_by_name(fn, current_guess):
         new_guess = 'PRX'
     elif 'gemini-' in fn.lower() and 'image' in fn.lower():
         new_guess = 'NanoBanana'
+    elif 'z-image' in fn.lower() or 'z_image' in fn.lower():
+        new_guess = 'Z-Image'
     if debug_load:
         shared.log.trace(f'Autodetect: method=name file="{fn}" previous="{current_guess}" current="{new_guess}"')
     return new_guess or current_guess

modules/sd_models.py

@@ -443,6 +443,10 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
             from pipelines.model_prx import load_prx
             sd_model = load_prx(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
+        elif model_type in ['Z-Image']:
+            from pipelines.model_z_image import load_z_image
+            sd_model = load_z_image(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:

modules/shared_items.py

@@ -47,6 +47,7 @@ pipelines = {
     'WanAI': getattr(diffusers, 'WanPipeline', None),
     'Qwen': getattr(diffusers, 'QwenImagePipeline', None),
     'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None),
+    'Z-Image': getattr(diffusers, 'ZImagePipeline', None),
     # dynamically imported and redefined later
     'Meissonic': getattr(diffusers, 'DiffusionPipeline', None),
     'Monetico': getattr(diffusers, 'DiffusionPipeline', None),

pipelines/model_z_image.py

@@ -0,0 +1,27 @@
+import diffusers
+from modules import shared, sd_models, sd_hijack_te, devices, model_quant
+from pipelines import generic
+
+
+def load_z_image(checkpoint_info, diffusers_load_config=None):
+    if diffusers_load_config is None:
+        diffusers_load_config = {}
+    repo_id = sd_models.path_to_repo(checkpoint_info)
+    sd_models.hf_auth_check(checkpoint_info)
+
+    shared.log.debug(f'Load model: type=Z-Image repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}')
+    transformer = generic.load_transformer(repo_id, cls_name=diffusers.ZImageTransformer2DModel, load_config=diffusers_load_config)
+
+    load_config, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+    pipe = diffusers.ZImagePipeline.from_pretrained(
+        repo_id,
+        cache_dir=shared.opts.diffusers_dir,
+        transformer=transformer,
+        **load_config,
+    )
+
+    del transformer
+    sd_hijack_te.init_hijack(pipe)
+
+    devices.torch_gc(force=True, reason='load')
+    return pipe