|
|
| from __future__ import annotations |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from torch.nn.attention.flex_attention import create_block_mask |
| from transformers.utils import ModelOutput |
| from dataclasses import dataclass |
| from transformers.cache_utils import Cache |
|
|
| @dataclass |
| class SimpleOutputWithPast(ModelOutput): |
| loss: torch.FloatTensor | None = None |
| logits: torch.FloatTensor | None = None |
| causal_logits: torch.FloatTensor | None = None |
| past_key_values: Cache | None = None |
| hidden_states: tuple[torch.FloatTensor, ...] | None = None |
| attentions: tuple[torch.FloatTensor, ...] | None = None |
| from .utils import maybe_truncate_last_dim, pad_along_last_dim |
| from .configuration_nlduni import get_ministral_shift_sbd_logits |
| from .sbd_attention import build_image_attention_doc_ids |
|
|
| def _resolve_shift_sbd_logits(model, shift_logits): |
| model_module = model.module if hasattr(model, 'module') else model |
| config = getattr(model_module, 'config', None) |
| configured = get_ministral_shift_sbd_logits(config) if config is not None else None |
| has_alignment_contract = config is not None |
| if shift_logits is None: |
| return bool(configured) if configured is not None else False |
| requested = bool(shift_logits) |
| if has_alignment_contract and requested != configured: |
| raise ValueError('Inference shift_logits must match config.shift_sbd_logits to avoid a one-token alignment error. Update the model config explicitly if this override is intentional.') |
| return requested |
|
|
| def wte(model, x, t2i_inference=False, gen_shape=None, x_gen=None, inputs_embeds_curr=None, new_token_mask=None, x_gen_child_embeddings=None, gen_parent_mask=None, gen_source_null_mask=None, gen_role=None): |
| if t2i_inference: |
| if (x_gen is None) == (x_gen_child_embeddings is None): |
| raise ValueError('provide exactly one of x_gen or x_gen_child_embeddings') |
| if new_token_mask is None: |
| new_token_mask = x >= INT_MAX |
| if gen_parent_mask is not None or gen_source_null_mask is not None or gen_role is not None: |
| raise ValueError('structured generation masks and roles require Infinity BSQ') |
| gen_latents_comp_embeds = model.call_gen_embedding(x_gen, gen_shape, child_embeddings=x_gen_child_embeddings) |
| if inputs_embeds_curr is None: |
| x_txt_only = x.clone() |
| x_txt_only[new_token_mask] = 0 |
| inputs_embeds_curr = model.embed_tokens(x_txt_only) |
| inputs_embeds_curr[new_token_mask] = pad_along_last_dim(gen_latents_comp_embeds, inputs_embeds_curr.shape[-1]).view(-1, inputs_embeds_curr.shape[-1]) |
| else: |
| inputs_embeds_curr = model.embed_tokens(x) |
| new_token_mask = None |
| return (inputs_embeds_curr, new_token_mask) |
| INT_MAX = 1000000 |
| SBD_FLEX_ATTENTION_PAD_MULTIPLE = 512 |
|
|
| def _round_up_to_multiple(length, multiple): |
| if isinstance(multiple, bool) or not isinstance(multiple, int) or multiple < 1: |
| raise ValueError('padding multiple must be a positive integer') |
| return (length + multiple - 1) // multiple * multiple |
|
|
| def single_stream_sbd_sampling_mask(block_size, past_length, b, h, q_idx, kv_idx, query_image_doc_id, key_value_image_doc_id): |
| del block_size, h |
| absolute_q = q_idx + past_length |
| token_causal = absolute_q >= kv_idx |
| image_query = query_image_doc_id[b, q_idx] > 0 |
| same_image_doc = image_query & (query_image_doc_id[b, q_idx] == key_value_image_doc_id[b, kv_idx]) |
| return token_causal | same_image_doc |
|
|
| def _create_single_stream_sbd_sampling_block_mask(input_embeddings, *, past_key_values, block_size, input_modality_indices, generation_token_mask=None, pad_to_multiple=None): |
| if isinstance(block_size, bool) or not isinstance(block_size, int) or block_size < 1: |
| raise ValueError('SBD diffusion sampling requires a positive integer block_size') |
| query_length = input_embeddings.shape[1] |
| past_length = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| key_value_length = past_length + query_length |
| padded_query_length = _round_up_to_multiple(query_length, pad_to_multiple) if pad_to_multiple is not None else query_length |
| padded_key_value_length = _round_up_to_multiple(key_value_length, pad_to_multiple) if pad_to_multiple is not None else key_value_length |
| if input_modality_indices is None: |
| raise ValueError('SBD diffusion sampling requires input_modality_indices to route same-block generation attention') |
| if tuple(input_modality_indices.shape) != tuple(input_embeddings.shape[:2]): |
| raise ValueError('SBD diffusion sampling requires input_modality_indices to align with the current input embeddings') |
| query_transformer_routing = input_modality_indices.to(dtype=torch.bool) |
| if generation_token_mask is None: |
| generation_token_mask = query_transformer_routing |
| elif tuple(generation_token_mask.shape) != tuple(input_embeddings.shape[:2]): |
| raise ValueError('SBD diffusion sampling requires generation_token_mask to align with the current input embeddings') |
| query_target_generation_routing = generation_token_mask.to(dtype=torch.bool) |
| cached_generation_routing = torch.zeros(query_transformer_routing.shape[0], past_length, dtype=torch.bool, device=query_transformer_routing.device) |
| padded_generation_routing = torch.zeros(query_transformer_routing.shape[0], padded_key_value_length - key_value_length, dtype=torch.bool, device=query_transformer_routing.device) |
| key_value_target_generation_routing = torch.cat((cached_generation_routing, query_target_generation_routing, padded_generation_routing), dim=1) |
| key_value_valid = (torch.arange(padded_key_value_length, device=query_transformer_routing.device).unsqueeze(0) < key_value_length).expand(query_transformer_routing.shape[0], -1).contiguous() |
| key_value_image_doc_id = build_image_attention_doc_ids(key_value_target_generation_routing, valid_tokens=key_value_valid) |
| query_image_doc_id = torch.cat((key_value_image_doc_id[:, past_length:past_length + query_length], torch.zeros(query_transformer_routing.shape[0], padded_query_length - query_length, dtype=key_value_image_doc_id.dtype, device=query_transformer_routing.device)), dim=1) |
| past_length_tensor = torch.tensor(past_length, dtype=torch.long, device=query_transformer_routing.device) |
| query_valid = (torch.arange(padded_query_length, device=query_transformer_routing.device).unsqueeze(0) < query_length).expand(query_transformer_routing.shape[0], -1).contiguous() |
|
|
| def mask_mod(b, h, q_idx, kv_idx): |
| logical_attention = single_stream_sbd_sampling_mask(block_size, past_length_tensor, b, h, q_idx, kv_idx, query_image_doc_id, key_value_image_doc_id) & query_valid[b, q_idx] & key_value_valid[b, kv_idx] |
| padded_query_attention = ~query_valid[b, q_idx] & (kv_idx == 0) |
| return logical_attention | padded_query_attention |
| return create_block_mask(mask_mod, B=input_embeddings.shape[0], H=None, Q_LEN=padded_query_length, KV_LEN=padded_key_value_length, device=input_embeddings.device) |
|
|
| def get_logits(model, input_emnbeddings, modality_indices=None, t2i_inference=False, past_key_values=None, gen_shape=None, timesteps=None, input_modality_indices=None, generation_attention_mask=None, sbd_diffusion_sampling=False, full_attention_approx=False, loop_k=None): |
| if t2i_inference: |
| resolved_loop_k = None |
| if input_modality_indices is None: |
| input_modality_indices = modality_indices |
| attention_kwargs = {} |
| if full_attention_approx: |
| attention_kwargs['overwrite_block_mask'] = 'full' |
| attention_kwargs['overwrite_attn_impl'] = 'flash_attn' |
| elif sbd_diffusion_sampling: |
| if generation_attention_mask is None: |
| generation_attention_mask = modality_indices |
| attention_kwargs['overwrite_block_mask'] = _create_single_stream_sbd_sampling_block_mask(input_emnbeddings, past_key_values=past_key_values, block_size=getattr(model.config, 'block_size', None), input_modality_indices=input_modality_indices, generation_token_mask=generation_attention_mask, pad_to_multiple=SBD_FLEX_ATTENTION_PAD_MULTIPLE) |
| else: |
| attention_kwargs['overwrite_attn_impl'] = 'flash_attn' |
| output = model(None, input_embeddings=input_emnbeddings, modality_indices=input_modality_indices, output_hidden_states=True, past_key_values=past_key_values, use_cache=None, single_stream_image_diffusion=True, loop_k=None, **attention_kwargs) |
| hidden_states = getattr(output, 'last_hidden_state', None) |
| if hidden_states is None: |
| hidden_states = output.hidden_states[-1] |
| gen_hidden_states = hidden_states[modality_indices] |
| gen_hidden_states = maybe_truncate_last_dim(gen_hidden_states, model.config.d_model_gen) |
| gen_logits = model.call_gen_predictor(gen_hidden_states, gen_shape, timesteps=timesteps) |
| seq_len_per_img = np.prod(gen_shape) |
| if len(gen_logits.shape) == 2: |
| gen_logits = gen_logits.view(-1, seq_len_per_img, gen_logits.shape[-1]) |
| else: |
| gen_logits = gen_logits.view(-1, seq_len_per_img, *gen_logits.shape[-2:]) |
| return gen_logits |
| else: |
| raise ValueError('get_logits is the image sampler helper') |
|
|
| def add_gumbel_noise(logits, temperature): |
| if temperature == 0: |
| return logits |
| logits = logits.to(torch.float64) |
| noise = torch.rand_like(logits, dtype=torch.float64) |
| gumbel_noise = (-torch.log(noise)) ** temperature |
| return logits.exp() / gumbel_noise |
|
|
| def get_transfer_index(logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None, neg_entropy=False): |
| logits_with_noise = add_gumbel_noise(logits, temperature=temperature) |
| x0 = torch.argmax(logits_with_noise, dim=-1) |
| if remasking == 'low_confidence': |
| p = F.softmax(logits, dim=-1) |
| x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) |
| elif remasking == 'top_p_margin': |
| p = F.softmax(logits, dim=-1) |
| top2 = torch.topk(p, k=2, dim=-1).values |
| margin = top2[..., 0] - top2[..., 1] |
| plus_inf = torch.full_like(margin, 1e309) |
| minus_inf = torch.full_like(margin, -1e309) |
| masked_for_min = torch.where(mask_index, margin, plus_inf) |
| masked_for_max = torch.where(mask_index, margin, minus_inf) |
| row_min = masked_for_min.amin(dim=1, keepdim=True) |
| row_max = masked_for_max.amax(dim=1, keepdim=True) |
| denom = row_max - row_min |
| normalized = torch.zeros_like(margin) |
| nonzero = denom > 0 |
| normalized = torch.where(mask_index & nonzero, (margin - row_min) / (denom + 1e-12), normalized) |
| normalized = torch.where(mask_index & ~nonzero, torch.ones_like(normalized), normalized) |
| x0_p = normalized |
| elif remasking == 'random': |
| x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device) |
| else: |
| raise NotImplementedError(remasking) |
| if neg_entropy: |
| p = F.softmax(logits, dim=-1) |
| epsilon = 1e-10 |
| log_probs = torch.log(p + epsilon) |
| confidence_scores = torch.sum(p * log_probs, dim=-1) |
| else: |
| confidence_scores = x0_p |
| x0 = torch.where(mask_index, x0, x) |
| confidence = torch.where(mask_index, confidence_scores, -np.inf) |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device) |
| if threshold is not None: |
| num_transfer_tokens = mask_index.sum(dim=1, keepdim=True) |
| for j in range(confidence.shape[0]): |
| (_, select_index) = torch.topk(confidence[j], k=num_transfer_tokens[j]) |
| transfer_index[j, select_index] = True |
| if threshold is not None: |
| for k in range(1, num_transfer_tokens[j]): |
| if confidence[j, select_index[k]] < threshold: |
| transfer_index[j, select_index[k]] = False |
| return (x0, transfer_index) |
|
|
| def get_num_transfer_tokens(mask_index, steps: int): |
| mask_num = mask_index.sum(dim=1, keepdim=True) |
| base = mask_num // steps |
| remainder = mask_num % steps |
| num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base |
| for i in range(mask_num.size(0)): |
| num_transfer_tokens[i, :int(remainder[i])] += 1 |
| return num_transfer_tokens |
|
|
| def simple_fwd(model, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, past_key_values=None, **kwargs): |
| enc_out = model.encoder(past_key_values=past_key_values, input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, overwrite_attn_impl='flash_attn', **kwargs) |
| logits = model.diffusion_head(enc_out.last_hidden_state) |
| return SimpleOutputWithPast(loss=logits, logits=logits, causal_logits=None, past_key_values=enc_out.past_key_values, hidden_states=None, attentions=None) |
|
|
| @torch.no_grad() |
| def generate_with_prefix_cache_block_diff(model, prompt=None, prompt_embeds=None, steps=128, gen_length=128, block_length=128, temperature=0.0, remasking='low_confidence', mask_id=126336, threshold=None, factor=None, shift_logits=None, neg_entropy=False, causal_context=False, eos_token_id=None, max_thinking_tokens=None, end_think_token_id=None): |
| shift_logits = _resolve_shift_sbd_logits(model, shift_logits) |
| dream_style = shift_logits |
| if (prompt is None) == (prompt_embeds is None): |
| raise ValueError('Exactly one of `prompt` or `prompt_embeds` must be provided.') |
| if prompt is not None: |
| prompt_ids = prompt |
| prompt_len = prompt_ids.shape[1] |
| x_accum = prompt_ids.clone() |
| B = prompt_ids.shape[0] |
| token_device = prompt_ids.device |
| token_dtype = prompt_ids.dtype |
| else: |
| prompt_ids = None |
| prompt_len = prompt_embeds.shape[1] |
| B = prompt_embeds.shape[0] |
| token_device = prompt_embeds.device |
| token_dtype = torch.long |
| x_accum = torch.full((B, prompt_len), mask_id, dtype=token_dtype, device=token_device) |
| if gen_length % block_length != 0: |
| gen_length = (gen_length // block_length + 1) * block_length |
| num_blocks = gen_length // block_length |
| assert steps % num_blocks == 0 |
| steps_per_block = steps // num_blocks |
| nfe = 0 |
| model_module = model.module if hasattr(model, 'module') else model |
| for layer in model_module.encoder.layers: |
| layer.self_attn.mode = 'bidirectional' |
| if causal_context: |
| for layer in model_module.encoder.layers: |
| if hasattr(layer.self_attn, 'diffusion_lm'): |
| layer.self_attn.diffusion_lm = False |
| output = simple_fwd(model, input_ids=prompt_ids, inputs_embeds=prompt_embeds, use_cache=True, use_causal_mask=causal_context) |
| past_key_values = output.past_key_values |
| if causal_context: |
| for layer in model_module.encoder.layers: |
| if hasattr(layer.self_attn, 'diffusion_lm'): |
| layer.self_attn.diffusion_lm = True |
| next_token = None |
| if causal_context: |
| last_logit = output.logits[:, -1, :] |
| if temperature > 0: |
| probs = torch.softmax(last_logit / temperature, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| else: |
| next_token = torch.argmax(last_logit, dim=-1, keepdim=True) |
| next_logits_context = None |
| if dream_style: |
| next_logits_context = output.logits[:, -1:, :] |
| for num_block in range(num_blocks): |
| mask_block = torch.ones((B, block_length), dtype=token_dtype, device=token_device) * mask_id |
| if causal_context: |
| mask_block[:, 0] = next_token[:, 0] |
| x_accum = torch.cat([x_accum, mask_block], dim=1) |
| current_block_start = prompt_len + num_block * block_length |
| block_slice = slice(current_block_start, current_block_start + block_length) |
| if end_think_token_id is not None and max_thinking_tokens is not None: |
| tokens_before_block = num_block * block_length |
| tokens_after_block = tokens_before_block + block_length |
| if tokens_after_block > max_thinking_tokens: |
| gen_so_far = x_accum[:, prompt_len:current_block_start] |
| has_end_think = (gen_so_far == end_think_token_id).any(dim=1) if gen_so_far.size(1) > 0 else torch.zeros(B, dtype=torch.bool, device=token_device) |
| if not has_end_think.all(): |
| if tokens_before_block < max_thinking_tokens: |
| offset = max_thinking_tokens - tokens_before_block |
| else: |
| offset = 0 |
| inject_pos = current_block_start + offset |
| for b in range(B): |
| if not has_end_think[b]: |
| x_accum[b, inject_pos] = end_think_token_id |
| mask_block_idx0 = x_accum[:, block_slice] == mask_id |
| if dream_style: |
| schedule_mask = mask_block_idx0 |
| else: |
| schedule_mask = mask_block_idx0 |
| num_transfer_tokens = get_num_transfer_tokens(schedule_mask, steps_per_block) |
| for i in range(steps_per_block): |
| mask_block_idx = x_accum[:, block_slice] == mask_id |
| if mask_block_idx.sum() == 0: |
| break |
| nfe += 1 |
| logits_block = simple_fwd(model, x_accum[:, block_slice], past_key_values=past_key_values, use_cache=False).logits |
| if dream_style: |
| if block_length == 1: |
| logits_use = next_logits_context |
| else: |
| logits_use = torch.cat([next_logits_context, logits_block[:, :-1, :]], dim=1) |
| mask_use = mask_block_idx |
| x_use = x_accum[:, block_slice] |
| (x0, transfer_idx) = get_transfer_index(logits_use, temperature, remasking, mask_use, x_use, num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, neg_entropy=neg_entropy) |
| cur = x_accum[:, block_slice].clone() |
| cur[transfer_idx] = x0[transfer_idx] |
| x_accum[:, block_slice] = cur |
| else: |
| (x0, transfer_idx) = get_transfer_index(logits_block, temperature, remasking, mask_block_idx, x_accum[:, block_slice], num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, neg_entropy=neg_entropy) |
| cur = x_accum[:, block_slice].clone() |
| cur[transfer_idx] = x0[transfer_idx] |
| x_accum[:, block_slice] = cur |
| if eos_token_id is not None: |
| block_tokens = x_accum[:, block_slice] |
| eos_mask = block_tokens == eos_token_id |
| any_eos = eos_mask.any(dim=1) |
| if any_eos.any(): |
| after_eos = eos_mask.cumsum(dim=1).bool() |
| mask_before = (block_tokens == mask_id) & ~after_eos |
| if (any_eos & ~mask_before.any(dim=1)).any(): |
| break |
| if causal_context: |
| for layer in model_module.encoder.layers: |
| if hasattr(layer.self_attn, 'diffusion_lm'): |
| layer.self_attn.diffusion_lm = False |
| output = simple_fwd(model, x_accum[:, block_slice], past_key_values=past_key_values, use_cache=True, use_causal_mask=causal_context) |
| past_key_values = output.past_key_values |
| nfe += 1 |
| if causal_context: |
| for layer in model_module.encoder.layers: |
| if hasattr(layer.self_attn, 'diffusion_lm'): |
| layer.self_attn.diffusion_lm = True |
| last_logit = output.logits[:, -1, :] |
| if temperature > 0: |
| probs = torch.softmax(last_logit / temperature, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| else: |
| next_token = torch.argmax(last_logit, dim=-1, keepdim=True) |
| if dream_style and num_block < num_blocks - 1: |
| next_logits_context = output.logits[:, -1:, :] |
| if eos_token_id is not None: |
| gen_so_far = x_accum[:, prompt_len:] |
| is_eos = gen_so_far == eos_token_id |
| has_eos = is_eos.any(dim=1) |
| if has_eos.all(): |
| first_eos_pos = is_eos.to(torch.int64).argmax(dim=1) |
| max_eos = first_eos_pos.max().item() |
| if prompt_ids is None: |
| return (x_accum[:, prompt_len:prompt_len + max_eos + 1], nfe) |
| return (x_accum[:, :prompt_len + max_eos + 1], nfe) |
| if prompt_ids is None: |
| return (x_accum[:, prompt_len:], nfe) |
| return (x_accum, nfe) |
|
|
