general te hijack

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

scripts/allegrovideo.py

@@ -2,7 +2,7 @@ import time
 import gradio as gr
 import transformers
 import diffusers
-from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer
+from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer, sd_hijack_te
 
 
 repo_id = 'rhymes-ai/Allegro'
@@ -19,16 +19,6 @@ def hijack_decode(*args, **kwargs):
     return res
 
 
-def hijack_encode_prompt(*args, **kwargs):
-    t0 = time.time()
-    res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
-    t1 = time.time()
-    timer.process.add('te', t1-t0)
-    shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
-    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
-    return res
-
-
 class Script(scripts.Script):
     def title(self):
         return 'Video: Allegro (Legacy)'
@@ -94,9 +84,9 @@ class Script(scripts.Script):
             shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
             shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
             shared.sd_model.vae.decode = hijack_decode
-            shared.sd_model.encode_prompt = hijack_encode_prompt
             shared.sd_model.vae.enable_tiling()
             # shared.sd_model.vae.enable_slicing()
+            sd_hijack_te.init_hijack(shared.sd_model)
 
         shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
         devices.torch_gc(force=True)

scripts/hunyuanvideo.py

@@ -3,7 +3,7 @@ import torch
 import gradio as gr
 import transformers
 import diffusers
-from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, sd_samplers, model_quant, timer
+from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, sd_samplers, model_quant, timer, sd_hijack_te
 
 
 default_template = """Describe the video by detailing the following aspects:
@@ -48,16 +48,6 @@ def hijack_decode(*args, **kwargs):
     return res
 
 
-def hijack_encode_prompt(*args, **kwargs):
-    t0 = time.time()
-    res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
-    t1 = time.time()
-    timer.process.add('te', t1-t0)
-    shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
-    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
-    return res
-
-
 class Script(scripts.Script):
     def title(self):
         return 'Video: Hunyuan Video (Legacy)'
@@ -135,10 +125,10 @@ class Script(scripts.Script):
             shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
             shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
             shared.sd_model.vae.decode = hijack_decode
-            shared.sd_model.encode_prompt = hijack_encode_prompt
             shared.sd_model.vae.enable_slicing()
             shared.sd_model.vae.enable_tiling()
             shared.sd_model.vae.use_framewise_decoding = True
+            sd_hijack_te.init_hijack(shared.sd_model)
             loaded_model = model
 
     def run(self, p: processing.StableDiffusionProcessing, model, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument

scripts/legacy_allegrovideo.py

@@ -2,7 +2,7 @@ import time
 import gradio as gr
 import transformers
 import diffusers
-from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer
+from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer, sd_hijack_te
 
 
 repo_id = 'rhymes-ai/Allegro'
@@ -19,16 +19,6 @@ def hijack_decode(*args, **kwargs):
     return res
 
 
-def hijack_encode_prompt(*args, **kwargs):
-    t0 = time.time()
-    res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
-    t1 = time.time()
-    timer.process.add('te', t1-t0)
-    shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
-    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
-    return res
-
-
 class Script(scripts.Script):
     def title(self):
         return 'Video: Allegro (Legacy)'
@@ -94,8 +84,8 @@ class Script(scripts.Script):
             shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
             shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
             shared.sd_model.vae.decode = hijack_decode
-            shared.sd_model.encode_prompt = hijack_encode_prompt
             shared.sd_model.vae.enable_tiling()
+            sd_hijack_te.init_hijack(shared.sd_model)
             # shared.sd_model.vae.enable_slicing()
 
         shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)

scripts/ltxvideo.py

@@ -4,7 +4,7 @@ import torch
 import gradio as gr
 import diffusers
 import transformers
-from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer
+from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer, sd_hijack_te
 
 
 repos = {
@@ -39,16 +39,6 @@ def hijack_decode(*args, **kwargs):
     return res
 
 
-def hijack_encode_prompt(*args, **kwargs):
-    t0 = time.time()
-    res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
-    t1 = time.time()
-    timer.process.add('te', t1-t0)
-    shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
-    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
-    return res
-
-
 class Script(scripts.Script):
     def title(self):
         return 'Video: LTX Video (Legacy)'
@@ -131,9 +121,10 @@ class Script(scripts.Script):
             shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
             shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
             shared.sd_model.vae.decode = hijack_decode
-            shared.sd_model.encode_prompt = hijack_encode_prompt
             shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(repo_id)
             shared.sd_model.sd_model_hash = None
+            sd_hijack_te.init_hijack(shared.sd_model)
+
         shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
         shared.sd_model.vae.enable_slicing()
         shared.sd_model.vae.enable_tiling()