sampling.py

# ADOBE CONFIDENTIAL
# Copyright 2025 Adobe
# All Rights Reserved.
# NOTICE: All information contained herein is, and remains
# the property of Adobe and its suppliers, if any. The intellectual
# and technical concepts contained herein are proprietary to Adobe
# and its suppliers and are protected by all applicable intellectual
# property laws, including trade secret and copyright laws.
# Dissemination of this information or reproduction of this material
# is strictly forbidden unless prior written permission is obtained
# from Adobe.


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