From 59efc95e0088e16194c548c3b3c45e00edaebe3a Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Mon, 14 Apr 2025 10:03:33 -0400
Subject: [PATCH] flux-cfgzero map autopipeline

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 modules/cfgzero/__init__.py | 6 ++++--
 modules/sd_hijack_te.py     | 2 +-
 modules/shared.py           | 1 +
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/modules/cfgzero/__init__.py b/modules/cfgzero/__init__.py
index f7d350ee3..8b9e51a3f 100644
--- a/modules/cfgzero/__init__.py
+++ b/modules/cfgzero/__init__.py
@@ -28,8 +28,12 @@ def apply(p: processing.StableDiffusionProcessing):
     orig_pipeline = shared.sd_model
 
     if cls == 'FluxPipeline':
+        from diffusers import pipelines
         from modules.cfgzero.flux_pipeline import FluxCFGZeroPipeline
         shared.sd_model = sd_models.switch_pipe(FluxCFGZeroPipeline, shared.sd_model)
+        pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["fluxcfgzero"] = FluxCFGZeroPipeline
+        pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["fluxcfgzero"] = pipelines.FluxImg2ImgPipeline
+        pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["fluxcfgzero"] = pipelines.FluxInpaintPipeline
     if cls == 'CogView4Pipeline':
         from modules.cfgzero.cogview4_pipeline import CogView4CFGZeroPipeline
         shared.sd_model = sd_models.switch_pipe(CogView4CFGZeroPipeline, shared.sd_model)
@@ -44,9 +48,7 @@ def apply(p: processing.StableDiffusionProcessing):
         shared.sd_model = sd_models.switch_pipe(WanCFGZeroPipeline, shared.sd_model)
     if cls == 'HunyuanVideoPipeline':
         from modules.cfgzero.hunyuan_t2v_pipeline import HunyuanVideoCFGZeroPipeline
-        from modules.model_hidream import init_hijack
         shared.sd_model = sd_models.switch_pipe(HunyuanVideoCFGZeroPipeline, shared.sd_model)
-        init_hijack(shared.sd_model)
 
     shared.log.debug(f'Apply CFGZero: cls={cls} init={shared.opts.cfgzero_enabled} star={shared.opts.cfgzero_star} steps={shared.opts.cfgzero_steps}')
     p.task_args['use_zero_init'] = shared.opts.cfgzero_enabled
diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py
index f80ac29ee..5653839f1 100644
--- a/modules/sd_hijack_te.py
+++ b/modules/sd_hijack_te.py
@@ -20,7 +20,7 @@ def hijack_encode_prompt(*args, **kwargs):
 
 
 def init_hijack(pipe):
-    if pipe is not None and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'):
+    if shared.opts.te_hijack and pipe is not None and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'):
         shared.log.debug(f'Model: cls={pipe.__class__.__name__} hijack encode')
         pipe.orig_encode_prompt = pipe.encode_prompt
         pipe.encode_prompt = hijack_encode_prompt
diff --git a/modules/shared.py b/modules/shared.py
index 61c7e72bd..8c51ac40d 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -443,6 +443,7 @@ options_templates.update(options_section(('text_encoder', "Text Encoder"), {
     "comma_padding_backtrack": OptionInfo(20, "Prompt padding", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1, "visible": not native }),
     "sd_textencder_linebreak": OptionInfo(True, "Use line break as prompt segment marker", gr.Checkbox),
     "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox),
+    "te_hijack": OptionInfo(True, "Offload after prompt encode", gr.Checkbox),
     "te_optional_sep": OptionInfo("<h2>Optional</h2>", "", gr.HTML),
     "te_pooled_embeds": OptionInfo(False, "SDXL: Use weighted pooled embeds"),
     "te_complex_human_instruction": OptionInfo(True, "Sana: Use complex human instructions"),