teacache for ltxvideo

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

modules/sd_samplers_common.py

@@ -52,9 +52,12 @@ def single_sample_to_image(sample, approximation=None):
         if len(sample.shape) == 4 and sample.shape[0]: # likely animatediff latent
             sample = sample.permute(1, 0, 2, 3)[0]
         if approximation == 2: # TAESD
-            if shared.opts.live_preview_downscale and (sample.shape[-1] > 128 or sample.shape[-2] > 128):
-                scale = 128 / max(sample.shape[-1], sample.shape[-2])
-                sample = torch.nn.functional.interpolate(sample.unsqueeze(0), scale_factor=[scale, scale], mode='bilinear', align_corners=False)[0]
+            if (len(sample.shape) == 3 or len(sample.shape) == 4) and shared.opts.live_preview_downscale and (sample.shape[-1] > 128 or sample.shape[-2] > 128):
+                try:
+                    scale = 128 / max(sample.shape[-1], sample.shape[-2])
+                    sample = torch.nn.functional.interpolate(sample.unsqueeze(0), scale_factor=[scale, scale], mode='bilinear', align_corners=False)[0]
+                except Exception:
+                    pass
             x_sample = sd_vae_taesd.decode(sample)
             x_sample = (1.0 + x_sample) / 2.0 # preview requires smaller range
         elif shared.sd_model_type == 'sc' and approximation != 3:

modules/teacache/teacache_ltx.py

@@ -0,0 +1,163 @@
+from typing import Any, Dict, Optional, Tuple
+import numpy as np
+import torch
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import is_torch_version, scale_lora_layers, unscale_lora_layers
+
+
+def teacache_forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_attention_mask: torch.Tensor,
+        num_frames: int,
+        height: int,
+        width: int,
+        rope_interpolation_scale: Optional[Tuple[float, float, float]] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor:
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        scale_lora_layers(self, lora_scale)
+
+        image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        batch_size = hidden_states.size(0)
+        hidden_states = self.proj_in(hidden_states)
+
+        temb, embedded_timestep = self.time_embed(
+            timestep.flatten(),
+            batch_size=batch_size,
+            hidden_dtype=hidden_states.dtype,
+        )
+
+        temb = temb.view(batch_size, -1, temb.size(-1))
+        embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1))
+
+        encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
+
+        if self.enable_teacache:
+            inp = hidden_states.clone()
+            temb_ = temb.clone()
+            inp = self.transformer_blocks[0].norm1(inp)
+            num_ada_params = self.transformer_blocks[0].scale_shift_table.shape[0]
+            ada_values = self.transformer_blocks[0].scale_shift_table[None, None] + temb_.reshape(batch_size, temb_.size(1), num_ada_params, -1)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
+            modulated_inp = inp * (1 + scale_msa) + shift_msa
+            if self.cnt == 0 or self.cnt == self.num_steps-1:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+            else:
+                coefficients = [2.14700694e+01, -1.28016453e+01,  2.31279151e+00,  7.92487521e-01, 9.69274326e-03]
+                rescale_func = np.poly1d(coefficients)
+                self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+                if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                    should_calc = False
+                else:
+                    should_calc = True
+                    self.accumulated_rel_l1_distance = 0
+            self.previous_modulated_input = modulated_inp
+            self.cnt += 1
+            if self.cnt == self.num_steps:
+                self.cnt = 0
+
+        if self.enable_teacache:
+            if not should_calc:
+                hidden_states += self.previous_residual
+            else:
+                ori_hidden_states = hidden_states.clone()
+                for block in self.transformer_blocks:
+                    if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                        def create_custom_forward(module, return_dict=None):
+                            def custom_forward(*inputs):
+                                if return_dict is not None:
+                                    return module(*inputs, return_dict=return_dict)
+                                else:
+                                    return module(*inputs)
+
+                            return custom_forward
+
+                        ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            hidden_states,
+                            encoder_hidden_states,
+                            temb,
+                            image_rotary_emb,
+                            encoder_attention_mask,
+                            **ckpt_kwargs,
+                        )
+                    else:
+                        hidden_states = block(
+                            hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                            temb=temb,
+                            image_rotary_emb=image_rotary_emb,
+                            encoder_attention_mask=encoder_attention_mask,
+                        )
+
+                scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
+                shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+
+                hidden_states = self.norm_out(hidden_states)
+                hidden_states = hidden_states * (1 + scale) + shift
+                self.previous_residual = hidden_states - ori_hidden_states
+        else:
+            for block in self.transformer_blocks:
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        image_rotary_emb,
+                        encoder_attention_mask,
+                        **ckpt_kwargs,
+                    )
+                else:
+                    hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        encoder_attention_mask=encoder_attention_mask,
+                    )
+
+            scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
+            shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+
+            hidden_states = self.norm_out(hidden_states)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+
+        output = self.proj_out(hidden_states)
+
+        unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)

