|
|
from diffusers import FluxPipeline
|
|
|
from typing import List, Dict, Callable, Union
|
|
|
import torch
|
|
|
|
|
|
def retrieve(io):
|
|
|
if isinstance(io, tuple):
|
|
|
if len(io) == 1:
|
|
|
return io[0]
|
|
|
elif len(io) ==2:
|
|
|
return io
|
|
|
elif len(io) ==3:
|
|
|
return io[0]
|
|
|
else:
|
|
|
raise ValueError("A tuple should have length of 1")
|
|
|
elif isinstance(io, torch.Tensor):
|
|
|
return io
|
|
|
else:
|
|
|
raise ValueError("Input/Output must be a tensor, or 1-element tuple")
|
|
|
|
|
|
|
|
|
class HookedDiffusionAbstractPipeline:
|
|
|
parent_cls = None
|
|
|
pipe = None
|
|
|
|
|
|
def __init__(self, pipe: parent_cls, use_hooked_scheduler: bool = False):
|
|
|
self.__dict__['pipe'] = pipe
|
|
|
self.use_hooked_scheduler = use_hooked_scheduler
|
|
|
|
|
|
@classmethod
|
|
|
def from_pretrained(cls, *args, **kwargs):
|
|
|
return cls(cls.parent_cls.from_pretrained(*args, **kwargs))
|
|
|
|
|
|
|
|
|
def run_with_hooks(self,
|
|
|
*args,
|
|
|
position_hook_dict: Dict[str, Union[Callable, List[Callable]]],
|
|
|
**kwargs
|
|
|
):
|
|
|
'''
|
|
|
Run the pipeline with hooks at specified positions.
|
|
|
Returns the final output.
|
|
|
|
|
|
Args:
|
|
|
*args: Arguments to pass to the pipeline.
|
|
|
position_hook_dict: A dictionary mapping positions to hooks.
|
|
|
The keys are positions in the pipeline where the hooks should be registered.
|
|
|
The values are either a single hook or a list of hooks to be registered at the specified position.
|
|
|
Each hook should be a callable that takes three arguments: (module, input, output).
|
|
|
**kwargs: Keyword arguments to pass to the pipeline.
|
|
|
'''
|
|
|
hooks = []
|
|
|
for position, hook in position_hook_dict.items():
|
|
|
if isinstance(hook, list):
|
|
|
for h in hook:
|
|
|
hooks.append(self._register_general_hook(position, h))
|
|
|
else:
|
|
|
hooks.append(self._register_general_hook(position, hook))
|
|
|
|
|
|
hooks = [hook for hook in hooks if hook is not None]
|
|
|
|
|
|
try:
|
|
|
output = self.pipe(*args, **kwargs)
|
|
|
finally:
|
|
|
for hook in hooks:
|
|
|
hook.remove()
|
|
|
if self.use_hooked_scheduler:
|
|
|
self.pipe.scheduler.pre_hooks = []
|
|
|
self.pipe.scheduler.post_hooks = []
|
|
|
|
|
|
return output
|
|
|
|
|
|
def run_with_cache(self,
|
|
|
*args,
|
|
|
positions_to_cache: List[str],
|
|
|
save_input: bool = False,
|
|
|
save_output: bool = True,
|
|
|
**kwargs
|
|
|
):
|
|
|
cache_input, cache_output = dict() if save_input else None, dict() if save_output else None
|
|
|
hooks = [
|
|
|
self._register_cache_hook(position, cache_input, cache_output) for position in positions_to_cache
|
|
|
]
|
|
|
hooks = [hook for hook in hooks if hook is not None]
|
|
|
output = self.pipe(*args, **kwargs)
|
|
|
for hook in hooks:
|
|
|
hook.remove()
|
|
|
if self.use_hooked_scheduler:
|
|
|
self.pipe.scheduler.pre_hooks = []
|
|
|
self.pipe.scheduler.post_hooks = []
|
|
|
|
|
|
cache_dict = {}
|
|
|
if save_input:
|
|
|
for position, block in cache_input.items():
|
|
|
cache_input[position] = torch.stack(block, dim=1)
|
|
|
cache_dict['input'] = cache_input
|
|
|
|
|
|
if save_output:
|
|
|
for position, block in cache_output.items():
|
|
|
|
|
|
cache_output[position] = block
|
|
|
cache_dict['output'] = cache_output
|
|
|
return output, cache_dict
|
|
|
|
|
|
def run_with_hooks_and_cache(self,
|
|
|
*args,
|
|
|
position_hook_dict: Dict[str, Union[Callable, List[Callable]]],
|
|
|
positions_to_cache: List[str] = [],
|
|
|
save_input: bool = False,
|
|
|
save_output: bool = True,
|
|
|
**kwargs
|
|
|
):
|
|
|
cache_input, cache_output = dict() if save_input else None, dict() if save_output else None
|
|
|
hooks = [
|
|
|
self._register_cache_hook(position, cache_input, cache_output) for position in positions_to_cache
|
|
|
]
|
|
|
|
|
|
for position, hook in position_hook_dict.items():
|
|
|
if isinstance(hook, list):
|
|
|
for h in hook:
|
|
|
hooks.append(self._register_general_hook(position, h))
|
|
|
else:
|
|
|
hooks.append(self._register_general_hook(position, hook))
|
|
|
|
|
|
hooks = [hook for hook in hooks if hook is not None]
|
|
|
output = self.pipe(*args, **kwargs)
|
|
|
for hook in hooks:
|
|
|
hook.remove()
|
|
|
if self.use_hooked_scheduler:
|
|
|
self.pipe.scheduler.pre_hooks = []
|
|
|
self.pipe.scheduler.post_hooks = []
|
|
|
|
|
|
cache_dict = {}
|
|
|
if save_input:
|
|
|
for position, block in cache_input.items():
|
|
|
cache_input[position] = torch.stack(block, dim=1)
|
|
|
cache_dict['input'] = cache_input
|
|
|
|
|
|
if save_output:
|
|
|
for position, block in cache_output.items():
|
|
|
cache_output[position] = torch.stack(block, dim=1)
|
|
|
cache_dict['output'] = cache_output
|
|
|
|
|
|
return output, cache_dict
|
|
|
|
|
|
|
|
|
def _locate_block(self, position: str):
|
|
|
block = self.pipe
|
|
|
for step in position.split('.'):
|
|
|
if step.isdigit():
|
|
|
step = int(step)
|
|
|
block = block[step]
|
|
|
else:
|
|
|
block = getattr(block, step)
|
|
|
return block
|
|
|
|
|
|
|
|
|
def _register_cache_hook(self, position: str, cache_input: Dict, cache_output: Dict):
|
|
|
|
|
|
if position.endswith('$self_attention') or position.endswith('$cross_attention'):
|
|
|
return self._register_cache_attention_hook(position, cache_output)
|
|
|
|
|
|
if position == 'noise':
|
|
|
def hook(model_output, timestep, sample, generator):
|
|
|
if position not in cache_output:
|
|
|
cache_output[position] = []
|
|
|
cache_output[position].append(sample)
|
|
|
|
|
|
if self.use_hooked_scheduler:
|
|
|
self.pipe.scheduler.post_hooks.append(hook)
|
|
|
else:
|
|
|
raise ValueError('Cannot cache noise without using hooked scheduler')
|
|
|
return
|
|
|
|
|
|
block = self._locate_block(position)
|
|
|
|
|
|
def hook(module, input, kwargs, output):
|
|
|
if cache_input is not None:
|
|
|
if position not in cache_input:
|
|
|
cache_input[position] = []
|
|
|
cache_input[position].append(retrieve(input))
|
|
|
|
|
|
if cache_output is not None:
|
|
|
if position not in cache_output:
|
|
|
cache_output[position] = []
|
|
|
cache_output[position].append(retrieve(output))
|
|
|
|
|
|
return block.register_forward_hook(hook, with_kwargs=True)
|
|
|
|
|
|
def _register_cache_attention_hook(self, position, cache):
|
|
|
attn_block = self._locate_block(position.split('$')[0])
|
|
|
if position.endswith('$self_attention'):
|
|
|
attn_block = attn_block.attn1
|
|
|
elif position.endswith('$cross_attention'):
|
|
|
attn_block = attn_block.attn2
|
|
|
else:
|
|
|
raise ValueError('Wrong attention type')
|
|
|
|
|
|
def hook(module, args, kwargs, output):
|
|
|
hidden_states = args[0]
|
|
|
encoder_hidden_states = kwargs['encoder_hidden_states']
|
|
|
attention_mask = kwargs['attention_mask']
|
|
|
batch_size, sequence_length, _ = hidden_states.shape
|
|
|
attention_mask = attn_block.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
|
query = attn_block.to_q(hidden_states)
|
|
|
|
|
|
|
|
|
if encoder_hidden_states is None:
|
|
|
encoder_hidden_states = hidden_states
|
|
|
elif attn_block.norm_cross is not None:
|
|
|
encoder_hidden_states = attn_block.norm_cross(encoder_hidden_states)
|
|
|
|
|
|
key = attn_block.to_k(encoder_hidden_states)
|
|
|
value = attn_block.to_v(encoder_hidden_states)
|
|
|
|
|
|
query = attn_block.head_to_batch_dim(query)
|
|
|
key = attn_block.head_to_batch_dim(key)
|
|
|
value = attn_block.head_to_batch_dim(value)
|
|
|
|
|
|
attention_probs = attn_block.get_attention_scores(query, key, attention_mask)
|
|
|
attention_probs = attention_probs.view(
|
|
|
batch_size,
|
|
|
attention_probs.shape[0] // batch_size,
|
|
|
attention_probs.shape[1],
|
|
|
attention_probs.shape[2]
|
|
|
)
|
|
|
if position not in cache:
|
|
|
cache[position] = []
|
|
|
cache[position].append(attention_probs)
|
|
|
|
|
|
return attn_block.register_forward_hook(hook, with_kwargs=True)
|
|
|
|
|
|
def _register_general_hook(self, position, hook):
|
|
|
if position == 'scheduler_pre':
|
|
|
if not self.use_hooked_scheduler:
|
|
|
raise ValueError('Cannot register hooks on scheduler without using hooked scheduler')
|
|
|
self.pipe.scheduler.pre_hooks.append(hook)
|
|
|
return
|
|
|
elif position == 'scheduler_post':
|
|
|
if not self.use_hooked_scheduler:
|
|
|
raise ValueError('Cannot register hooks on scheduler without using hooked scheduler')
|
|
|
self.pipe.scheduler.post_hooks.append(hook)
|
|
|
return
|
|
|
|
|
|
block = self._locate_block(position)
|
|
|
return block.register_forward_hook(hook)
|
|
|
|
|
|
def to(self, *args, **kwargs):
|
|
|
self.pipe = self.pipe.to(*args, **kwargs)
|
|
|
return self
|
|
|
|
|
|
def __getattr__(self, name):
|
|
|
return getattr(self.pipe, name)
|
|
|
|
|
|
def __setattr__(self, name, value):
|
|
|
return setattr(self.pipe, name, value)
|
|
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
|
return self.pipe(*args, **kwargs)
|
|
|
|
|
|
|
|
|
class HookedFluxPipeline(HookedDiffusionAbstractPipeline):
|
|
|
parent_cls = FluxPipeline |
