From 6cf445d317ce8792485139171b986b32380743ef Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 18 Feb 2025 18:47:42 -0500 Subject: [PATCH] add ras-sd35 experimental Signed-off-by: Vladimir Mandic --- .pylintrc | 3 +- .ruff.toml | 3 +- CHANGELOG.md | 5 +- modules/interrogate/deepseek.py | 2 +- modules/interrogate/interrogate.py | 2 +- modules/interrogate/joytag.py | 2 +- modules/interrogate/openclip.py | 10 +- modules/masking.py | 1 - modules/postprocess/dcc.py | 4 +- modules/processing_diffusers.py | 4 +- modules/prompt_parser.py | 2 +- modules/ras/__init__.py | 36 ++++ modules/ras/ras_attention.py | 257 +++++++++++++++++++++++++ modules/ras/ras_forward.py | 139 +++++++++++++ modules/ras/ras_manager.py | 104 ++++++++++ modules/ras/ras_scheduler.py | 215 +++++++++++++++++++++ modules/schedulers/scheduler_ufogen.py | 2 +- modules/shared.py | 3 + 18 files changed, 777 insertions(+), 17 deletions(-) create mode 100644 modules/ras/__init__.py create mode 100644 modules/ras/ras_attention.py create mode 100644 modules/ras/ras_forward.py create mode 100644 modules/ras/ras_manager.py create mode 100644 modules/ras/ras_scheduler.py diff --git a/.pylintrc b/.pylintrc index 6395e3422..b472d26a0 100644 --- a/.pylintrc +++ b/.pylintrc @@ -23,12 +23,14 @@ ignore-paths=/usr/lib/.*$, modules/k-diffusion, modules/ldsr, modules/meissonic, + modules/mod, modules/omnigen, modules/onnx_impl, modules/pag, modules/pixelsmith, modules/prompt_parser_xhinker.py, modules/pulid/eva_clip, + modules/ras, modules/rife, modules/schedulers, modules/taesd, @@ -36,7 +38,6 @@ ignore-paths=/usr/lib/.*$, modules/todo, modules/unipc, modules/xadapter, - modules/mod, repositories, extensions-builtin/Lora, extensions-builtin/sd-webui-agent-scheduler, diff --git a/.ruff.toml b/.ruff.toml index 1cca9cebb..8e3d13e64 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -17,12 +17,14 @@ exclude = [ "modules/k-diffusion", "modules/ldsr", "modules/meissonic", + "modules/mod", "modules/omnigen", "modules/pag", "modules/pixelsmith", "modules/postprocess/aurasr_arch.py", "modules/prompt_parser_xhinker.py", "modules/pulid/eva_clip", + "modules/ras", "modules/rife", "modules/schedulers", "modules/segmoe", @@ -31,7 +33,6 @@ exclude = [ "modules/todo", "modules/unipc", "modules/xadapter", - "modules/mod", "repositories", "extensions-builtin/Lora", "extensions-builtin/sd-extension-chainner/nodes", diff --git a/CHANGELOG.md b/CHANGELOG.md index 425a05460..47861bf0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,13 +50,16 @@ We're back with another update with nearly 100 commits! result is model smaller than *Flux.1-Dev*, but with similar capabilities fully supports offloading and on-the-fly quantization simply select from *networks -> models -> reference* -- **Pipelines** +- **Functions** - [Mixture-of-Diffusers](https://huggingface.co/posts/elismasilva/251775641926329) Regional tiling type of a solution for SDXL models select from *scripts -> mixture of diffusers* - [Automatic Color Inpaint] Automatically creates mask based on selected color and triggers inpaint simply select in *scripts -> automatic color inpaint* when in img2img mode + - [RAS: Region-Adaptive Sampling](https://github.com/microsoft/RAS) *experimental* + Speeds up SD3.5 models by sampling only regions of interest + Enable in *settings -> pipeline modifiers -> ras* - **Interrogate/Captioning** - Redesigned captioning UI split from Process tab into separate tab diff --git a/modules/interrogate/deepseek.py b/modules/interrogate/deepseek.py index 892f269dc..5138c5693 100644 --- a/modules/interrogate/deepseek.py +++ b/modules/interrogate/deepseek.py @@ -20,7 +20,7 @@ vl_gpt = None vl_chat_processor = None -class fake_attrdict(object): +class fake_attrdict(): class AttrDict(dict): # dot notation access to dictionary attributes __getattr__ = dict.get __setattr__ = dict.__setitem__ diff --git a/modules/interrogate/interrogate.py b/modules/interrogate/interrogate.py index 45d485cef..f68212d91 100644 --- a/modules/interrogate/interrogate.py +++ b/modules/interrogate/interrogate.py @@ -28,7 +28,7 @@ def interrogate(image): elif shared.opts.interrogate_default_type == 'VLM': shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} vlm="{shared.opts.interrogate_vlm_model}" prompt="{shared.opts.interrogate_vlm_prompt}"') from modules.interrogate import vqa - prompt = vqa.interrogate(image=image, model_name=shared.opts.interrogate_vlm_model, question=shared.opts.interrogate_vlm_prompt) + prompt = vqa.interrogate(image=image, model_name=shared.opts.interrogate_vlm_model, question=shared.opts.interrogate_vlm_prompt, prompt=None) shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') return prompt else: diff --git a/modules/interrogate/joytag.py b/modules/interrogate/joytag.py index 075d002e4..ae48cb9c6 100644 --- a/modules/interrogate/joytag.py +++ b/modules/interrogate/joytag.py @@ -896,7 +896,7 @@ class CNNStem(nn.Module): ty, line = line.split(":") if ":" in line else (line, "") options = line.split(",") options = [o.split("=") for o in options] if line else [] - options = {k: v for k, v in options} # noqa: C416 + options = {k: v for k, v in options} # pylint: disable=unnecessary-comprehension # noqa: C416 if ty == 'conv': layers.append(nn.Conv2d( in_channels=channels, diff --git a/modules/interrogate/openclip.py b/modules/interrogate/openclip.py index b8b43092c..6d3656fab 100644 --- a/modules/interrogate/openclip.py +++ b/modules/interrogate/openclip.py @@ -239,11 +239,11 @@ class BatchWriter: def update_interrogate_params(): if ci is not None: - ci.caption_max_length=shared.opts.interrogate_clip_max_length, - ci.chunk_size=shared.opts.interrogate_clip_chunk_size, - ci.flavor_intermediate_count=shared.opts.interrogate_clip_flavor_count, - ci.clip_offload=shared.opts.interrogate_offload, - ci.caption_offload=shared.opts.interrogate_offload, + ci.caption_max_length=shared.opts.interrogate_clip_max_length + ci.chunk_size=shared.opts.interrogate_clip_chunk_size + ci.flavor_intermediate_count=shared.opts.interrogate_clip_flavor_count + ci.clip_offload=shared.opts.interrogate_offload + ci.caption_offload=shared.opts.interrogate_offload def get_clip_models(): diff --git a/modules/masking.py b/modules/masking.py index 8364e9369..f9af9106a 100644 --- a/modules/masking.py +++ b/modules/masking.py @@ -408,7 +408,6 @@ def run_mask(input_image: Image.Image, input_mask: Image.Image = None, return_ty debug(f'Mask args legacy: blur={mask_blur} padding={mask_padding}') if invert is not None: opts.invert = invert - print('HERE', opts.mask_blur, mask_blur) if mask_blur is not None: # compatibility with old img2img values which uses px values opts.mask_blur = round(4 * mask_blur / size, 3) if mask_padding is not None: # compatibility with old img2img values which uses px values diff --git a/modules/postprocess/dcc.py b/modules/postprocess/dcc.py index 7b9112f75..8e76d1bb9 100644 --- a/modules/postprocess/dcc.py +++ b/modules/postprocess/dcc.py @@ -1,8 +1,8 @@ import numpy as np -def DetectDirect(A, type, k, T): - if type == 1: +def DetectDirect(A, dcc_type, k, T): + if dcc_type == 1: # 45 degree diagonal direction t1 = abs(A[2,0]-A[0,2]) t2 = abs(A[4,0]-A[2,2])+abs(A[2,2]-A[0,4]) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index b2a7034d9..06d5f7eb2 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -5,7 +5,7 @@ import numpy as np import torch import torchvision.transforms.functional as TF from PIL import Image -from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, hidiffusion, timer, modelstats, extra_networks +from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, hidiffusion, timer, modelstats, extra_networks, ras from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name from modules.processing_args import set_pipeline_args from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed @@ -93,6 +93,7 @@ def process_base(p: processing.StableDiffusionProcessing): sd_models.move_model(shared.sd_model.transformer, devices.device) extra_networks.activate(p, exclude=['text_encoder', 'text_encoder_2', 'text_encoder_3']) hidiffusion.apply(p, shared.sd_model_type) + ras.apply(shared.sd_model, p) timer.process.record('move') if hasattr(shared.sd_model, 'tgate') and getattr(p, 'gate_step', -1) > 0: base_args['gate_step'] = p.gate_step @@ -106,6 +107,7 @@ def process_base(p: processing.StableDiffusionProcessing): if hasattr(output, 'images'): shared.history.add(output.images, info=processing.create_infotext(p), ops=p.ops) timer.process.record('pipeline') + ras.unapply(shared.sd_model) hidiffusion.unapply() sd_models_compile.openvino_post_compile(op="base") # only executes on compiled vino models sd_models_compile.check_deepcache(enable=False) diff --git a/modules/prompt_parser.py b/modules/prompt_parser.py index a5eb3ade2..43330c8aa 100644 --- a/modules/prompt_parser.py +++ b/modules/prompt_parser.py @@ -354,7 +354,7 @@ def parse_prompt_attention(text): # log.trace(f'Prompt: text="{text[m.start():m.end()]}" section="{section}" weight="{weight}"') if len(section) == 0: continue - elif section.startswith('\\'): + if section.startswith('\\'): if len(res) > 0 and text[m.start()-1] != ' ': res[-1][0] += section[1:] # append literal character to the last section else: diff --git a/modules/ras/__init__.py b/modules/ras/__init__.py new file mode 100644 index 000000000..891dae240 --- /dev/null +++ b/modules/ras/__init__.py @@ -0,0 +1,36 @@ +# source +# original: + +from modules import shared, processing + + +def apply(pipe, p: processing.StableDiffusionProcessing): + if shared.sd_model_type != "sd3" or not shared.opts.ras_enable: + return + from .ras_manager import MANAGER + from .ras_scheduler import RASFlowMatchEulerDiscreteScheduler + from .ras_attention import RASJointAttnProcessor2_0 + from .ras_forward import ras_forward + scheduler = RASFlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config) + pipe.scheduler = scheduler + MANAGER.num_steps = p.steps + MANAGER.scheduler_end_step = p.steps + MANAGER.width = p.width + MANAGER.height = p.height + MANAGER.error_reset_steps = [int(1*p.steps/3), int(2*p.steps/3)] + shared.log.info(f'RAS: scheduler={pipe.scheduler.__class__.__name__} {str(MANAGER)}') + MANAGER.reset_cache() + MANAGER.generate_skip_token_list() + pipe.transformer.old_forward = pipe.transformer.forward + pipe.transformer.forward = ras_forward.__get__(pipe.transformer, pipe.transformer.__class__) # pylint: disable=no-value-for-parameter + for block in pipe.transformer.transformer_blocks: + block.attn.set_processor(RASJointAttnProcessor2_0()) + + +def unapply(pipe): + if hasattr(pipe, 'transformer') and hasattr(pipe.transformer, "old_forward"): + from diffusers.models.attention_processor import JointAttnProcessor2_0 + pipe.transformer.forward = pipe.transformer.old_forward + del pipe.transformer.old_forward + for block in pipe.transformer.transformer_blocks: + block.attn.set_processor(JointAttnProcessor2_0()) diff --git a/modules/ras/ras_attention.py b/modules/ras/ras_attention.py new file mode 100644 index 000000000..ca3a083e2 --- /dev/null +++ b/modules/ras/ras_attention.py @@ -0,0 +1,257 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional +import math +import torch +import torch.nn.functional as F +from diffusers.models.attention_processor import Attention +from . import ras_manager + + +class RASLuminaAttnProcessor2_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is + used in the LuminaNextDiT model. It applies a s normalization layer and rotary embedding on query and key vector. + """ + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + if ras_manager.MANAGER.sample_ratio < 1.0: + self.k_cache = None + self.v_cache = None + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + query_rotary_emb: Optional[torch.Tensor] = None, + key_rotary_emb: Optional[torch.Tensor] = None, + base_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + from diffusers.models.embeddings import apply_rotary_emb + + is_self_attention = True if hidden_states.shape == encoder_hidden_states.shape else False + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + batch_size, sequence_length, _ = hidden_states.shape + + # Get Query-Key-Value Pair + query = attn.to_q(hidden_states) + value = attn.to_v(encoder_hidden_states) + key = attn.to_k(encoder_hidden_states) + + query_dim = query.shape[-1] + inner_dim = key.shape[-1] + head_dim = query_dim // attn.heads + dtype = query.dtype + + # Get key-value heads + kv_heads = inner_dim // head_dim + + # Apply Query-Key Norm if needed + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + query = query.view(batch_size, -1, attn.heads, head_dim) + key = key.view(batch_size, -1, kv_heads, head_dim) + value = value.view(batch_size, -1, kv_heads, head_dim) + + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.current_step == 0 and is_self_attention: + self.k_cache = None + self.v_cache = None + + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.current_step > ras_manager.MANAGER.scheduler_end_step and is_self_attention: + self.k_cache = None + self.v_cache = None + + # Apply RoPE if needed + if query_rotary_emb is not None: + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_RAS_step: + query = apply_rotary_emb(query, ras_manager.MANAGER.image_rotary_emb_skip, use_real=False) + else: + query = apply_rotary_emb(query, query_rotary_emb, use_real=False) + if key_rotary_emb is not None: + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_RAS_step: + key = apply_rotary_emb(key, ras_manager.MANAGER.image_rotary_emb_skip, use_real=False) + else: + key = apply_rotary_emb(key, key_rotary_emb, use_real=False) + + if ras_manager.MANAGER.sample_ratio < 1.0 and (ras_manager.MANAGER.current_step == ras_manager.MANAGER.scheduler_start_step - 1 or ras_manager.MANAGER.current_step in ras_manager.MANAGER.error_reset_steps) and is_self_attention: + self.k_cache = key + self.v_cache = value + + if ras_manager.MANAGER.sample_ratio < 1.0 and is_self_attention and ras_manager.MANAGER.is_RAS_step: + self.k_cache[:, ras_manager.MANAGER.other_patchified_index] = key + self.v_cache[:, ras_manager.MANAGER.other_patchified_index] = value + key = self.k_cache + value = self.v_cache + + query, key = query.to(dtype), key.to(dtype) + + if ras_manager.MANAGER.sample_ratio < 1.0 and is_self_attention and ras_manager.MANAGER.is_RAS_step: + if is_self_attention: + sequence_length = key.shape[1] + else: + sequence_length = base_sequence_length + + + # Apply proportional attention if true + if key_rotary_emb is None: + softmax_scale = None + else: + if base_sequence_length is not None: + softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale + else: + softmax_scale = attn.scale + + # perform Grouped-qurey Attention (GQA) # TODO replace with GQA + n_rep = attn.heads // kv_heads + if n_rep >= 1: + key = key.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) + value = value.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) + + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.bool().view(batch_size, 1, 1, -1) + if ras_manager.MANAGER.sample_ratio < 1.0 and is_self_attention and ras_manager.MANAGER.is_RAS_step: + attention_mask = attention_mask.expand(-1, attn.heads, query.shape[1], -1) + else: + attention_mask = attention_mask.expand(-1, attn.heads, sequence_length, -1) + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, scale=softmax_scale + ) + hidden_states = hidden_states.transpose(1, 2).to(dtype) + return hidden_states + + +class RASJointAttnProcessor2_0: + """Attention processor used typically in processing the SD3-like self-attention projections.""" + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + if ras_manager.MANAGER.sample_ratio < 1.0: + self.k_cache = None + self.v_cache = None + + def __call__( + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + attention_mask: Optional[torch.FloatTensor] = None, + *args, + **kwargs, + ) -> torch.FloatTensor: + residual = hidden_states + + batch_size = hidden_states.shape[0] + + # `sample` projections. + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.current_step == 0: + self.k_cache = None + self.v_cache = None + + if ras_manager.MANAGER.sample_ratio < 1.0 and (ras_manager.MANAGER.current_step == ras_manager.MANAGER.scheduler_start_step - 1 or ras_manager.MANAGER.current_step in ras_manager.MANAGER.error_reset_steps): + self.k_cache = key + self.v_cache = value + + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_RAS_step: + self.k_cache[:, :, ras_manager.MANAGER.other_patchified_index] = key + self.v_cache[:, :, ras_manager.MANAGER.other_patchified_index] = value + key = self.k_cache + value = self.v_cache + + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.current_step > ras_manager.MANAGER.scheduler_end_step: + self.k_cache = None + self.v_cache = None + + # `context` projections. + if encoder_hidden_states is not None: + encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) + encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) + encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) + + encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( + batch_size, -1, attn.heads, head_dim + ).transpose(1, 2) + encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( + batch_size, -1, attn.heads, head_dim + ).transpose(1, 2) + encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( + batch_size, -1, attn.heads, head_dim + ).transpose(1, 2) + + if attn.norm_added_q is not None: + encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) + if attn.norm_added_k is not None: + encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) + + query = torch.cat([query, encoder_hidden_states_query_proj], dim=2) + key = torch.cat([key, encoder_hidden_states_key_proj], dim=2) + value = torch.cat([value, encoder_hidden_states_value_proj], dim=2) + + hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + if encoder_hidden_states is not None: + # Split the attention outputs. + hidden_states, encoder_hidden_states = ( + hidden_states[:, : residual.shape[1]], + hidden_states[:, residual.shape[1] :], + ) + if not attn.context_pre_only: + encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + if encoder_hidden_states is not None: + return hidden_states, encoder_hidden_states + else: + return hidden_states diff --git a/modules/ras/ras_forward.py b/modules/ras/ras_forward.py new file mode 100644 index 000000000..63c71428e --- /dev/null +++ b/modules/ras/ras_forward.py @@ -0,0 +1,139 @@ +# Copyright 2024 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, List, Optional, Union +import torch +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers +from . import ras_manager + + +def ras_forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + pooled_projections: torch.FloatTensor = None, + timestep: torch.LongTensor = None, + block_controlnet_hidden_states: List = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + skip_layers: Optional[List[int]] = None, + ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + """ + The [`SD3Transformer2DModel`] 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. + skip_layers (`list` of `int`, *optional*): + A list of layer indices to skip during the forward pass. + + 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 + + scale_lora_layers(self, lora_scale) + + height, width = hidden_states.shape[-2:] + + hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too. + temb = self.time_text_embed(timestep, pooled_projections) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_RAS_step: + hidden_states = hidden_states[:, ras_manager.MANAGER.other_patchified_index] + + 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, ip_temb = self.image_proj(ip_adapter_image_embeds, timestep) + + joint_attention_kwargs.update(ip_hidden_states=ip_hidden_states, temb=ip_temb) + + for index_block, block in enumerate(self.transformer_blocks): + # Skip specified layers + is_skip = True if skip_layers is not None and index_block in skip_layers else False + + if torch.is_grad_enabled() and self.gradient_checkpointing and not is_skip: + encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + joint_attention_kwargs, + ) + elif not is_skip: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if block_controlnet_hidden_states is not None and block.context_pre_only is False: + interval_control = len(self.transformer_blocks) / len(block_controlnet_hidden_states) + hidden_states = hidden_states + block_controlnet_hidden_states[int(index_block / interval_control)] + + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # unpatchify + patch_size = self.config.patch_size + height = height // patch_size + width = width // patch_size + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_RAS_step: + final_hidden_states = torch.zeros( + (hidden_states.shape[0], height * width, hidden_states.shape[2]), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + final_hidden_states[:, ras_manager.MANAGER.other_patchified_index] = hidden_states + hidden_states = final_hidden_states + + hidden_states = hidden_states.reshape( + shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size) + ) + + 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) diff --git a/modules/ras/ras_manager.py b/modules/ras/ras_manager.py new file mode 100644 index 000000000..13196e00e --- /dev/null +++ b/modules/ras/ras_manager.py @@ -0,0 +1,104 @@ +class ras_manager: + def __init__(self): + ## configurable + self.metric = "std" + self.patch_size = 2 + self.scheduler_start_step = 4 + self.sample_ratio = 0.5 + self.starvation_scale = 1.0 + self.vae_size = 8 + self.high_ratio = 0.3 + self.skip_num_step = [] + self.skip_num_step_length = 0 + + # applied by sdnext pipeline in ras/__init__.py + self.scheduler_end_step = 0 + self.error_reset_steps = [0, 0] + self.num_steps = 0 + self.height = 0 + self.width = 0 + + ## dynamic + self.current_step = 0 + self.is_RAS_step = False + self.is_next_RAS_step = False + self.cached_index = None + self.other_index = None + self.cached_patchified_index = None + self.other_patchified_index = None + self.image_rotary_emb_skip = None + self.cached_scaled_noise = None + self.skip_token_num_list = [] + + def __str__(self): + return f'steps={self.num_steps} start={self.scheduler_start_step} end={self.scheduler_end_step} patch={self.patch_size} metric={self.metric} reset={self.error_reset_steps} ratio={self.sample_ratio} starvation={self.starvation_scale} vae={self.vae_size} high={self.high_ratio} skip={self.skip_num_step}length={self.skip_num_step_length}' + + def set_parameters(self, args): + self.patch_size = args.patch_size + self.scheduler_start_step = args.scheduler_start_step + self.scheduler_end_step = args.scheduler_end_step + self.metric = args.metric + self.error_reset_steps = [int(i.strip()) for i in args.error_reset_steps.split(",")] + self.sample_ratio = args.sample_ratio + self.num_steps = args.num_inference_steps + self.skip_num_step = args.skip_num_step + self.skip_num_step_length = args.skip_num_step_length + self.height = args.height + self.width = args.width + self.high_ratio = args.high_ratio + self.generate_skip_token_list() + + + def generate_skip_token_list(self): + avg_skip_token_num = int((1 - self.sample_ratio) * ((self.height // self.patch_size) // self.vae_size) * ((self.width // self.patch_size) // self.vae_size)) + if self.skip_num_step_length == 0: # static dropping + self.skip_token_num_list = [avg_skip_token_num for i in range(self.num_steps)] + for i in self.error_reset_steps: + self.skip_token_num_list[i] = 0 + for i in range(self.scheduler_start_step): + self.skip_token_num_list[i] = 0 + return + for i in range(0, self.num_steps // self.skip_num_step_length + 1): + for j in range(self.skip_num_step_length): + if i * self.skip_num_step_length + j >= self.num_steps: + break + temp_skip_num = avg_skip_token_num + self.skip_num_step * (i - (((self.num_steps + self.scheduler_start_step) // self.skip_num_step_length) // 2)) + temp_skip_num = (temp_skip_num // 64) * 64 + self.skip_token_num_list.append(temp_skip_num) + for i in range(self.scheduler_start_step): + self.skip_token_num_list[i] = 0 + for i in self.error_reset_steps: + self.skip_token_num_list[i] = 0 + for i in range(len(self.skip_token_num_list)): + assert self.skip_token_num_list[i] >= 0, "Skip token number should be positive" + assert self.skip_token_num_list[i] <= ((self.height // self.patch_size) // self.vae_size) * ((self.width // self.patch_size) // self.vae_size) + + def reset_cache(self): + self.cached_index = None + self.other_index = None + self.cached_patchified_index = None + self.other_patchified_index = None + self.image_rotary_emb_skip = None + self.cached_scaled_noise = None + self.current_step = 0 + if self.current_step >= self.scheduler_start_step and self.current_step <= self.scheduler_end_step and self.current_step not in self.error_reset_steps: + self.is_RAS_step = True + else: + self.is_RAS_step = False + if self.current_step + 1 >= self.scheduler_start_step and self.current_step + 1 <= self.scheduler_end_step and self.current_step + 1 not in self.error_reset_steps: + self.is_next_RAS_step = True + else: + self.is_next_RAS_step = False + + def increase_step(self): + self.current_step += 1 + if self.current_step >= self.scheduler_start_step and self.current_step <= self.scheduler_end_step and self.current_step not in self.error_reset_steps: + self.is_RAS_step = True + else: + self.is_RAS_step = False + if self.current_step + 1 >= self.scheduler_start_step and self.current_step + 1 < self.scheduler_end_step and self.current_step + 1 not in self.error_reset_steps: + self.is_next_RAS_step = True + else: + self.is_next_RAS_step = False + +MANAGER = ras_manager() diff --git a/modules/ras/ras_scheduler.py b/modules/ras/ras_scheduler.py new file mode 100644 index 000000000..a5143a067 --- /dev/null +++ b/modules/ras/ras_scheduler.py @@ -0,0 +1,215 @@ +# This file is a modified version of the original file from the HuggingFace/diffusers library. + +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import Optional, Tuple, Union +import torch +from diffusers.configuration_utils import register_to_config +from diffusers.utils import BaseOutput, logging +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from . import ras_manager + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class RASFlowMatchEulerDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class RASFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler): + """ + RAS Euler scheduler. + + This model inherits from ['FlowMatchEulerDiscreteScheduler']. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting=False, + base_shift: Optional[float] = 0.5, + max_shift: Optional[float] = 1.15, + base_image_seq_len: Optional[int] = 256, + max_image_seq_len: Optional[int] = 4096, + invert_sigmas: bool = False, + ): + super().__init__(num_train_timesteps=num_train_timesteps, + shift=shift, + use_dynamic_shifting=use_dynamic_shifting, + base_shift=base_shift, + max_shift=max_shift, + base_image_seq_len=base_image_seq_len, + max_image_seq_len=max_image_seq_len, + # invert_sigmas=invert_sigmas + ) + self.drop_cnt = None + + + def _init_ras_config(self, latents): + self.drop_cnt = torch.zeros((latents.shape[-2] // ras_manager.MANAGER.patch_size * latents.shape[-1] // ras_manager.MANAGER.patch_size), device=latents.device) - len(self.sigmas) + + def extract_latents_index_from_patched_latents_index(self, indices, height): + flattened_indices = indices // (height // ras_manager.MANAGER.patch_size) * ras_manager.MANAGER.patch_size * height + indices % (height // ras_manager.MANAGER.patch_size) *ras_manager.MANAGER.patch_size + flattened_indices = (flattened_indices[:, None] + torch.tensor([0, height + 1, 1, height], dtype=indices.dtype, device=indices.device)[None, :]).flatten() + return flattened_indices + + def ras_selection(self, sample, diff, height, width): + diff = diff.squeeze(0).permute(1, 2, 0) + # calculate the metric for each patch + if ras_manager.MANAGER.metric == "std": + metric = torch.std(diff, dim=-1).view(height // ras_manager.MANAGER.patch_size, ras_manager.MANAGER.patch_size, width // ras_manager.MANAGER.patch_size, ras_manager.MANAGER.patch_size).transpose(-2, -3).mean(-1).mean(-1).view(-1) + elif ras_manager.MANAGER.metric == "l2norm": + metric = torch.norm(diff, p=2, dim=-1).view(height // ras_manager.MANAGER.patch_size, ras_manager.MANAGER.patch_size, width // ras_manager.MANAGER.patch_size, ras_manager.MANAGER.patch_size).transpose(-2, -3).mean(-1).mean(-1).view(-1) + else: + raise ValueError("Unknown metric") + + # scale the metric with the drop count to avoid starvation + metric *= torch.exp(ras_manager.MANAGER.starvation_scale * self.drop_cnt) + current_skip_num = ras_manager.MANAGER.skip_token_num_list[self._step_index + 1] + assert ras_manager.MANAGER.high_ratio >= 0 and ras_manager.MANAGER.high_ratio <= 1, "High ratio should be in the range of [0, 1]" + indices = torch.sort(metric, dim=0, descending=False).indices + low_bar = int(current_skip_num * (1 - ras_manager.MANAGER.high_ratio)) + high_bar = int(current_skip_num * ras_manager.MANAGER.high_ratio) + cached_patchified_indices = torch.cat([indices[:low_bar], indices[-high_bar:]]) + other_patchified_indices = indices[low_bar:-high_bar] + self.drop_cnt[cached_patchified_indices] += 1 + latent_cached_indices = self.extract_latents_index_from_patched_latents_index(cached_patchified_indices, height) + + return latent_cached_indices, other_patchified_indices + + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[RASFlowMatchEulerDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + s_churn (`float`): + s_tmin (`float`): + s_tmax (`float`): + s_noise (`float`, defaults to 1.0): + Scaling factor for noise added to the sample. + generator (`torch.Generator`, *optional*): + A random number generator. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + if self.drop_cnt is None or self._step_index == 0: + self._init_ras_config(sample) + + if self._step_index == 0: + ras_manager.MANAGER.reset_cache() + + latent_dim, height, width = sample.shape[-3:] + + assert ras_manager.MANAGER.sample_ratio > 0.0 and ras_manager.MANAGER.sample_ratio <= 1.0 + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_RAS_step: + model_output.squeeze(0).view(latent_dim, -1)[:, ras_manager.MANAGER.cached_index] = ras_manager.MANAGER.cached_scaled_noise + model_output = model_output.transpose(0, 1).view(latent_dim, height, width).unsqueeze(0) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + + diff = (sigma_next - sigma) * model_output + + prev_sample = sample + diff + # Cast sample back to model compatible dtype + prev_sample = prev_sample.to(model_output.dtype) + + if ras_manager.MANAGER.sample_ratio < 1.0 and ras_manager.MANAGER.is_next_RAS_step: + latent_cached_indices, other_patchified_indices = self.ras_selection(sample, diff, height, width) + ras_manager.MANAGER.cached_scaled_noise = model_output.squeeze(0).view(latent_dim, -1)[:, latent_cached_indices] + ras_manager.MANAGER.cached_index = latent_cached_indices + ras_manager.MANAGER.other_patchified_index = other_patchified_indices + + # upon completion increase step index by one + self._step_index += 1 + ras_manager.MANAGER.increase_step() + if ras_manager.MANAGER.current_step >= ras_manager.MANAGER.num_steps: + ras_manager.MANAGER.reset_cache() + + if not return_dict: + return (prev_sample,) + + return RASFlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) diff --git a/modules/schedulers/scheduler_ufogen.py b/modules/schedulers/scheduler_ufogen.py index f4d8aee97..ff5f27eb3 100644 --- a/modules/schedulers/scheduler_ufogen.py +++ b/modules/schedulers/scheduler_ufogen.py @@ -520,4 +520,4 @@ class UFOGenScheduler(SchedulerMixin, ConfigMixin): ) prev_t = timestep - self.config.num_train_timesteps // num_inference_steps - return prev_t \ No newline at end of file + return prev_t diff --git a/modules/shared.py b/modules/shared.py index c327d1270..d1aa82ba2 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -662,6 +662,9 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "linfusion_sep": OptionInfo("

Batch

", "", gr.HTML), "enable_linfusion": OptionInfo(False, "Apply LinFusion distillation on load"), + "ras_sep": OptionInfo("

RAS: Region-Adaptive Sampling

", "", gr.HTML), + "ras_enable": OptionInfo(False, "RAS enabled"), + "inference_batch_sep": OptionInfo("

Batch

", "", gr.HTML), "sequential_seed": OptionInfo(True, "Batch mode uses sequential seeds"), "batch_frame_mode": OptionInfo(False, "Parallel process images in batch"),