Release calibrated Image 2.1 Turbo FP8 transformer, accelerated SM120 runtime, quality evidence and real-time demo
ffc54ec verified | """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 | |