From 8469da6a7dc1ff562c26acba8f7ff7db524e2dc5 Mon Sep 17 00:00:00 2001 From: LucipherDev Date: Fri, 10 Jan 2025 00:16:19 +0530 Subject: [PATCH] Add TeaCache --- README.md | 22 +++- example_workflow.json | 20 ++- nodes.py | 144 ++++++++++++++------ pyproject.toml | 4 +- tangoflux/model.py | 299 ++++++++++++++++++++++++++++++++++++++++-- 5 files changed, 432 insertions(+), 57 deletions(-) diff --git a/README.md b/README.md index cc077c8..f663339 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ python install.py ## Example Workflow -![example_workflow](https://github.com/user-attachments/assets/3293cf59-d3bf-4d48-8e40-2b9a74ea035a) +![example_workflow](https://github.com/user-attachments/assets/afbf7b53-d712-4c9c-a538-53f0dc001f45) ## Usage @@ -39,6 +39,18 @@ python install.py The nodes can be found in "TangoFlux" category as `TangoFluxLoader`, `TangoFluxSampler`, `TangoFluxVAEDecodeAndPlay`. +![teacache_options](https://github.com/user-attachments/assets/29e676d9-902b-4ea2-9f72-18d3607996e8) + +> [TeaCache](https://github.com/LiewFeng/TeaCache) can speedup TangoFlux 2x without much audio quality degradation, in a training-free manner. +> +> +> ## 📈 Inference Latency Comparisons on a Single A800 +> +> +> | TangoFlux | TeaCache (0.25) | TeaCache (0.4) | +> |:-------------------:|:----------------------------:|:--------------------:| +> | ~4.08 s | ~2.42 s | ~1.95 s | + ## Citation ```bibtex @@ -52,3 +64,11 @@ The nodes can be found in "TangoFlux" category as `TangoFluxLoader`, `TangoFluxS url={https://arxiv.org/abs/2412.21037}, } ``` +``` +@article{liu2024timestep, + title={Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model}, + author={Liu, Feng and Zhang, Shiwei and Wang, Xiaofeng and Wei, Yujie and Qiu, Haonan and Zhao, Yuzhong and Zhang, Yingya and Ye, Qixiang and Wan, Fang}, + journal={arXiv preprint arXiv:2411.19108}, + year={2024} +} +``` diff --git a/example_workflow.json b/example_workflow.json index f693961..b784106 100644 --- a/example_workflow.json +++ b/example_workflow.json @@ -11,7 +11,7 @@ ], "size": [ 210, - 46 + 102 ], "flags": {}, "order": 0, @@ -38,7 +38,10 @@ "properties": { "Node name for S&R": "TangoFluxLoader" }, - "widgets_values": [] + "widgets_values": [ + false, + 0.25 + ] }, { "id": 13, @@ -115,7 +118,7 @@ 50, 3, 10, - 714702298733222, + 106139285587780, "randomize", 1 ] @@ -152,10 +155,13 @@ "extra": { "ds": { "scale": 0.9480295566502464, - "offset": { - "0": -200.8333282470703, - "1": -127.03429412841797 - } + "offset": [ + -200.83333333333337, + -102.2460379319304 + ] + }, + "node_versions": { + "comfyui-tangoflux": "1.0.2" } }, "version": 0.4 diff --git a/nodes.py b/nodes.py index f4090cf..38807fe 100644 --- a/nodes.py +++ b/nodes.py @@ -6,13 +6,13 @@ import torch import torchaudio import re -from diffusers import AutoencoderOobleck +from diffusers import AutoencoderOobleck, FluxTransformer2DModel from huggingface_hub import snapshot_download from comfy.utils import load_torch_file, ProgressBar import folder_paths -from .tangoflux.model import TangoFlux +from .tangoflux.model import TangoFlux, teacache_forward log = logging.getLogger("TangoFlux") @@ -27,11 +27,18 @@ folder_paths.folder_names_and_paths["tangoflux"] = ( ) TEXT_ENCODER_DIR = os.path.join(folder_paths.models_dir, "text_encoders") + class TangoFluxLoader: @classmethod def INPUT_TYPES(cls): return { - "required": {}, + "required": { + "enable_teacache": ("BOOLEAN", {"default": False}), + "rel_l1_thresh": ( + "FLOAT", + {"default": 0.25, "min": 0.0, "max": 10.0, "step": 0.01}, + ), + }, } RETURN_TYPES = ("TANGOFLUX_MODEL", "TANGOFLUX_VAE") @@ -41,14 +48,26 @@ class TangoFluxLoader: CATEGORY = "TangoFlux" FUNCTION = "load_tangoflux" DESCRIPTION = "Load TangoFlux model" - + def __init__(self): self.model = None self.vae = None + self.enable_teacache = False + self.rel_l1_thresh = 0.25 + self.original_forward = FluxTransformer2DModel.forward - def load_tangoflux(self, tangoflux_path=TANGOFLUX_DIR, text_encoder_path=TEXT_ENCODER_DIR, device="cuda"): - if None in [self.model, self.vae]: - pbar = ProgressBar(5) + def load_tangoflux( + self, + enable_teacache=False, + rel_l1_thresh=0.25, + tangoflux_path=TANGOFLUX_DIR, + text_encoder_path=TEXT_ENCODER_DIR, + device="cuda", + ): + if self.model is None or self.enable_teacache != enable_teacache: + del self.model + + pbar = ProgressBar(6) if not os.path.exists(tangoflux_path): log.info(f"Downloading TangoFlux models to: {tangoflux_path}") @@ -61,16 +80,20 @@ class TangoFluxLoader: pbar.update(1) - log.info(f"Loading config") + log.info("Loading config") with open(os.path.join(tangoflux_path, "config.json"), "r") as f: config = json.load(f) - + pbar.update(1) - - text_encoder = re.sub(r'[<>:"/\\|?*]', '-', config.get("text_encoder_name", "google/flan-t5-large")) + + text_encoder = re.sub( + r'[<>:"/\\|?*]', + "-", + config.get("text_encoder_name", "google/flan-t5-large"), + ) text_encoder_path = os.path.join(text_encoder_path, text_encoder) - + if not os.path.exists(text_encoder_path): log.info(f"Downloading text encoders to: {text_encoder_path}") snapshot_download( @@ -82,28 +105,57 @@ class TangoFluxLoader: pbar.update(1) - log.info(f"Loading TangoFlux models") + log.info("Loading TangoFlux models") model_weights = load_torch_file( - os.path.join(tangoflux_path, "tangoflux.safetensors"), device=torch.device(device) + os.path.join(tangoflux_path, "tangoflux.safetensors"), + device=torch.device(device), ) + + pbar.update(1) + + if enable_teacache: + log.info("Enabling TeaCache") + FluxTransformer2DModel.forward = teacache_forward + else: + log.info("Disabling TeaCache") + FluxTransformer2DModel.forward = self.original_forward + self.model = TangoFlux(config, text_encoder_path) + self.model.load_state_dict(model_weights, strict=False) self.model.to(device) - pbar.update(1) - - log.info(f"Loading TangoFlux VAE") - - vae_weights = load_torch_file( - os.path.join(tangoflux_path, "vae.safetensors") - ) - self.vae = AutoencoderOobleck() - self.vae.load_state_dict(vae_weights) - self.vae.to(device) + if enable_teacache: + self.model.transformer.__class__.enable_teacache = True + self.model.transformer.__class__.cnt = 0 + self.model.transformer.__class__.rel_l1_thresh = rel_l1_thresh + self.model.transformer.__class__.accumulated_rel_l1_distance = 0 + self.model.transformer.__class__.previous_modulated_input = None + self.model.transformer.__class__.previous_residual = None pbar.update(1) + + self.enable_teacache = enable_teacache + self.rel_l1_thresh = rel_l1_thresh + + if self.vae is None: + log.info("Loading TangoFlux VAE") + vae_weights = load_torch_file( + os.path.join(tangoflux_path, "vae.safetensors") + ) + self.vae = AutoencoderOobleck() + self.vae.load_state_dict(vae_weights) + self.vae.to(device) + + pbar.update(1) + + if self.enable_teacache == True and self.rel_l1_thresh != rel_l1_thresh: + self.model.transformer.__class__.rel_l1_thresh = rel_l1_thresh + + self.rel_l1_thresh = rel_l1_thresh + return (self.model, self.vae) @@ -134,14 +186,28 @@ class TangoFluxSampler: DESCRIPTION = "Sampler for TangoFlux" def sample( - self, model, prompt, steps=50, guidance_scale=3, duration=10, seed=0, batch_size=1, device="cuda" + self, + model, + prompt, + steps=50, + guidance_scale=3, + duration=10, + seed=0, + batch_size=1, + device="cuda", ): pbar = ProgressBar(steps) with torch.no_grad(): model.to(device) - - log.info(f"Generating latents with TangoFlux") + + try: + if model.transformer.__class__.enable_teacache: + model.transformer.__class__.num_steps = steps + except: + pass + + log.info("Generating latents with TangoFlux") latents = model( prompt, @@ -178,16 +244,16 @@ class TangoFluxVAEDecodeAndPlay: CATEGORY = "TangoFlux" FUNCTION = "play" DESCRIPTION = "Decoder and Player for TangoFlux" - + def decode(self, vae, latents): results = [] - + for latent in latents: decoded = vae.decode(latent.unsqueeze(0).transpose(2, 1)).sample.cpu() results.append(decoded) - + results = torch.cat(results, dim=0) - + return results def play( @@ -226,9 +292,9 @@ class TangoFluxVAEDecodeAndPlay: latents = latents["latents"] vae.to(device) - - log.info(f"Decoding Tangoflux latents") - + + log.info("Decoding Tangoflux latents") + waves = self.decode(vae, latents) pbar.update(1) @@ -236,15 +302,17 @@ class TangoFluxVAEDecodeAndPlay: for wave in waves: waveform_end = int(duration * vae.config.sampling_rate) wave = wave[:, :waveform_end] - + file = f"{filename}_{counter:05}_.{format}" - torchaudio.save(os.path.join(full_output_folder, file), wave, sample_rate=44100) - + torchaudio.save( + os.path.join(full_output_folder, file), wave, sample_rate=44100 + ) + counter += 1 audios.append({"filename": file, "subfolder": subfolder, "type": type}) - + pbar.update(1) return { diff --git a/pyproject.toml b/pyproject.toml index d1fcc00..d91a68c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,9 @@ [project] name = "comfyui-tangoflux" description = "ComfyUI Custom Nodes for 'TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching'. This generates high-quality 44.1kHz audio up to 30 seconds using just a text prompt." -version = "1.0.2" +version = "1.0.3" license = {file = "LICENSE"} -dependencies = ["torch", "torchaudio", "transformers", "diffusers", "accelerate", "tqdm"] +dependencies = ["torchaudio", "diffusers", "accelerate", "tqdm"] [project.urls] Repository = "https://github.com/LucipherDev/ComfyUI-TangoFlux" diff --git a/tangoflux/model.py b/tangoflux/model.py index 8b0dd7f..b44cec8 100644 --- a/tangoflux/model.py +++ b/tangoflux/model.py @@ -7,9 +7,285 @@ import copy import numpy as np from tqdm import tqdm -from typing import Optional, Union, List +from typing import Optional, Union, List, Any, Dict from math import pi import inspect +from diffusers.models import FluxTransformer2DModel +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.utils import ( + USE_PEFT_BACKEND, + is_torch_version, + scale_lora_layers, + unscale_lora_layers, +) + + +# Code from https://github.com/ali-vilab/TeaCache/blob/main/TeaCache4TangoFlux/teacache_tango_flux.py +def teacache_forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + pooled_projections: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, +) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + """ + The [`FluxTransformer2DModel`] forward method. + + Args: + hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`): + Input `hidden_states`. + encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected + from the embeddings of input conditions. + timestep ( `torch.LongTensor`): + Used to indicate denoising step. + block_controlnet_hidden_states: (`list` of `torch.Tensor`): + A list of tensors that if specified are added to the residuals of transformer blocks. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + if joint_attention_kwargs is not None: + joint_attention_kwargs = joint_attention_kwargs.copy() + lora_scale = joint_attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + else: + if ( + joint_attention_kwargs is not None + and joint_attention_kwargs.get("scale", None) is not None + ): + logger.warning( + "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." + ) + hidden_states = self.x_embedder(hidden_states) + + timestep = timestep.to(hidden_states.dtype) * 1000 + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + else: + guidance = None + temb = ( + self.time_text_embed(timestep, pooled_projections) + if guidance is None + else self.time_text_embed(timestep, guidance, pooled_projections) + ) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + ids = torch.cat((txt_ids, img_ids)) + image_rotary_emb = self.pos_embed(ids) + + if self.enable_teacache: + inp = hidden_states.clone() + temb_ = temb.clone() + modulated_inp, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.transformer_blocks[0].norm1(inp, emb=temb_) + ) + if self.cnt == 0 or self.cnt == self.num_steps - 1: + should_calc = True + self.accumulated_rel_l1_distance = 0 + else: + coefficients = [ + 4.98651651e02, + -2.83781631e02, + 5.58554382e01, + -3.82021401e00, + 2.64230861e-01, + ] + 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 index_block, block in enumerate(self.transformer_blocks): + if self.training 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 {} + ) + encoder_hidden_states, hidden_states = ( + torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + **ckpt_kwargs, + ) + ) + + else: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + for index_block, block in enumerate(self.single_transformer_blocks): + if self.training 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, + temb, + image_rotary_emb, + **ckpt_kwargs, + ) + + else: + hidden_states = block( + hidden_states=hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] + self.previous_residual = hidden_states - ori_hidden_states + else: + for index_block, block in enumerate(self.transformer_blocks): + if self.training 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 {} + ) + encoder_hidden_states, hidden_states = ( + torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + **ckpt_kwargs, + ) + ) + + else: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + for index_block, block in enumerate(self.single_transformer_blocks): + if self.training 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, + temb, + image_rotary_emb, + **ckpt_kwargs, + ) + + else: + hidden_states = block( + hidden_states=hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) class StableAudioPositionalEmbedding(nn.Module): @@ -150,8 +426,12 @@ class TangoFlux(nn.Module): self.noise_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) self.noise_scheduler_copy = copy.deepcopy(self.noise_scheduler) self.max_text_seq_len = 64 - self.text_encoder = T5EncoderModel.from_pretrained(text_encoder_dir if text_encoder_dir is not None else self.text_encoder_name) - self.tokenizer = T5TokenizerFast.from_pretrained(text_encoder_dir if text_encoder_dir is not None else self.text_encoder_name) + self.text_encoder = T5EncoderModel.from_pretrained( + text_encoder_dir if text_encoder_dir is not None else self.text_encoder_name + ) + self.tokenizer = T5TokenizerFast.from_pretrained( + text_encoder_dir if text_encoder_dir is not None else self.text_encoder_name + ) self.text_embedding_dim = self.text_encoder.config.d_model self.fc = nn.Sequential( @@ -270,6 +550,10 @@ class TangoFlux(nn.Module): """Only tested for single inference. Haven't test for batch inference""" torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True bsz = num_samples_per_prompt device = self.transformer.device @@ -322,10 +606,7 @@ class TangoFlux(nn.Module): txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(device) audio_ids = ( - torch.arange(self.audio_seq_len) - .unsqueeze(-1) - .repeat(1, 3) - .to(device) + torch.arange(self.audio_seq_len).unsqueeze(-1).repeat(1, 3).to(device) ) timesteps = timesteps.to(device) @@ -357,9 +638,9 @@ class TangoFlux(nn.Module): ) latents = scheduler.step(noise_pred, t, latents).prev_sample - + progress_bar.update(1) - + if callback_on_step_end is not None: callback_on_step_end()