|
|
| from __future__ import annotations |
| import math |
| from numbers import Integral |
| import torch |
| from torch import nn |
| GEN_IMAGE_ROW_MARKERS = 'gen_image_row_markers' |
| GEN_IMAGE_FIXED_PREFIX_TOKENS = 'gen_image_fixed_prefix_tokens' |
| GENERATION_SKIP_DOWNSAMPLE = 'generation_skip_down_sample' |
| ROW_MARKER_OFFSET = 0 |
| IMAGE_MASK_OFFSET = 1 |
| IMAGE_END_OFFSET = 2 |
| FIXED_PREFIX_OFFSET = 3 |
| RESERVED_GENERATION_ROWS = 256 |
| MARKERS_PER_RAW_ROW = 2 |
| UNCOMPRESSED_MARKERS_PER_RAW_ROW = 1 |
| _SUPPORTED_MARKERS_PER_ROW = (UNCOMPRESSED_MARKERS_PER_RAW_ROW, MARKERS_PER_RAW_ROW) |
|
|
| def _validate_markers_per_row(value) -> int: |
| if isinstance(value, bool) or not isinstance(value, Integral): |
| raise TypeError('markers_per_row must be an integer') |
| value = int(value) |
| if value not in _SUPPORTED_MARKERS_PER_ROW: |
| raise ValueError('markers_per_row must be 1 for an uncompressed image grid or 2 for a grid that will undergo 2x spatial compression') |
| return value |
|
|
| def _validate_fixed_prefix_count(value) -> int: |
| if value is None: |
| return 0 |
| if isinstance(value, bool) or not isinstance(value, Integral): |
| raise TypeError(f'config.{GEN_IMAGE_FIXED_PREFIX_TOKENS} must be an integer') |
| value = int(value) |
| max_tokens = RESERVED_GENERATION_ROWS - FIXED_PREFIX_OFFSET |
| if not 0 <= value <= max_tokens: |
| raise ValueError(f'config.{GEN_IMAGE_FIXED_PREFIX_TOKENS} must be between 0 and {max_tokens}') |
| return value |
|
|
| def image_row_markers_enabled(config) -> bool: |
| value = getattr(config, GEN_IMAGE_ROW_MARKERS, False) |
| if value is None: |
| return False |
| if not isinstance(value, bool): |
| raise TypeError(f'config.{GEN_IMAGE_ROW_MARKERS} must be boolean') |
| return value |
|
|
| def resolve_generation_skip_down_sample(config, *, runtime_value=None) -> bool: |
| configured = getattr(config, GENERATION_SKIP_DOWNSAMPLE, None) |
| if configured is not None and (not isinstance(configured, bool)): |
| raise TypeError(f'config.{GENERATION_SKIP_DOWNSAMPLE} must be boolean or null') |
| runtime_configured = runtime_value is not None |
| runtime_enabled = bool(runtime_value) |
| if configured is None: |
| return runtime_enabled |
| if runtime_configured and runtime_enabled != configured: |
| raise ValueError(f'SKIP_DOWN_SAMPLE conflicts with checkpoint-backed config.{GENERATION_SKIP_DOWNSAMPLE}: runtime={runtime_enabled}, checkpoint={configured}') |
| return configured |
|
|
| def image_generation_special_ids(codebook_size: int) -> tuple[int, int, int]: |
| if isinstance(codebook_size, bool) or not isinstance(codebook_size, Integral): |
| raise TypeError('codebook_size must be a positive integer') |
| codebook_size = int(codebook_size) |
| if codebook_size < 1: |
| raise ValueError('codebook_size must be a positive integer') |
| return (codebook_size + ROW_MARKER_OFFSET, codebook_size + IMAGE_MASK_OFFSET, codebook_size + IMAGE_END_OFFSET) |
|
|
| def image_fixed_prefix_token_count(config) -> int: |
| return _validate_fixed_prefix_count(getattr(config, GEN_IMAGE_FIXED_PREFIX_TOKENS, 0)) |
|
|
| def image_generation_fixed_prefix_ids(codebook_size: int, count: int) -> tuple[int, ...]: |
| image_generation_special_ids(codebook_size) |
| count = _validate_fixed_prefix_count(count) |
| start = int(codebook_size) + FIXED_PREFIX_OFFSET |
| return tuple(range(start, start + count)) |
|
|
| def split_generation_fixed_prefix(generation_placeholder_mask: torch.Tensor, *, codebook_size: int, count: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| if not torch.is_tensor(generation_placeholder_mask): |
| raise TypeError('generation_placeholder_mask must be a tensor') |
| if generation_placeholder_mask.ndim != 2: |
| raise ValueError('generation_placeholder_mask must have shape [batch, sequence]') |
| placeholder_mask = generation_placeholder_mask.to(dtype=torch.bool) |
| fixed_ids = image_generation_fixed_prefix_ids(codebook_size, count) |
| if not fixed_ids: |
| return (torch.zeros_like(placeholder_mask), placeholder_mask, torch.empty(0, dtype=torch.long, device=placeholder_mask.device)) |
| previous_placeholder = torch.cat((torch.zeros_like(placeholder_mask[:, :1]), placeholder_mask[:, :-1]), dim=1) |
| run_starts = placeholder_mask & ~previous_placeholder |
| positions = torch.arange(placeholder_mask.shape[1], device=placeholder_mask.device, dtype=torch.long).unsqueeze(0).expand_as(placeholder_mask) |
| most_recent_start = torch.where(run_starts, positions, positions.new_full((), -1)).cummax(dim=1).values |
| run_offsets = positions - most_recent_start |
| fixed_mask = placeholder_mask & run_offsets.lt(len(fixed_ids)) |
| image_mask = placeholder_mask & ~fixed_mask |
| next_placeholder = torch.cat((placeholder_mask[:, 1:], torch.zeros_like(placeholder_mask[:, :1])), dim=1) |
| run_ends = placeholder_mask & ~next_placeholder |
| run_lengths = run_offsets[run_ends] + 1 |
| if bool(run_lengths.le(len(fixed_ids)).any()): |
| shortest = int(run_lengths.min().item()) |
| raise ValueError(f'each generated-image placeholder run must contain the fixed prefix plus at least one image slot; shortest run={shortest}, fixed prefix={len(fixed_ids)}') |
| dense_ids = int(codebook_size) + FIXED_PREFIX_OFFSET + run_offsets |
| return (fixed_mask, image_mask, dense_ids[fixed_mask]) |
|
|
| def generation_embedding_capacity(generation_model) -> int: |
| embedding = getattr(generation_model, 'gen_embedding', None) |
| if isinstance(embedding, nn.Embedding): |
| return int(embedding.num_embeddings) |
| if isinstance(embedding, nn.Module): |
| for module in embedding.modules(): |
| if isinstance(module, nn.Embedding): |
| return int(module.num_embeddings) |
| raise TypeError('generation model does not expose an nn.Embedding lookup') |
|
|
| def image_row_marker_count(generation_model, *, skip_down_sample=False) -> int: |
| compression_available = getattr(generation_model, 'packed_image_codec', None) is not None or getattr(generation_model, 'downsample_gen', None) is not None |
| if compression_available and (not skip_down_sample): |
| return MARKERS_PER_RAW_ROW |
| return UNCOMPRESSED_MARKERS_PER_RAW_ROW |
|
|
| def validate_image_row_marker_support(generation_model, *, skip_down_sample=False): |
| config = generation_model.config |
| if 'emu3' not in 'emu3'.lower(): |
| raise NotImplementedError("generation image row markers currently require Emu3's single-codebook tokenizer") |
| image_row_marker_count(generation_model, skip_down_sample=skip_down_sample) |
| codebook_size = int(generation_model.vqvae.config.num_vq_embeddings) |
| (row_id, mask_id, end_id) = image_generation_special_ids(codebook_size) |
| capacity = generation_embedding_capacity(generation_model) |
| if end_id >= capacity: |
| raise ValueError(f'generation embedding has no spare rows for row/end markers: required ID {end_id}, capacity {capacity}') |
| return (row_id, mask_id, end_id) |
|
|
| def validate_image_fixed_prefix_support(generation_model) -> tuple[int, ...]: |
| config = generation_model.config |
| count = image_fixed_prefix_token_count(config) |
| if count == 0: |
| return () |
| if getattr(config, 'generation_embedding_scale', None) is None: |
| raise ValueError('generation fixed-prefix tokens require factorized generation scaling so uncompressed tokens receive the common input scale') |
| if 'emu3' not in 'emu3'.lower(): |
| raise NotImplementedError("generation fixed-prefix tokens currently require Emu3's single-codebook tokenizer") |
| codebook_size = int(generation_model.vqvae.config.num_vq_embeddings) |
| fixed_ids = image_generation_fixed_prefix_ids(codebook_size, count) |
| capacity = generation_embedding_capacity(generation_model) |
| if fixed_ids[-1] >= capacity: |
| raise ValueError(f'generation embedding has no spare rows for the requested fixed prefix: required ID {fixed_ids[-1]}, capacity {capacity}') |
| special_ids = set(image_generation_special_ids(codebook_size)) |
| if special_ids.intersection(fixed_ids): |
| raise RuntimeError('generation fixed-prefix IDs overlap image mask/row/end markers') |
| return fixed_ids |
|
|
| def augmented_image_shape(gen_shape, *, markers_per_row: int=MARKERS_PER_RAW_ROW) -> tuple[int, int]: |
| if gen_shape is None or len(gen_shape) != 2: |
| raise ValueError('gen_shape must contain image height and width') |
| markers_per_row = _validate_markers_per_row(markers_per_row) |
| (height, width) = (int(value) for value in gen_shape) |
| if height < 1 or width < 1: |
| raise ValueError('image row markers require a positive image grid') |
| if markers_per_row == MARKERS_PER_RAW_ROW and (height % 2 or width % 2): |
| raise ValueError('two-column image row markers require an even-height, even-width grid for 2x spatial compression') |
| return (height, width + markers_per_row) |
|
|
| def append_image_row_markers(token_ids: torch.Tensor, gen_shape, *, row_marker_id: int, image_end_id: int, markers_per_row: int=MARKERS_PER_RAW_ROW) -> tuple[torch.Tensor, tuple[int, int], torch.Tensor]: |
| if not torch.is_tensor(token_ids) or token_ids.ndim != 2: |
| raise ValueError('token_ids must have shape [batch, raw_image_tokens]') |
| markers_per_row = _validate_markers_per_row(markers_per_row) |
| (height, width) = (int(value) for value in gen_shape) |
| augmented_shape = augmented_image_shape((height, width), markers_per_row=markers_per_row) |
| if token_ids.shape[-1] != height * width: |
| raise ValueError(f'raw image token count does not match gen_shape: {token_ids.shape[-1]} != {height * width}') |
| grid = token_ids.reshape(token_ids.shape[0], height, width) |
| markers = torch.full((token_ids.shape[0], height, markers_per_row), int(row_marker_id), dtype=token_ids.dtype, device=token_ids.device) |
| markers[:, -1, :] = int(image_end_id) |
| augmented = torch.cat((grid, markers), dim=-1).flatten(1) |
| structural_mask = torch.zeros_like(augmented, dtype=torch.bool) |
| structural_mask.reshape(token_ids.shape[0], *augmented_shape)[:, :, -markers_per_row:] = True |
| return (augmented, augmented_shape, structural_mask) |
|
|
| def strip_image_row_markers(token_ids: torch.Tensor, original_gen_shape, *, markers_per_row: int=MARKERS_PER_RAW_ROW) -> torch.Tensor: |
| if not torch.is_tensor(token_ids) or token_ids.ndim != 2: |
| raise ValueError('token_ids must have shape [batch, augmented_tokens]') |
| markers_per_row = _validate_markers_per_row(markers_per_row) |
| (height, width) = (int(value) for value in original_gen_shape) |
| augmented_shape = augmented_image_shape((height, width), markers_per_row=markers_per_row) |
| if token_ids.shape[-1] != math.prod(augmented_shape): |
| raise ValueError(f'augmented image token count does not match marker grid: {token_ids.shape[-1]} != {math.prod(augmented_shape)}') |
| return token_ids.reshape(token_ids.shape[0], *augmented_shape)[:, :, :width].flatten(1) |
|
|
