File size: 4,938 Bytes
ffc54ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """Avoid rebuilding immutable prefix layout during cached image decode."""
import math
import types
import torch
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache
def metadata(model, kw):
cache = kw['kv_cache']
if not hasattr(cache, '_fast_layout'):
mask = kw['img_mask']
repeats = torch.where(mask, 4, 1)[0]
image_pad_mask = torch.repeat_interleave(mask[0], repeats)
shapes = kw['img_shapes']
rope = model.pos_embed(shapes[0], image_pad_mask, device=mask.device)
_, target_mask = model.build_token_metadata(image_pad_mask, shapes[0])
target_tokens = math.prod(shapes[0][-1])
prefix = rope.shape[0] - target_tokens
valid = None
enc_mask = kw.get('encoder_hidden_states_mask')
if enc_mask is not None:
valid = torch.ones(kw['hidden_states'].shape[0], image_pad_mask.shape[0], dtype=torch.bool, device=mask.device)
valid[:, (~image_pad_mask).nonzero(as_tuple=True)[0]] = enc_mask.bool()[:, ~mask[0][:enc_mask.shape[1]]]
valid = valid[:, None, None, :]
cache._fast_layout = (rope[prefix:], target_mask[prefix:], valid, target_tokens)
return cache._fast_layout
def core(model, hidden_states, timestep, cache, layout):
rope, target_mask, attention_mask, target_tokens = layout
# Keep the original img_in shape, including condition-image tokens for edits.
image = model.img_in(hidden_states)
hidden = image[:, -target_tokens:]
timestep = timestep.to(image.dtype)
timestep = torch.cat([timestep, timestep.new_zeros(1)], dim=0)
temb = model.time_text_embed(timestep, image)
modulation = model.modulation(temb)
for i, block in enumerate(model.transformer_blocks):
hidden = block(hidden_states=hidden, modulation=modulation, rotary_emb=rope,
attention_mask=attention_mask, target_token_mask=target_mask,
layer_cache=cache.get_layer(i), kv_cache_mode='cached',
cache_write_slice=None, segments=None, key_valid=None)
hidden = model.norm_out(hidden, temb, target_mask)
return model.proj_out(hidden)
class Graph:
def __init__(self, model, hidden, timestep, cache, layout):
self.hidden = hidden.clone()
self.timestep = timestep.clone()
self.cache = QwenImage21KVCache(len(cache.layer_caches))
for dst, src in zip(self.cache.layer_caches, cache.layer_caches):
dst.store(src.k.clone(), src.v.clone())
self.layout = tuple(x.clone() if isinstance(x, torch.Tensor) else x for x in layout)
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
for _ in range(2):
core(model, self.hidden, self.timestep, self.cache, self.layout)
torch.cuda.current_stream().wait_stream(stream)
self.graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(self.graph, stream=stream):
self.output = core(model, self.hidden, self.timestep, self.cache, self.layout)
self.owner = cache
def run(self, hidden, timestep, cache, layout):
self.hidden.copy_(hidden)
self.timestep.copy_(timestep)
if self.owner is not cache:
for dst, src in zip(self.cache.layer_caches, cache.layer_caches):
dst.k.copy_(src.k)
dst.v.copy_(src.v)
for dst, src in zip(self.layout, layout):
if isinstance(dst, torch.Tensor):
dst.copy_(src)
self.owner = cache
self.graph.replay()
return self.output
def dispatch(self, **kw):
if kw.get('kv_cache_mode') != 'cached' or kw.get('attention_kwargs'):
return self._original_forward(**kw)
cache = kw['kv_cache']
layout = metadata(self, kw)
hidden, timestep = kw['hidden_states'], kw['timestep']
if self._fast_use_graph:
key = (tuple(hidden.shape), tuple(cache.get_layer(0).k.shape),
tuple(kw['img_shapes'][0]), layout[2] is not None)
if key not in self._fast_graphs:
if len(self._fast_graphs) >= 2:
self._fast_graphs.pop(next(iter(self._fast_graphs)))
self._fast_graphs[key] = Graph(self, hidden, timestep, cache, layout)
result = self._fast_graphs[key].run(hidden, timestep, cache, layout)
else:
result = core(self, hidden, timestep, cache, layout)
return (result,) if not kw.get('return_dict', True) else Transformer2DModelOutput(sample=result)
def enable(pipe, graph=False):
m = pipe.transformer
if not hasattr(m, '_original_forward'):
m._original_forward = m.forward
m.forward = types.MethodType(dispatch, m)
m._fast_graphs = {}
m._fast_use_graph = graph
return pipe
|