diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index c4a2712f3..b01b847d8 100644 --- a/modules/sd_samplers_common.py +++ b/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: diff --git a/modules/teacache/teacache_ltx.py b/modules/teacache/teacache_ltx.py new file mode 100644 index 000000000..d0e3c942e --- /dev/null +++ b/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) diff --git a/scripts/ltxvideo.py b/scripts/ltxvideo.py index 7b12a3a86..e6464970d 100644 --- a/scripts/ltxvideo.py +++ b/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()