nodes/LucipherDev-ComfyUI-TangoFlux
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add TeaCache

Browse Source
This commit is contained in:
LucipherDev
2025-01-10 00:16:19 +05:30
parent 9c54679a1c
commit 8469da6a7d
5 changed files with 432 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+21 -1
View File
@@ -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}
}
```
example_workflow.json
+13 -7
View File
@@ -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
nodes.py
+106 -38
View File
@@ -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 {
pyproject.toml
+2 -2
View File
@@ -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"
tangoflux/model.py
+290 -9
View File
@@ -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()