| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| from __future__ import annotations |
| import os |
| from .image_utils import compute_num_image_tokens, fix_one_prompt, process_images, tokenizer_image_token |
| from .token_constants import IMAGE_TOKEN_INDEX, IMAGE_PAD_TOKEN_INDEX |
| from .conversation import conv_templates |
| from PIL import Image |
| import copy |
| import functools |
| import math |
| import torch |
| from .schedules import cosine_schedule_2, logit_normal_schedule |
| import time |
| from .chat_utils import wte as wte_ministral |
| from .chat_utils import get_logits as get_logits_ministral |
| from .chat_utils import _resolve_shift_sbd_logits as resolve_ministral_shift_sbd_logits |
| from .sbd_attention import generated_image_attention_window |
| from tqdm.auto import tqdm |
| import numpy as np |
| import torch.distributions as dists |
| import torch.nn.functional as F |
| from .image_utils import pad_to_square_and_resize |
| from .utils import pad_along_last_dim |
| from .padding import _assert_source_padding_vq_locked, _prefill_source_padding_vq, _restore_source_padding_vq, _resolve_source_native_size, _source_padding_vq_mask, resolve_source_padding_vq_lock |
| from .refinement import refinement_mask_fraction |
| from .markers import append_image_row_markers, augmented_image_shape, image_fixed_prefix_token_count, image_generation_special_ids, image_row_marker_count, image_row_markers_enabled, resolve_generation_skip_down_sample, split_generation_fixed_prefix, strip_image_row_markers, validate_image_fixed_prefix_support, validate_image_row_marker_support |
| from .tree import BalancedSemanticTree, advance_revisable_semantic_tree, build_semantic_tree_embeddings, production_semantic_tree_contract, resolve_semantic_tree_atomic_block_size, resolve_semantic_tree_refinement_gap, resolve_semantic_tree_steps, semantic_tree_activation_counts, semantic_tree_decode_schedule, semantic_tree_noise_for_ages, semantic_tree_refinement_due |
| from .pmj import stratified_random |
|
|
| def _load_prediction_image(image): |
| if isinstance(image, (str, os.PathLike)): |
| image = Image.open(image) |
| return image.convert('RGB') |
|
|
| def _validate_text_to_image_batch_size(batch_size): |
| if batch_size != 1: |
| raise ValueError(f'text_to_image currently supports batch_size=1 only; got {batch_size}') |
|
|
| def _prediction_constant_version(model_config): |
| model_type = str(getattr(model_config, 'model_type', '')) |
| return 'ministral' |
| return None |
|
|
| def _move_image_to_device(value, *, device, dtype): |
| if torch.is_tensor(value): |
| return value.to(device=device, dtype=dtype) |
| if isinstance(value, tuple): |
| if not value: |
| return value |
| return (_move_image_to_device(value[0], device=device, dtype=dtype), *value[1:]) |
| if isinstance(value, list): |
| return [_move_image_to_device(item, device=device, dtype=dtype) for item in value] |
| raise TypeError(f'Unsupported processed image type: {type(value).__name__}') |
|
|
| def _tokenize_prediction_prompt_with_images(prompt, tokenizer, model, image_processor, image_tensor): |
| image_aspect_ratio = 'radio_h_native' |
| native_image_token_shapes = None |
| if 'native' in image_aspect_ratio: |
| native_images = image_tensor if isinstance(image_tensor, list) else [image_tensor] |
| native_image_token_shapes = [compute_num_image_tokens([item], image_aspect_ratio, image_processor) for item in native_images] |
| prompt = fix_one_prompt(prompt, model.config, native_image_token_shapes) |
| input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt', constant_version=_prediction_constant_version(model.config)) |
| if native_image_token_shapes is not None: |
| expected_slots = sum((math.prod(shape) for shape in native_image_token_shapes)) |
| actual_slots = int(input_ids.eq(IMAGE_PAD_TOKEN_INDEX).sum().item()) |
| if actual_slots != expected_slots: |
| raise RuntimeError(f'Native-image T2I prompt/token mismatch: expected {expected_slots} image slots but tokenization produced {actual_slots}.') |
| return (prompt, input_ids) |
|
|
| def _transformer_image_token_count(model, raw_token_count, gen_shape): |
| raw_token_count = int(raw_token_count) |
| raw_grid_tokens = math.prod(gen_shape) |
| if raw_token_count != raw_grid_tokens: |
| raise ValueError(f'Raw generation token count must match gen_shape: {raw_token_count} != {raw_grid_tokens}') |
| generation_model = model.get_model() |
| skip_down_sample = resolve_generation_skip_down_sample(generation_model.config, runtime_value=False) |
| compression_active = (getattr(generation_model, 'packed_image_codec', None) is not None or getattr(generation_model, 'downsample_gen', None) is not None) and (not skip_down_sample) |
| if not compression_active: |
| return raw_token_count |
| (height, width) = (int(value) for value in gen_shape) |
| return height // 2 * (width // 2) |
|
|
| def gumbel_noise(t, generator=None): |
| noise = torch.zeros_like(t).uniform_(0, 1, generator=generator) |
| return -torch.log(-torch.log(noise)) |
|
|
| def top_k_top_p_filtering(logits: torch.Tensor, top_k: int=0, top_p: float=1.0, filter_value: float=-1e309, min_tokens_to_keep: int=1) -> torch.Tensor: |
| if top_k > 0: |
| top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) |
| indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] |
| logits[indices_to_remove] = filter_value |
| if top_p < 1.0: |
| (sorted_logits, sorted_indices) = torch.sort(logits, descending=True) |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) |
| sorted_indices_to_remove = cumulative_probs > top_p |
| if min_tokens_to_keep > 1: |
| sorted_indices_to_remove[..., :min_tokens_to_keep] = 0 |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() |
| sorted_indices_to_remove[..., 0] = 0 |
| indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) |
| logits[indices_to_remove] = filter_value |
| return logits |
| from .constants import NemotronMinistralDLMConstant |
| NC = NemotronMinistralDLMConstant |
|
|
| def _resolve_ministral_sbd_diffusion_sampling(model, requested): |
| model_module = model.module if hasattr(model, 'module') else model |
| config = model_module.config |
| if requested is None: |
| requested = True |
| enabled = bool(requested) |
| return enabled |
|
|
| def _configure_ministral_t2i_attention(model_module, *, sbd_diffusion_sampling): |
| config = model_module.config |
| for layer in model_module.encoder.layers: |
| attention = layer.self_attn |
| if sbd_diffusion_sampling: |
| attention.set_attention_mode('sbd_block_diff', block_size=config.block_size) |
| else: |
| attention.mode = 'bidirectional' |
| if hasattr(attention, 'diffusion_lm'): |
| attention.diffusion_lm = True |
|
|
| def _restore_attention_state_after_prediction(function): |
|
|
| @functools.wraps(function) |
| def wrapped(model, *args, **kwargs): |
| model_module = model.module if hasattr(model, 'module') else model |
| encoder = getattr(model_module, 'encoder', None) |
| layers = getattr(encoder, 'layers', ()) |
| saved_states = [] |
| for layer in layers: |
| attention = layer.self_attn |
| saved_states.append((attention, getattr(attention, 'mode', None), getattr(attention, 'block_size', None), getattr(attention, 'diffusion_lm', None))) |
| try: |
| return function(model, *args, **kwargs) |
| finally: |
| for (attention, mode, block_size, diffusion_lm) in saved_states: |
| if hasattr(attention, 'mode'): |
| attention.mode = mode |
| if hasattr(attention, 'block_size'): |
| attention.block_size = block_size |
| if hasattr(attention, 'diffusion_lm'): |
| attention.diffusion_lm = diffusion_lm |
| return wrapped |
|
|
| def _load_semantic_tree_path(model_config, *, cache_owner, device, vocab_size): |
| artifact_path = getattr(model_config, 'image_token_tree_path', None) |
| if not isinstance(artifact_path, str) or not artifact_path.strip(): |
| raise ValueError('semantic-tree token sampling requires config.image_token_tree_path') |
| artifact_path = os.path.abspath(os.path.expanduser(artifact_path.strip())) |
| cache_key = (artifact_path, str(torch.device(device)), int(vocab_size)) |
| cache_entry = cache_owner.__dict__.get('_semantic_tree_token_path_cache') |
| if isinstance(cache_entry, tuple) and len(cache_entry) == 2 and (cache_entry[0] == cache_key): |
| return cache_entry[1] |
| tree = BalancedSemanticTree.from_artifact(production_semantic_tree_contract(artifact_path), expected_vocab_size=int(vocab_size)).to(device) |
| cache_owner.__dict__['_semantic_tree_token_path_cache'] = (cache_key, tree) |
| return tree |
|
|
| def _select_within_unstarted(scores, unstarted, count, *, multinomial=False): |
| if scores.ndim != 2 or unstarted.shape != scores.shape: |
| raise ValueError('scores and unstarted must have matching [batch, token] shapes') |
| if unstarted.dtype != torch.bool: |
| raise TypeError('unstarted must use torch.bool') |
| if scores.shape[0] != 1: |
| raise ValueError('delayed refinement selection currently requires batch size 1') |
| count = int(count.item()) if torch.is_tensor(count) else int(count) |
| eligible = unstarted[0].nonzero(as_tuple=False).flatten() |
| if count < 0 or count > eligible.numel(): |
| raise ValueError('selection count exceeds the unstarted-token set') |
| if count == 0: |
| return eligible[:0] |
| if count == eligible.numel(): |
| return eligible |
| eligible_scores = scores[0, eligible] |
| if multinomial: |
| local_indices = torch.multinomial(eligible_scores, num_samples=count) |
| else: |
| local_indices = torch.topk(eligible_scores, k=count).indices |
| return eligible[local_indices] |
|
|
| def _combine_edit_guidance_logits(logits_full, logits_unconditional, logits_vq_dropped, *, guidance_scale, guidance_scale_image, edit_mode, guidance_mode='legacy', logits_image_only=None): |
| if guidance_mode == 'legacy': |
| if edit_mode in [0, 3]: |
| logits_cond = (1 + guidance_scale_image) * logits_full - guidance_scale_image * logits_vq_dropped |
| elif edit_mode in [1, 2]: |
| if guidance_scale_image <= -1: |
| raise ValueError('legacy edit modes 1/2 require guidance_scale_image > -1') |
| logits_cond = (logits_full + guidance_scale_image * logits_vq_dropped) / (1 + guidance_scale_image) |
| else: |
| raise ValueError(f'Unsupported edit_mode: {edit_mode}') |
| return (1 + guidance_scale) * logits_cond - guidance_scale * logits_unconditional |
| if guidance_mode == 'independent': |
| return logits_full + guidance_scale * (logits_full - logits_unconditional) + guidance_scale_image * (logits_full - logits_vq_dropped) |
| if guidance_mode == 'ip2p': |
| if logits_image_only is None: |
| raise ValueError('ip2p guidance requires image-only logits') |
| return logits_unconditional + (1 + guidance_scale_image) * (logits_image_only - logits_unconditional) + (1 + guidance_scale) * (logits_full - logits_image_only) |
| raise ValueError("edit_guidance_mode must be 'legacy', 'independent', or 'ip2p'") |
|
|
| def _resolve_guidance_scale_schedule(guidance_scale_schedule, n_steps): |
| if guidance_scale_schedule is None: |
| return None |
| if not isinstance(guidance_scale_schedule, list): |
| raise TypeError('guidance_scale_schedule must be a list') |
| if isinstance(n_steps, bool) or not isinstance(n_steps, int) or n_steps < 1: |
| raise ValueError('guidance_scale_schedule requires a positive integer n_steps') |
| if len(guidance_scale_schedule) != n_steps: |
| raise ValueError(f'guidance_scale_schedule must contain exactly n_steps values: {len(guidance_scale_schedule)} != {n_steps}') |
| resolved = [] |
| for value in guidance_scale_schedule: |
| if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)): |
| raise TypeError('guidance_scale_schedule values must be real numbers') |
| value = float(value) |
| if not math.isfinite(value): |
| raise ValueError('guidance_scale_schedule values must be finite') |
| resolved.append(value) |
| return tuple(resolved) |
|
|
| def _resolve_semantic_tree_guidance_scale_range(guidance_scale_range): |
| if guidance_scale_range is None: |
| return None |
| if not isinstance(guidance_scale_range, (list, tuple)): |
| raise TypeError('semantic_tree_guidance_scale_range must be a two-value sequence') |
| if len(guidance_scale_range) != 2: |
| raise ValueError('semantic_tree_guidance_scale_range must contain exactly two values') |
| resolved = [] |
| for value in guidance_scale_range: |
| if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)): |
| raise TypeError('semantic_tree_guidance_scale_range values must be real numbers') |
| value = float(value) |
| if not math.isfinite(value): |
| raise ValueError('semantic_tree_guidance_scale_range values must be finite') |
| resolved.append(value) |
| return tuple(resolved) |
|
|
| def _resolve_semantic_tree_global_guidance_scale_range(guidance_scale_range): |
| if guidance_scale_range is None: |
| return None |
| if not isinstance(guidance_scale_range, (list, tuple)): |
| raise TypeError('semantic_tree_global_guidance_scale_range must be a two-value sequence') |
| if len(guidance_scale_range) != 2: |
| raise ValueError('semantic_tree_global_guidance_scale_range must contain exactly two values') |
| resolved = [] |
| for value in guidance_scale_range: |
| if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)): |
| raise TypeError('semantic_tree_global_guidance_scale_range values must be real numbers') |
| value = float(value) |
| if not math.isfinite(value): |
| raise ValueError('semantic_tree_global_guidance_scale_range values must be finite') |
| resolved.append(value) |
| return tuple(resolved) |
|
|
| def _resolve_semantic_tree_cfg_noise_interval(noise_interval): |
| if noise_interval is None: |
| return None |
| if not isinstance(noise_interval, (list, tuple)): |
| raise TypeError('semantic_tree_cfg_noise_interval must be a two-value sequence') |
| if len(noise_interval) != 2: |
| raise ValueError('semantic_tree_cfg_noise_interval must contain exactly two values') |
| resolved = [] |
| for value in noise_interval: |
| if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)): |
| raise TypeError('semantic_tree_cfg_noise_interval values must be real numbers') |
| value = float(value) |
| if not math.isfinite(value): |
| raise ValueError('semantic_tree_cfg_noise_interval values must be finite') |
| resolved.append(value) |
| (minimum, maximum) = resolved |
| if not 0.0 <= minimum <= maximum <= 1.0: |
| raise ValueError('semantic_tree_cfg_noise_interval must satisfy 0 <= min_q <= max_q <= 1') |
| return (minimum, maximum) |
|
|
| @torch.no_grad() |
| @_restore_attention_state_after_prediction |
| def text_to_image(model, prompt, confidence_policy='mmada', guidance_scale=5, n_steps=20, tokenizer=None, image_resolution=512, n_tokens=1024, shift=3, alg_temp=1, schedule='shift', min_temperature=0.01, dynamic_temperature=False, micro_cond='ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; SCORE : 6.5', temperature=1, schedule_temp='linear', shift_alg=None, top_p=None, top_k=None, unmask_order=None, cfg_interval=[0, 1], edit_image=None, image_processor=None, lock_source_padding_vq=None, edit_source_native_size=None, guidance_scale_image=5, edit_guidance_mode='legacy', edit_mode=0, template='Generate an image with the caption:\n <prompt>', plan=None, conv_template='llada', use_cache=None, cache_prompt=None, causal_context=True, edit_threshold=-1, disable_tqdm=False, ministral_sbd_sampling=None, ministral_full_attention_approx=True, ministral_shift_sbd_logits=None, image_embedding_refinement_steps=None, loop_k=None): |
| _validate_text_to_image_batch_size(1) |
| lock_source_padding_vq = resolve_source_padding_vq_lock(model.get_model().config, lock_source_padding_vq, has_edit_image=edit_image is not None) |
| if edit_source_native_size is not None and (not lock_source_padding_vq): |
| raise ValueError('edit_source_native_size is valid only with lock_source_padding_vq=true') |
| semantic_tree_refinement_gap = resolve_semantic_tree_refinement_gap(0) |
| semantic_tree_atomic_block_size = resolve_semantic_tree_atomic_block_size(1) |
| if shift_alg is None: |
| shift_alg = shift |
| resolved_loop_k = None |
| guidance_scale_schedule = _resolve_guidance_scale_schedule(None, n_steps) |
| semantic_tree_guidance_scale_range = _resolve_semantic_tree_guidance_scale_range(None) |
| semantic_tree_global_guidance_scale_range = _resolve_semantic_tree_global_guidance_scale_range(None) |
| semantic_tree_cfg_noise_interval = _resolve_semantic_tree_cfg_noise_interval(None) |
| guidance_scale_modes = sum((value is not None for value in (None, None, None))) |
| if guidance_scale_modes > 1: |
| raise ValueError('guidance_scale_schedule, semantic_tree_guidance_scale_range, and semantic_tree_global_guidance_scale_range are mutually exclusive') |
| device = model.get_model().device |
| use_image_local_ar_diffusion = False |
| if isinstance(edit_image, (str, os.PathLike)): |
| edit_image = _load_prediction_image(edit_image) |
| source_native_size = None |
| if lock_source_padding_vq: |
| source_native_size = _resolve_source_native_size(edit_source_native_size, fallback=edit_image.size) |
| conv_template = conv_template |
| use_mdm_prime = False |
| use_semantic_tree_tokens = True |
| semantic_tree_noise_cfg_requested = any((value is not None for value in (None, None, None))) |
| if semantic_tree_noise_cfg_requested: |
| if edit_image is not None: |
| raise NotImplementedError('semantic-tree noise-dependent CFG currently supports T2I only') |
| image_embedding_refinement_steps = resolve_semantic_tree_steps(model.get_model().config, image_embedding_refinement_steps) |
| use_refinement_state = True |
| reserve_token = '<|reserved_token_5|>' |
| reserve_token_2 = '<|reserved_token_6|>' |
| reserve_id = 126089 |
| reserve_id2 = 126090 |
| INT_MAX = 1000000 |
| is_emu3 = 'emu3' in 'emu3'.lower() |
| use_image_row_markers = image_row_markers_enabled(model.get_model().config) |
| fixed_prefix_count = image_fixed_prefix_token_count(model.get_model().config) |
| if fixed_prefix_count: |
| validate_image_fixed_prefix_support(model.get_model()) |
| row_marker_id = None |
| image_end_id = None |
| markers_per_row = None |
| skip_down_sample = resolve_generation_skip_down_sample(model.get_model().config, runtime_value=False) |
| (row_marker_id, _, image_end_id) = validate_image_row_marker_support(model.get_model(), skip_down_sample=skip_down_sample) |
| markers_per_row = image_row_marker_count(model.get_model(), skip_down_sample=skip_down_sample) |
| image_discrete_path = 'tree_metric' |
| use_tree_metric_path = False |
| use_affine_mixture_path = False |
| use_endpoint_probability_path = False |
| config = model.get_model().config |
| semantic_block_size = getattr(config, 'block_size', None) |
| if isinstance(semantic_block_size, bool) or not isinstance(semantic_block_size, int) or semantic_block_size < 1: |
| raise NotImplementedError('semantic_tree_token_diffusion requires a positive integer block_size') |
| if ministral_sbd_sampling is False: |
| raise NotImplementedError('semantic_tree_token_diffusion cannot disable checkpoint-native Ministral SBD sampling') |
| semantic_edit_controls_requested = edit_threshold > 0 |
| if semantic_edit_controls_requested: |
| raise NotImplementedError('semantic_tree_token_diffusion does not support post-commit edit/refinement controls') |
| codebook_size = int(model.get_model().vqvae.config.num_vq_embeddings) |
| (_, img_mask_id, _) = image_generation_special_ids(codebook_size) |
| txt_mask_id = 126336 |
| eot_id = 126384 |
| model_type = 'llada' |
| model_type = 'ministral' |
| is_ministral = True |
| use_ministral_sbd_sampling = False |
| _wte = wte_ministral |
| img_begin = NC.gen_im_start_token |
| img_end = NC.gen_im_end_token |
| if use_cache is None: |
| use_cache = True |
| if cache_prompt is None: |
| cache_prompt = True |
| conv_template = 'ministral' |
| reserve_id = NC.reserve_id |
| reserve_id2 = NC.reserve_id_enc |
| reserve_token = NC.reserve_id_token |
| reserve_token_2 = NC.reserve_id_token_enc |
| img_begin = NC.gen_im_start_token |
| img_end = NC.gen_im_end_token |
| txt_mask_id = NC.mask_id |
| eot_id = NC.eos_id |
| _wte = wte_ministral |
| model_module = model.module if hasattr(model, 'module') else model |
| use_ministral_sbd_sampling = _resolve_ministral_sbd_diffusion_sampling(model_module, ministral_sbd_sampling) |
| resolve_ministral_shift_sbd_logits(model_module, ministral_shift_sbd_logits) |
| _configure_ministral_t2i_attention(model_module, sbd_diffusion_sampling=use_ministral_sbd_sampling) |
| gen_shape_map = {1024: (64, 64), 512: (32, 32), 256: (16, 16)} |
| content_gen_shape = gen_shape_map[image_resolution] |
| if int(n_tokens) != math.prod(content_gen_shape): |
| raise ValueError(f'Raw generation token count must match the content image grid: {n_tokens} != {math.prod(content_gen_shape)}') |
| gen_shape = augmented_image_shape(content_gen_shape, markers_per_row=markers_per_row) |
| micro_cond = micro_cond |
| if edit_image is not None: |
| micro_cond = '' |
| prompt = f'{prompt} {micro_cond}' |
| question = template.replace('<prompt>', prompt) |
| image_tokens_txt = _transformer_image_token_count(model, math.prod(gen_shape), gen_shape) |
| n_tokens_txt = image_tokens_txt + fixed_prefix_count |
| if plan is None: |
| plan = '' |
| edit_source_vq_codes = None |
| if edit_image is not None: |
| image = edit_image |
| image_tensor = process_images([image], image_processor, model.config) |
| image_tensor = _move_image_to_device(image_tensor, device=model.device, dtype=torch.bfloat16) |
| modalities = ['image'] |
| image_sizes = [image.size] |
| conv = copy.deepcopy(conv_templates[conv_template]) |
| image = image.convert('RGB') |
| image_1024 = pad_to_square_and_resize(image, image_resolution) |
| (enc_latents, _gen_shape) = model.encode_image_gen(model.model.image_processor_gen.preprocess(image_1024).to(model.device, model.dtype), enc=True) |
| if lock_source_padding_vq: |
| if tuple((int(value) for value in _gen_shape)) != tuple(content_gen_shape): |
| raise RuntimeError(f'edit source VQ grid does not match generation grid: {tuple(_gen_shape)} != {tuple(content_gen_shape)}') |
| if enc_latents.ndim != 2 or tuple(enc_latents.shape) != (1, n_tokens): |
| raise RuntimeError(f'edit source VQ codes must have shape [1, {n_tokens}], got {tuple(enc_latents.shape)}') |
| edit_source_vq_codes = enc_latents.detach().clone() |
| enc_embeddings = model.model.call_gen_embedding(enc_latents, _gen_shape, enc=True) |
| conv.append_message(conv.roles[0], f'<image> {reserve_token_2 * enc_embeddings.shape[1]}\n {prompt} ') |
| conv.append_message(conv.roles[1], f'Sure {plan}{img_begin}{reserve_token * n_tokens_txt}{img_end}') |
| prompt_question = conv.get_prompt() |
| prompt_question.removesuffix('<|start_header_id|>assistant<|end_header_id|>\n\n') |
| (prompt_question, input_ids) = _tokenize_prediction_prompt_with_images(prompt_question, tokenizer, model, image_processor, image_tensor) |
| input_ids = input_ids.unsqueeze(0).to(model.device) |
| attention_mask = None |
| position_ids = None |
| image_sizes = [image.size] |
| modalities = ['image'] |
| (inputs, position_ids, attention_mask, _, inputs_embeds, _, raw_input_ids) = model.prepare_inputs_labels_for_multimodal(input_ids, position_ids, attention_mask, None, None, image_tensor, modalities, image_sizes=image_sizes, return_inputs=True) |
| expected_encoder_slots = int(enc_embeddings.shape[0] * enc_embeddings.shape[1]) |
| actual_encoder_slots = int(raw_input_ids.eq(reserve_id2).sum().item()) |
| if actual_encoder_slots != expected_encoder_slots: |
| raise RuntimeError(f'T2I edit encoder prompt/token mismatch: expected {expected_encoder_slots} slots but tokenization produced {actual_encoder_slots}.') |
| inputs_embeds[raw_input_ids == reserve_id2] = 0 |
| _enc = pad_along_last_dim(enc_embeddings, inputs_embeds.shape[-1]) |
| inputs_embeds[raw_input_ids == reserve_id2] = _enc.flatten(0, 1) |
| is_eot = torch.where(raw_input_ids == eot_id)[1] |
| assert len(is_eot) == 3 |
| prompt_cutoff = is_eot[1] |
| is_prompt = torch.zeros_like(raw_input_ids, dtype=torch.bool) |
| is_prompt[:, :prompt_cutoff + 1] = True |
| is_gen = raw_input_ids == reserve_id |
| is_gen_enc = raw_input_ids == reserve_id2 |
| else: |
| conv = copy.deepcopy(conv_templates[conv_template]) |
| conv.append_message(conv.roles[0], question) |
| conv.append_message(conv.roles[1], f'Sure {plan}{img_begin}{reserve_token * n_tokens_txt}{img_end}') |
| prompt_question = conv.get_prompt() |
| prompt_question.removesuffix('<|start_header_id|>assistant<|end_header_id|>\n\n') |
| input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(device) |
| is_gen = input_ids == reserve_id |
| is_gen_enc = input_ids == reserve_id2 |
| is_eot = torch.where(input_ids == eot_id)[1] |
| assert len(is_eot) == 3 |
| prompt_cutoff = is_eot[1] |
| is_prompt = torch.zeros_like(input_ids, dtype=torch.bool) |
| is_prompt[:, :prompt_cutoff + 1] = True |
| raw_input_ids = input_ids |
| (inputs_embeds, _) = _wte(model.model, raw_input_ids) |
| generation_placeholder_mask = is_gen |
| generated_prompt_slots = int(generation_placeholder_mask.sum().item()) |
| if generated_prompt_slots != n_tokens_txt: |
| raise RuntimeError(f'T2I prompt/token mismatch: expected {n_tokens_txt} generated-image slots but tokenization produced {generated_prompt_slots}.') |
| (fixed_prefix_mask, is_gen, fixed_prefix_ids) = split_generation_fixed_prefix(generation_placeholder_mask, codebook_size=int(model.get_model().vqvae.config.num_vq_embeddings), count=fixed_prefix_count) |
| if int(is_gen.sum().item()) != image_tokens_txt: |
| raise RuntimeError(f'T2I image-token mismatch after removing the fixed prefix: expected {image_tokens_txt}, found {int(is_gen.sum().item())}.') |
| generation_attention_mask = fixed_prefix_mask | is_gen |
| if fixed_prefix_ids.numel(): |
| fixed_prefix_embeddings = model.get_model().call_gen_fixed_token_embedding(fixed_prefix_ids) |
| inputs_embeds[fixed_prefix_mask] = pad_along_last_dim(fixed_prefix_embeddings, inputs_embeds.shape[-1]).reshape(-1, inputs_embeds.shape[-1]) |
| if plan: |
| is_gen_min = torch.where(raw_input_ids == reserve_id)[1][0] |
| is_prompt[:, :is_gen_min] = True |
| ministral_sbd_attention_window = None |
| use_exact_ministral_flash = bool(ministral_full_attention_approx) |
| use_ministral_prefix_cache = bool(use_cache and cache_prompt) |
| if use_ministral_sbd_sampling: |
| ministral_sbd_attention_window = generated_image_attention_window(generation_attention_mask) |
| use_ministral_prefix_cache = bool(use_cache and cache_prompt and causal_context) |
| if use_exact_ministral_flash and (not use_ministral_prefix_cache): |
| use_exact_ministral_flash = False |
| inputs_embeds_uncond = inputs_embeds.clone() |
| noise_embed = model.model.embed_tokens(torch.tensor([txt_mask_id]).to(model.device)) |
| inputs_embeds_uncond[is_prompt] = noise_embed |
| if edit_image is not None: |
| inputs_embeds_image_only = inputs_embeds_uncond.clone() |
| inputs_embeds_image_only[is_gen_enc] = inputs_embeds[is_gen_enc] |
| inputs_embeds_uncond_enc = inputs_embeds.clone() |
| if edit_mode == 0: |
| inputs_embeds_uncond_enc[~is_gen_enc] = noise_embed |
| is_gen_enc_ccc = is_gen_enc |
| elif edit_mode == 1: |
| inputs_embeds_uncond_enc[is_gen_enc] = noise_embed |
| is_gen_enc_ccc = torch.zeros_like(is_gen_enc, dtype=torch.bool) |
| elif edit_mode == 2: |
| inputs_embeds_uncond_enc[is_gen_enc | (raw_input_ids < 0)] = noise_embed |
| is_gen_enc_ccc = torch.zeros_like(is_gen_enc, dtype=torch.bool) |
| elif edit_mode == 3: |
| inputs_embeds_uncond_enc[~is_gen_enc & (raw_input_ids > 0)] = noise_embed |
| is_gen_enc_ccc = is_gen_enc |
| else: |
| raise ValueError('Not Supported edit_mode') |
| inputs_embeds_uncond_enc[fixed_prefix_mask] = inputs_embeds[fixed_prefix_mask] |
| image_gen_latents_offset = torch.zeros(1, n_tokens, dtype=torch.long, device=input_ids.device) |
| is_unitok = False |
| structural_mask = None |
| endpoint_probability_path = None |
| semantic_tree = None |
| semantic_representatives = None |
| image_vocab_size = int(model.get_model().vqvae.config.num_vq_embeddings) |
| semantic_tree = _load_semantic_tree_path(model.get_model().config, cache_owner=model.get_model(), device=input_ids.device, vocab_size=image_vocab_size) |
| if image_embedding_refinement_steps > semantic_tree.depth: |
| raise ValueError(f'semantic-tree delayed commit steps cannot exceed tree depth: {image_embedding_refinement_steps} > {semantic_tree.depth}') |
| semantic_representatives = semantic_tree.sample_representatives(torch.zeros_like(image_gen_latents_offset), levels=0) |
| image_gen_latents_offset[:] = img_mask_id |
| xt = image_gen_latents_offset.clone() |
| (xt, marker_gen_shape, structural_mask) = append_image_row_markers(xt, content_gen_shape, row_marker_id=row_marker_id, image_end_id=image_end_id, markers_per_row=markers_per_row) |
| if marker_gen_shape != gen_shape: |
| raise RuntimeError('inference marker grid does not match the prompt grid') |
| (semantic_representatives, semantic_marker_shape, semantic_structural_mask) = append_image_row_markers(semantic_representatives, content_gen_shape, row_marker_id=row_marker_id, image_end_id=image_end_id, markers_per_row=markers_per_row) |
| if semantic_marker_shape != gen_shape or not torch.equal(semantic_structural_mask, structural_mask): |
| raise RuntimeError('semantic-tree marker layout does not match generation') |
| if structural_mask is None: |
| content_positions = torch.arange(xt.shape[-1], device=xt.device, dtype=torch.long) |
| else: |
| content_positions = (~structural_mask[0]).nonzero(as_tuple=False).flatten() |
| source_padding_raw_mask = None |
| source_padding_lock_mask = None |
| source_padding_locked_codes = None |
| source_padding_locked_embeddings = None |
| if lock_source_padding_vq: |
| if edit_source_vq_codes is None or source_native_size is None: |
| raise RuntimeError('source-padding VQ locking lost its source codes or geometry') |
| source_padding_raw_mask = _source_padding_vq_mask(source_native_size, content_gen_shape, device=xt.device).expand(edit_source_vq_codes.shape[0], -1) |
| (_, source_padding_lock_mask, source_padding_locked_codes) = _prefill_source_padding_vq(xt, edit_source_vq_codes.to(device=xt.device, dtype=xt.dtype), source_padding_raw_mask, content_positions) |
| source_padding_locked_embeddings = model.get_model().call_gen_token_embedding(source_padding_locked_codes) |
| sampling_protected_mask = structural_mask |
| refinement_child_embeddings = None |
| refinement_stage = None |
| semantic_next_refinement_round = None |
| refinement_stage = torch.full_like(xt, -1) |
| semantic_next_refinement_round = torch.zeros_like(refinement_stage) |
| if structural_mask is not None: |
| refinement_stage[structural_mask] = image_embedding_refinement_steps |
| semantic_noise_levels = semantic_tree_noise_for_ages(refinement_stage.clamp_min(0), steps=image_embedding_refinement_steps) |
| refinement_child_embeddings = build_semantic_tree_embeddings(semantic_representatives, semantic_noise_levels, model.get_model().call_gen_token_embedding, img_mask_id) |
| prime_digits = None |
| endpoint_sampling_times = None |
| mask_idx = xt == img_mask_id |
| n_mask = mask_idx.sum() |
| mask_idx_sch = mask_idx |
| semantic_counts = [semantic_tree_activation_counts(int(row_mask.sum().item()), activation_steps=n_steps, schedule=schedule, shift=shift, device=mask_idx_sch.device) for row_mask in mask_idx_sch] |
| if any((row_counts.shape != semantic_counts[0].shape for row_counts in semantic_counts[1:])): |
| raise ValueError('semantic-tree inference requires equal activation-step counts across batch rows') |
| num_transfer_tokens = torch.stack(semantic_counts) |
| activation_steps = num_transfer_tokens.shape[1] |
| decode_num_transfer_tokens = semantic_tree_decode_schedule(num_transfer_tokens, refinement_steps=image_embedding_refinement_steps, refinement_gap=0) |
| sch_temperatures = torch.linspace(0, 1, n_steps, device='cpu').numpy() |
| if schedule_temp == 'linear': |
| sch_temperatures = (1 - sch_temperatures) * (1 - min_temperature) + min_temperature |
| elif schedule_temp == 'cosine2': |
| sch_temperatures = cosine_schedule_2(1 - sch_temperatures) * (1 - min_temperature) + min_temperature |
| elif schedule_temp == 'shift': |
| sch_temperatures = logit_normal_schedule(shift_alg, 1 - sch_temperatures) * (1 - min_temperature) + min_temperature |
| else: |
| raise NotImplementedError(f'Unknown schedule_temp: {schedule_temp}') |
| sch_temperatures_samp = torch.linspace(0, 1, n_steps, device='cpu').numpy() |
| sch_temperatures_samp = (1 - sch_temperatures_samp) * (1 - 1) + 1 |
| sch_temperatures = torch.tensor(sch_temperatures, device=device, dtype=torch.float32) |
| sch_temperatures_samp = torch.tensor(sch_temperatures_samp, device=device, dtype=torch.float32) |
| temp_idx = 0 |
| cfg_start = int(cfg_interval[0] * n_steps) |
| cfg_end = int(cfg_interval[1] * n_steps) |
| use_optional_guidance_scale = any((value is not None for value in (None, None, None, None))) |
| scheduled_run_cfg_branches = None |
| if isinstance(unmask_order, str): |
| unmask_order = torch.load(unmask_order).tolist() |
| eligible_parent_positions = None |
| if confidence_policy == 'stratified' and unmask_order is None: |
| stratified_tokens = n_tokens |
| _dim = int(np.sqrt(stratified_tokens)) |
| if _dim * _dim != stratified_tokens: |
| raise ValueError('stratified local AR sampling requires a square parent grid') |
| seq = stratified_random(n=_dim, seed=42, shuffle_blocks=True) |
| unmask_order = seq |
| stratified_token_count = n_tokens |
| if confidence_policy == 'stratified': |
| assert unmask_order is not None |
| n_c = 0 |
| cat_cfg_time = 0 |
| post_cfg_time = 0 |
| wte_time = 0 |
| fwd_time = 0 |
| sample_time = 0 |
| all_time = 0 |
| total_edited = 0 |
| edit_transitions = [] |
| last_edit_step = torch.full_like(xt, -10000) |
| distribution_trace = {'initial_commits': [], 'edit_proposals': [], 'edits': []} |
| initial_trace_top_ids = None |
| initial_trace_top_logits = None |
| initial_trace_top_probs = None |
| initial_trace_top_sampling_probs = None |
| previous_trace_top_ids = None |
| previous_trace_top_logits = None |
| previous_trace_top_probs = None |
| intermediate_x0s = [] |
| for (decode_step_idx, num_transfer) in tqdm(enumerate(decode_num_transfer_tokens[0]), total=decode_num_transfer_tokens.shape[1], disable=disable_tqdm): |
| t0 = time.time() |
| schedule_temp_idx = temp_idx |
| schedule_temp_idx = min(schedule_temp_idx, min(activation_steps, n_steps) - 1) |
| local_temp = sch_temperatures[schedule_temp_idx] |
| local_temp_samp = sch_temperatures_samp[schedule_temp_idx] |
| temp_idx += 1 |
| schedule_progress_idx = temp_idx |
| schedule_progress_idx = min(schedule_progress_idx, activation_steps) |
| if schedule_progress_idx / n_steps > 100: |
| confidence_policy = 'mmada' |
| hard_mask_idx = refinement_stage < image_embedding_refinement_steps |
| mask_idx = refinement_stage < 0 |
| n_mask = mask_idx.sum() |
| timestep_refinement_stage = refinement_stage[:, content_positions] |
| timesteps = refinement_mask_fraction(timestep_refinement_stage, steps=image_embedding_refinement_steps) |
| semantic_tree_cfg_q = None |
| cfg_step_idx = temp_idx |
| cfg_step_idx = min(cfg_step_idx, activation_steps) |
| guidance_scale_for_step = guidance_scale |
| if scheduled_run_cfg_branches is None: |
| run_cfg_branches = guidance_scale != 0 or guidance_scale_image != 0 |
| else: |
| run_cfg_branches = scheduled_run_cfg_branches |
| do_cfg = run_cfg_branches and cfg_start <= cfg_step_idx and (cfg_end >= cfg_step_idx) |
| model_xt = _restore_source_padding_vq(xt, source_padding_lock_mask, source_padding_locked_codes) |
| if run_cfg_branches: |
| if edit_image is not None: |
| if edit_guidance_mode == 'ip2p': |
| input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds_image_only, inputs_embeds]).clone() |
| else: |
| input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds_uncond_enc, inputs_embeds]).clone() |
| xt_input = torch.cat([model_xt, model_xt, model_xt]) |
| new_token_mask = is_gen.repeat(3, 1) |
| fixed_prefix_routing = fixed_prefix_mask.repeat(3, 1) |
| is_gen_enc_null = torch.zeros_like(is_gen_enc, dtype=torch.bool) |
| if edit_guidance_mode == 'ip2p': |
| is_gen_enc_mask = torch.cat([is_gen_enc_null, is_gen_enc, is_gen_enc]) |
| else: |
| is_gen_enc_mask = torch.cat([is_gen_enc_null, is_gen_enc_ccc, is_gen_enc]) |
| timesteps = timesteps.repeat(3) |
| else: |
| input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds]).clone() |
| xt_input = torch.cat([model_xt, model_xt]) |
| new_token_mask = is_gen.repeat(2, 1) |
| fixed_prefix_routing = fixed_prefix_mask.repeat(2, 1) |
| is_gen_enc_mask = is_gen_enc.repeat(2, 1) |
| is_gen_enc_mask[0, :] = False |
| timesteps = timesteps.repeat(2) |
| else: |
| input_embeddings_input = inputs_embeds.clone() |
| new_token_mask = is_gen |
| fixed_prefix_routing = fixed_prefix_mask |
| xt_input = model_xt |
| is_gen_enc_mask = is_gen_enc |
| enc_use_image_branch = True |
| generation_attention_mask_input = fixed_prefix_routing | new_token_mask |
| if enc_use_image_branch: |
| modality_indices = generation_attention_mask_input | is_gen_enc_mask |
| else: |
| modality_indices = generation_attention_mask_input |
| t1 = time.time() |
| x_gen_input = xt_input |
| x_gen_child_embeddings = None |
| if xt_input.shape[0] % refinement_child_embeddings.shape[0] != 0: |
| raise RuntimeError('CFG batch does not align with the refinement state batch') |
| cfg_repeats = xt_input.shape[0] // refinement_child_embeddings.shape[0] |
| step_refinement_child_embeddings = refinement_child_embeddings |
| if source_padding_lock_mask is not None: |
| if source_padding_locked_embeddings.shape != refinement_child_embeddings.shape: |
| raise RuntimeError('source-padding embeddings do not align with the refinement state') |
| step_refinement_child_embeddings = refinement_child_embeddings.clone() |
| step_refinement_child_embeddings[source_padding_lock_mask] = source_padding_locked_embeddings[source_padding_lock_mask] |
| x_gen_input = None |
| x_gen_child_embeddings = step_refinement_child_embeddings.repeat(cfg_repeats, 1, 1) |
| (all_input_embeddings, new_token_mask) = _wte(model.get_model(), None, True, x_gen=x_gen_input, gen_shape=gen_shape, inputs_embeds_curr=input_embeddings_input, new_token_mask=new_token_mask, x_gen_child_embeddings=x_gen_child_embeddings) |
| t2 = time.time() |
| if use_ministral_prefix_cache: |
| if decode_step_idx == 0: |
| cache_cutoff = ministral_sbd_attention_window[0] if ministral_sbd_attention_window is not None else prompt_cutoff |
| if causal_context: |
| for layer in model_module.encoder.layers: |
| if hasattr(layer.self_attn, 'diffusion_lm'): |
| layer.self_attn.diffusion_lm = False |
| if cache_cutoff > 0: |
| output = model.get_model()(None, input_embeddings=all_input_embeddings[:, :cache_cutoff], modality_indices=modality_indices[:, :cache_cutoff], output_hidden_states=True, past_key_values=None, use_cache=True, overwrite_attn_impl='flash_attn') |
| past_key_values = output.past_key_values |
| cache_len = past_key_values.get_seq_length() |
| if cache_len != cache_cutoff: |
| raise RuntimeError(f'SBD prefix cache length does not match its generated-image boundary: {cache_len} != {cache_cutoff}') |
| else: |
| past_key_values = None |
| cache_len = 0 |
| if causal_context: |
| for layer in model_module.encoder.layers: |
| if hasattr(layer.self_attn, 'diffusion_lm'): |
| layer.self_attn.diffusion_lm = True |
| else: |
| pass |
| else: |
| past_key_values = None |
| cache_len = 0 |
| sampling_end = ministral_sbd_attention_window[1] if ministral_sbd_attention_window is not None and use_ministral_prefix_cache else all_input_embeddings.shape[1] |
| if not cache_len <= sampling_end <= all_input_embeddings.shape[1]: |
| raise RuntimeError(f'Invalid exact-image SBD sampling window: cache_len={cache_len}, sampling_end={sampling_end}, sequence_length={all_input_embeddings.shape[1]}') |
| logits = get_logits_ministral(model.get_model(), all_input_embeddings[:, cache_len:sampling_end], new_token_mask[:, cache_len:sampling_end], True, gen_shape=gen_shape, input_modality_indices=modality_indices[:, cache_len:sampling_end], timesteps=timesteps, past_key_values=past_key_values, generation_attention_mask=generation_attention_mask_input[:, cache_len:sampling_end], sbd_diffusion_sampling=use_ministral_sbd_sampling, full_attention_approx=use_exact_ministral_flash, loop_k=None) |
| t3 = time.time() |
| if run_cfg_branches: |
| if edit_image is not None: |
| (new_token_mask, _, _) = new_token_mask.chunk(3) |
| (logits_un, logits_middle, logits) = logits.chunk(3) |
| if do_cfg: |
| logits_is_ninf = logits == -np.inf |
| logits = _combine_edit_guidance_logits(logits, logits_un, logits_middle if edit_guidance_mode != 'ip2p' else None, guidance_scale=guidance_scale_for_step, guidance_scale_image=guidance_scale_image, edit_mode=edit_mode, guidance_mode=edit_guidance_mode, logits_image_only=logits_middle if edit_guidance_mode == 'ip2p' else None) |
| logits[logits_is_ninf] = -np.inf |
| else: |
| (new_token_mask, _) = new_token_mask.chunk(2) |
| (logits_un, logits) = logits.chunk(2) |
| if do_cfg: |
| logits_is_ninf = logits == -np.inf |
| if use_optional_guidance_scale: |
| active_guidance_scale = guidance_scale |
| logits = (1 + active_guidance_scale) * logits - active_guidance_scale * logits_un |
| else: |
| logits = (1 + guidance_scale) * logits - guidance_scale * logits_un |
| logits[logits_is_ninf] = -np.inf |
| t4 = time.time() |
| if top_p is not None or top_k is not None: |
| (_b, _l, _v) = logits.shape |
| logits = logits.view(_b * _l, _v) |
| logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p, filter_value=-np.inf, min_tokens_to_keep=1) |
| logits = logits.view(_b, _l, _v) |
| probs = logits.softmax(dim=-1) |
| trace_top_ids = None |
| trace_top_logits = None |
| trace_top_probs = None |
| _temperature = temperature |
| sampling_distribution = dists.Categorical(logits=logits / _temperature) |
| x0 = sampling_distribution.sample() |
| del sampling_distribution |
| x0_p = torch.gather(probs, -1, x0.long()[..., None]).squeeze(-1) |
| if edit_threshold <= 0: |
| x0 = torch.where(hard_mask_idx, x0, xt) |
| if decode_step_idx >= activation_steps: |
| select_index = torch.empty(0, dtype=torch.long, device=x0.device) |
| elif confidence_policy == 'mask_git': |
| _alg_temp = alg_temp |
| if dynamic_temperature: |
| _alg_temp = _alg_temp * local_temp |
| confidence = x0_p / _alg_temp |
| confidence = torch.where(mask_idx, confidence, -np.inf) |
| confidence = torch.softmax(confidence, dim=-1) |
| select_index = _select_within_unstarted(confidence, mask_idx, num_transfer, multinomial=True) |
| elif confidence_policy == 'mmada': |
| _alg_temp = alg_temp |
| if dynamic_temperature: |
| _alg_temp = _alg_temp * local_temp |
| confidence = torch.log(x0_p.clamp(1e-20)) + _alg_temp * gumbel_noise(x0_p, generator=None) |
| confidence = torch.where(mask_idx, confidence, -np.inf) |
| select_index = _select_within_unstarted(confidence, mask_idx, num_transfer) |
| elif confidence_policy == 'stratified': |
| start = stratified_token_count - int(n_mask.item()) |
| select_index = torch.tensor(unmask_order[start:start + int(num_transfer)], device=x0.device, dtype=torch.long) |
| if structural_mask is not None: |
| select_index = content_positions[select_index] |
| else: |
| raise NotImplementedError(f'Unknown confidence_policy: {confidence_policy}') |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device) |
| j = 0 |
| transfer_index[j, select_index] = True |
| proposal_child_embeddings = None |
| activation_index = transfer_index & (refinement_stage < 0) |
| refinement_stage[activation_index] = 0 |
| semantic_next_refinement_round[activation_index] = decode_step_idx + 0 |
| pending_index = (refinement_stage >= 0) & (refinement_stage < image_embedding_refinement_steps) |
| proposal_child_embeddings = model.get_model().call_gen_token_embedding(x0) |
| pending_index &= semantic_tree_refinement_due(refinement_stage, semantic_next_refinement_round, decode_round=decode_step_idx, refinement_steps=image_embedding_refinement_steps) |
| commit_index = torch.zeros_like(pending_index) |
| if pending_index.any(): |
| semantic_advance = advance_revisable_semantic_tree(semantic_representatives[pending_index], x0[pending_index], refinement_stage[pending_index], semantic_tree, steps=image_embedding_refinement_steps) |
| semantic_representatives[pending_index] = semantic_advance.representatives |
| refinement_stage[pending_index] = semantic_advance.ages |
| semantic_next_refinement_round[pending_index] = decode_step_idx + 0 + 1 |
| commit_index[pending_index] = semantic_advance.committed |
| semantic_noise_levels = semantic_tree_noise_for_ages(refinement_stage.clamp_min(0), steps=image_embedding_refinement_steps) |
| refinement_child_embeddings = build_semantic_tree_embeddings(semantic_representatives, semantic_noise_levels, model.get_model().call_gen_token_embedding, img_mask_id) |
| xt[commit_index] = semantic_representatives[commit_index] |
| transfer_index = commit_index |
| xt_is_mask = xt == img_mask_id |
| edit_progress = (decode_step_idx + 1) / decode_num_transfer_tokens.shape[1] |
| edits_enabled_this_round = edit_threshold > 0 and 0.0 <= edit_progress <= 1.0 and (image_embedding_refinement_steps == 1 or decode_step_idx < activation_steps) |
| if edits_enabled_this_round: |
| editable_position = ~xt_is_mask & ~transfer_index |
| if structural_mask is not None: |
| editable_position &= ~structural_mask |
| editing_confidence = torch.where(editable_position, x0_p, -torch.inf) |
| high_conf_editing = editing_confidence > edit_threshold |
| margin_accepted = torch.ones_like(editable_position) |
| proposal_logit_margin = None |
| token_changed = (x0 != xt) & high_conf_editing & margin_accepted |
| num_edited = token_changed.sum() |
| if num_edited > 0: |
| xt[token_changed] = x0[token_changed] |
| last_edit_step[token_changed] = decode_step_idx |
| refinement_child_embeddings[token_changed] = proposal_child_embeddings[token_changed] |
| refinement_stage[token_changed] = image_embedding_refinement_steps |
| total_edited += num_edited.item() |
| t5 = time.time() |
| cat_cfg_time += t1 - t0 |
| post_cfg_time += t2 - t1 |
| wte_time += t3 - t2 |
| fwd_time += t4 - t3 |
| sample_time += t5 - t4 |
| all_time += t5 - t0 |
| n_c += 1 |
| unresolved_refinement = refinement_stage < image_embedding_refinement_steps |
| if unresolved_refinement.any() or (xt == img_mask_id).any(): |
| stage_counts = {stage: (refinement_stage == stage).sum().item() for stage in range(-1, image_embedding_refinement_steps + 1)} |
| raise RuntimeError(f'image embedding refinement ended with unresolved tokens: unfinished_stages={unresolved_refinement.sum().item()}, committed_mask_tokens={(xt == img_mask_id).sum().item()}, stage_counts={stage_counts}') |
| xt = _restore_source_padding_vq(xt, source_padding_lock_mask, source_padding_locked_codes) |
| _assert_source_padding_vq_locked(xt, source_padding_lock_mask, source_padding_locked_codes) |
| if structural_mask is not None: |
| expected_row_id = int(row_marker_id) |
| expected_end_id = int(image_end_id) |
| marker_grid = xt.reshape(xt.shape[0], *gen_shape) |
| if not bool((marker_grid[:, :-1, -markers_per_row:] == expected_row_id).all() and (marker_grid[:, -1:, -markers_per_row:] == expected_end_id).all()): |
| raise RuntimeError('structural generation image markers changed during sampling') |
| xt = strip_image_row_markers(xt, content_gen_shape, markers_per_row=markers_per_row) |
| x0_img = xt |
| result = Image.fromarray(model.decode_image_gen(x0_img, image_resolution, image_resolution)[0]) |
| return result |
|
|