scripts/ltxvideo.py

@@ -5,6 +5,7 @@ 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.teacache.teacache_ltx import teacache_forward
 
 
 repos = {
@@ -83,6 +84,9 @@ class Script(scripts.Script):
         with gr.Row():
             num_frames = gr.Slider(label='Frames', minimum=9, maximum=257, step=1, value=41)
             sampler = gr.Checkbox(label='Override sampler', value=True)
+        with gr.Row():
+            teacache_enable = gr.Checkbox(label='Enable TeaCache', value=False)
+            teacache_threshold = gr.Slider(label='Threshold', minimum=0.01, maximum=0.1, step=0.01, value=0.03)
         with gr.Row():
             model_custom = gr.Textbox(value='', label='Path to model file', visible=False)
         with gr.Row():
@@ -94,9 +98,9 @@ class Script(scripts.Script):
             mp4_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=False)
         video_type.change(fn=video_type_change, inputs=[video_type], outputs=[duration, gif_loop, mp4_pad, mp4_interpolate])
         model.change(fn=model_change, inputs=[model], outputs=[model_custom])
-        return [model, model_custom, decode, sampler, num_frames, video_type, duration, gif_loop, mp4_pad, mp4_interpolate]
+        return [model, model_custom, decode, sampler, num_frames, video_type, duration, gif_loop, mp4_pad, mp4_interpolate, teacache_enable, teacache_threshold]
 
-    def run(self, p: processing.StableDiffusionProcessing, model, model_custom, decode, sampler, num_frames, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument
+    def run(self, p: processing.StableDiffusionProcessing, model, model_custom, decode, sampler, num_frames, video_type, duration, gif_loop, mp4_pad, mp4_interpolate, teacache_enable, teacache_threshold): # pylint: disable=arguments-differ, unused-argument
         # set params
         image = getattr(p, 'init_images', None)
         image = None if image is None or len(image) == 0 else image[0]
@@ -130,6 +134,7 @@ class Script(scripts.Script):
             kwargs = {}
             kwargs = model_quant.create_bnb_config(kwargs)
             kwargs = model_quant.create_ao_config(kwargs)
+            diffusers.LTXVideoTransformer3DModel.forward = teacache_forward
             if os.path.isfile(repo_id):
                 shared.sd_model = cls.from_single_file(
                     repo_id,
@@ -156,7 +161,16 @@ class Script(scripts.Script):
         shared.sd_model.vae.enable_slicing()
         shared.sd_model.vae.enable_tiling()
         devices.torch_gc(force=True)
-        shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} args={p.task_args}')
+
+        shared.sd_model.transformer.cnt = 0
+        shared.sd_model.transformer.accumulated_rel_l1_distance = 0
+        shared.sd_model.transformer.previous_modulated_input = None
+        shared.sd_model.transformer.previous_residual = None
+        shared.sd_model.transformer.enable_teacache = teacache_enable
+        shared.sd_model.transformer.rel_l1_thresh = teacache_threshold
+        shared.sd_model.transformer.num_steps = p.steps
+
+        shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} args={p.task_args} steps={p.steps} teacache={teacache_enable} threshold={teacache_threshold}')
 
         # run processing
         t0 = time.time()