diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e8daac8d..0945143e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,18 +4,22 @@ - Video: add FasterCache and PAB support to WanDB and LTX models - ZLUDA: add more GPUs to recognized list -- LoRA: obey configured device when performing calculations - Progress: add additional fields to progress API - Progress: use batch-count for progress - Grid: add of max-rows and max-columns in settings to control grid format +- LoRA: add option to force-reload LoRA on every generate - Gallery: add max-columns in settings for gradio gallery components - Styles: resize and bring quick-ui to forward on hover -- Logging: fix debug logging - Logging: logging cleanup - Params: Reset default guidance-rescale from 0.7 to 0.0 -- Diag: add get-server-status to ui generate context menu +- Diag: add get-server-status to UI generate context menu - Pipe: [SoftFill](https://github.com/zacheryvaughn/softfill-pipelines) select in scripts, available for sdxl in inpaint model +- Flux: TeaCache for Flux.1 +- Fix: LoRA obey configured device when performing calculations +- Fix: ZLUDA startup +- Fix: balanced offload remove non-blocking move op +- Fix: debug logging ## Update for 2025-04-03 diff --git a/TODO.md b/TODO.md index 9c03982ed..92d24f497 100644 --- a/TODO.md +++ b/TODO.md @@ -17,6 +17,7 @@ N/A - Video: API support - Video: STG: - Video: SmoothCache: https://github.com/huggingface/diffusers/issues/11135 +- TeaCache: https://github.com/ali-vilab/TeaCache ## Code TODO diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index aad71c50e..faccc7d6e 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -139,6 +139,8 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers)] def changed(self, requested: List[str], include: List[str], exclude: List[str]): + if shared.opts.lora_force_reload: + return True sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) if not hasattr(sd_model, 'loaded_loras'): sd_model.loaded_loras = {} diff --git a/modules/model_flux.py b/modules/model_flux.py index 3dfe83ff2..71b5f1e96 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -223,6 +223,10 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch shared.sd_model = None devices.torch_gc(force=True) + if shared.opts.teacache_enabled: + from modules import teacache + diffusers.FluxTransformer2DModel.forward = teacache.teacache_forward + # load overrides if any if shared.opts.sd_unet != 'Default': try: diff --git a/modules/para_attention.py b/modules/para_attention.py index 5ca53962c..f5c6e8635 100644 --- a/modules/para_attention.py +++ b/modules/para_attention.py @@ -4,17 +4,17 @@ from modules import shared supported_models = ['Flux', 'HunyuanVideo', 'CogVideoX', 'Mochi'] -def apply_first_block_cache(p): +def apply_first_block_cache(): if not shared.opts.para_cache_enabled or not shared.native: return - if not any(p.sd_model.__class__.__name__.startswith(x) for x in supported_models): + if not any(shared.sd_model.__class__.__name__.startswith(x) for x in supported_models): return from installer import install install('para_attn') try: from para_attn.first_block_cache import diffusers_adapters - diffusers_adapters.apply_cache_on_pipe(p.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold) - shared.log.info(f'Applying para-attn first-block-cache: diff-threshold={shared.opts.para_diff_threshold} cls={p.sd_model.__class__.__name__}') + diffusers_adapters.apply_cache_on_pipe(shared.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold) + shared.log.info(f'Transformers cache: type=paraattn rdt={shared.opts.para_diff_threshold} cls={shared.sd_model.__class__.__name__}') except Exception as e: - shared.log.error(f'Applying para-attn first-block-cache: {e}') + shared.log.error(f'Transformers cache: type=paraattn {e}') return diff --git a/modules/processing.py b/modules/processing.py index 619b544dc..f12fe1215 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -169,9 +169,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed: shared.prompt_styles.extract_comments(p) if shared.opts.cuda_compile_backend == 'none': token_merge.apply_token_merging(p.sd_model) - from modules import sd_hijack_freeu, para_attention + from modules import sd_hijack_freeu, para_attention, teacache sd_hijack_freeu.apply_freeu(p, not shared.native) - para_attention.apply_first_block_cache(p) + para_attention.apply_first_block_cache() + teacache.apply_teacache(p) if p.width is not None: p.width = 8 * int(p.width / 8) diff --git a/modules/shared.py b/modules/shared.py index 8e84c71eb..8c7a3a268 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -560,13 +560,13 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "pag_apply_layers": OptionInfo("m0", "PAG layer names"), "pab_sep": OptionInfo("

PAB: Pyramid attention broadcast

", "", gr.HTML), - "pab_enabled": OptionInfo(False, "Attention cache enabled"), - "pab_spacial_skip_range": OptionInfo(2, "FC spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}), - "pab_spacial_skip_start": OptionInfo(100, "FC spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}), - "pab_spacial_skip_end": OptionInfo(800, "FC spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}), + "pab_enabled": OptionInfo(False, "PAB cache enabled"), + "pab_spacial_skip_range": OptionInfo(2, "PAB spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}), + "pab_spacial_skip_start": OptionInfo(100, "PAB spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}), + "pab_spacial_skip_end": OptionInfo(800, "PAB spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}), "faster_cache__sep": OptionInfo("

Faster Cache

", "", gr.HTML), - "faster_cache_enabled": OptionInfo(False, "Faster cache enabled"), + "faster_cache_enabled": OptionInfo(False, "FC cache enabled"), "fc_spacial_skip_range": OptionInfo(2, "FC spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}), "fc_spacial_skip_start": OptionInfo(0, "FC spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}), "fc_spacial_skip_end": OptionInfo(681, "FC spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.01}), @@ -581,6 +581,10 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "para_cache_enabled": OptionInfo(False, "First-block cache enabled"), "para_diff_threshold": OptionInfo(0.1, "Residual diff threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), + "teacache_sep": OptionInfo("

TeaCache

", "", gr.HTML), + "teacache_enabled": OptionInfo(False, "TC cache enabled"), + "teacache_thresh": OptionInfo(0.6, "TC L1 threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), + "hypertile_sep": OptionInfo("

HyperTile

", "", gr.HTML), "hypertile_unet_enabled": OptionInfo(False, "UNet Enabled"), "hypertile_hires_only": OptionInfo(False, "HiRes pass only"), @@ -932,6 +936,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "lora_fuse_diffusers": OptionInfo(True, "LoRA fuse directly to model"), "lora_apply_gpu": OptionInfo(False, "LoRA load directly on GPU"), "lora_legacy": OptionInfo(not native, "LoRA load using legacy method"), + "lora_force_reload": OptionInfo(False, "LoRA force reload always"), "lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA load using Diffusers method"), "lora_maybe_diffusers": OptionInfo(False, "LoRA load using Diffusers method for selected models"), "lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}), diff --git a/modules/teacache/__init__.py b/modules/teacache/__init__.py new file mode 100644 index 000000000..84fbd3d0f --- /dev/null +++ b/modules/teacache/__init__.py @@ -0,0 +1 @@ +from .teacache_flux import apply_teacache, teacache_forward diff --git a/modules/teacache/teacache_flux.py b/modules/teacache/teacache_flux.py new file mode 100644 index 000000000..8e57341fd --- /dev/null +++ b/modules/teacache/teacache_flux.py @@ -0,0 +1,322 @@ +from typing import Any, Dict, Optional, Union +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers +import torch +import numpy as np + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +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, + controlnet_block_samples=None, + controlnet_single_block_samples=None, + return_dict: bool = True, + controlnet_blocks_repeat: bool = False, + ) -> 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) + + if txt_ids.ndim == 3: + logger.warning( + "Passing `txt_ids` 3d torch.Tensor is deprecated." + "Please remove the batch dimension and pass it as a 2d torch Tensor" + ) + txt_ids = txt_ids[0] + if img_ids.ndim == 3: + logger.warning( + "Passing `img_ids` 3d torch.Tensor is deprecated." + "Please remove the batch dimension and pass it as a 2d torch Tensor" + ) + img_ids = img_ids[0] + + ids = torch.cat((txt_ids, img_ids), dim=0) + image_rotary_emb = self.pos_embed(ids) + + if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs: + ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds") + ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds) + joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states}) + + 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.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 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 torch.is_grad_enabled() and self.gradient_checkpointing: + + def create_custom_forward4(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_forward4(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, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if controlnet_block_samples is not None: + interval_control = len(self.transformer_blocks) / len(controlnet_block_samples) + interval_control = int(np.ceil(interval_control)) + # For Xlabs ControlNet. + if controlnet_blocks_repeat: + hidden_states = ( + hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)] + ) + else: + hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + for index_block, block in enumerate(self.single_transformer_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + + def create_custom_forward2(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_forward2(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, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if controlnet_single_block_samples is not None: + interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples) + interval_control = int(np.ceil(interval_control)) + hidden_states[:, encoder_hidden_states.shape[1] :, ...] = ( + hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + controlnet_single_block_samples[index_block // interval_control] + ) + + 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 torch.is_grad_enabled() and self.gradient_checkpointing: + + def create_custom_forward1(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_forward1(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, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if controlnet_block_samples is not None: + interval_control = len(self.transformer_blocks) / len(controlnet_block_samples) + interval_control = int(np.ceil(interval_control)) + # For Xlabs ControlNet. + if controlnet_blocks_repeat: + hidden_states = ( + hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)] + ) + else: + hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + for index_block, block in enumerate(self.single_transformer_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + + def create_custom_forward3(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_forward3(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, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if controlnet_single_block_samples is not None: + interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples) + interval_control = int(np.ceil(interval_control)) + hidden_states[:, encoder_hidden_states.shape[1] :, ...] = ( + hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + controlnet_single_block_samples[index_block // interval_control] + ) + + 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) + + +def apply_teacache(p): + from modules import shared + if not shared.native or not shared.opts.teacache_enabled or not shared.sd_model.__class__.__name__.startswith('Flux'): + return + shared.sd_model.transformer.__class__.enable_teacache = shared.opts.teacache_thresh > 0 + shared.sd_model.transformer.__class__.cnt = 0 + shared.sd_model.transformer.__class__.num_steps = p.steps + shared.sd_model.transformer.__class__.rel_l1_thresh = shared.opts.teacache_thresh # 0.25 for 1.5x speedup, 0.4 for 1.8x speedup, 0.6 for 2.0x speedup, 0.8 for 2.25x speedup + shared.sd_model.transformer.__class__.accumulated_rel_l1_distance = 0 + shared.sd_model.transformer.__class__.previous_modulated_input = None + shared.sd_model.transformer.__class__.previous_residual = None + shared.log.info(f'Transformers cache: type=teacache thresh={shared.opts.teacache_thresh} cls={shared.sd_model.__class__.__name__}') diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py index 33911afdf..cae9daef5 100644 --- a/scripts/xyz_grid_classes.py +++ b/scripts/xyz_grid_classes.py @@ -212,4 +212,5 @@ axis_options = [ AxisOption("[IY] Scale", float, apply_task_arg('infusenet_conditioning_scale')), AxisOption("[IY] Start", float, apply_task_arg('infusenet_guidance_start')), AxisOption("[IY] End", float, apply_task_arg('infusenet_guidance_end')), + AxisOption("[TeaCache] Threshold", float, apply_setting('teacache_thresh')), ]