|
|
| from __future__ import annotations |
| import numpy as np |
| import torch |
|
|
| def resolve_source_padding_vq_lock(model_config, override=None, *, has_edit_image): |
| if not isinstance(has_edit_image, bool): |
| raise TypeError('has_edit_image must be boolean') |
| if override is None: |
| resolved = getattr(model_config, 'lock_source_padding_vq', False) |
| if resolved is None: |
| resolved = False |
| if not isinstance(resolved, bool): |
| raise TypeError('config.lock_source_padding_vq must be boolean or null') |
| if resolved and (not has_edit_image): |
| return False |
| return resolved |
| if not isinstance(override, bool): |
| raise TypeError('lock_source_padding_vq override must be boolean or null') |
| if override and (not has_edit_image): |
| raise ValueError('lock_source_padding_vq requires an edit_image') |
| return override |
|
|
| def _resolve_source_native_size(value, *, fallback): |
| resolved = fallback if value is None else value |
| if not isinstance(resolved, (tuple, list)) or len(resolved) != 2 or any((isinstance(item, bool) or not isinstance(item, (int, np.integer)) or int(item) < 1 for item in resolved)): |
| raise ValueError('edit_source_native_size must contain positive integer width and height') |
| return (int(resolved[0]), int(resolved[1])) |
|
|
| def _source_padding_vq_mask(source_native_size, raw_grid_shape, *, device): |
| (width, height) = _resolve_source_native_size(source_native_size, fallback=source_native_size) |
| if not isinstance(raw_grid_shape, (tuple, list)) or len(raw_grid_shape) != 2 or any((isinstance(item, bool) or not isinstance(item, (int, np.integer)) or int(item) < 1 for item in raw_grid_shape)): |
| raise ValueError('raw_grid_shape must contain positive integer height and width') |
| (grid_height, grid_width) = (int(value) for value in raw_grid_shape) |
| if grid_height != grid_width: |
| raise ValueError('source-padding VQ locking requires a square raw VQ grid') |
| side = max(width, height) |
| left = round((side - width) * 0.5) |
| top = round((side - height) * 0.5) |
| right = left + width |
| bottom = top + height |
| columns = torch.arange(grid_width, dtype=torch.long, device=device) |
| rows = torch.arange(grid_height, dtype=torch.long, device=device) |
| horizontally_outside = ((columns + 1) * side <= left * grid_width) | (columns * side >= right * grid_width) |
| vertically_outside = ((rows + 1) * side <= top * grid_height) | (rows * side >= bottom * grid_height) |
| return (vertically_outside[:, None] | horizontally_outside[None, :]).reshape(1, -1) |
|
|
| def _augment_source_padding_vq_mask(state, raw_padding_mask, content_positions): |
| if not torch.is_tensor(state) or state.ndim != 2: |
| raise ValueError('generation state must have shape [batch, tokens]') |
| if not torch.is_tensor(raw_padding_mask) or raw_padding_mask.dtype is not torch.bool or raw_padding_mask.ndim != 2 or (raw_padding_mask.shape[0] != state.shape[0]) or (raw_padding_mask.device != state.device): |
| raise TypeError('raw_padding_mask must be device-local bool with the state batch size') |
| if not torch.is_tensor(content_positions) or content_positions.ndim != 1 or content_positions.dtype is not torch.long or (content_positions.device != state.device) or (content_positions.numel() != raw_padding_mask.shape[1]): |
| raise TypeError('content_positions must be a device-local long vector covering raw cells') |
| if bool((content_positions < 0).any()) or bool((content_positions >= state.shape[1]).any()): |
| raise ValueError('content_positions contains an out-of-range state index') |
| locked_mask = torch.zeros_like(state, dtype=torch.bool) |
| locked_mask[:, content_positions] = raw_padding_mask |
| return locked_mask |
|
|
| def _prefill_source_padding_vq(state, source_codes, raw_padding_mask, content_positions): |
| if not torch.is_tensor(state) or state.ndim != 2: |
| raise ValueError('generation state must have shape [batch, tokens]') |
| if not torch.is_tensor(source_codes) or source_codes.ndim != 2 or source_codes.dtype != state.dtype or (source_codes.device != state.device): |
| raise TypeError('source_codes must match the state integer dtype and device') |
| if source_codes.shape != raw_padding_mask.shape: |
| raise TypeError('raw_padding_mask must be bool and match source_codes') |
| locked_mask = _augment_source_padding_vq_mask(state, raw_padding_mask, content_positions) |
| locked_codes = torch.zeros_like(state) |
| locked_codes[:, content_positions] = source_codes |
| prefilled = state.clone() |
| prefilled[locked_mask] = locked_codes[locked_mask] |
| return (prefilled, locked_mask, locked_codes) |
|
|
| def _restore_source_padding_vq(state, locked_mask, locked_codes): |
| if locked_mask is None: |
| return state |
| if not torch.is_tensor(state) or state.ndim != 2 or (not torch.is_tensor(locked_mask)) or (locked_mask.dtype is not torch.bool) or (locked_mask.shape != state.shape) or (locked_mask.device != state.device) or (not torch.is_tensor(locked_codes)) or (locked_codes.shape != state.shape) or (locked_codes.dtype != state.dtype) or (locked_codes.device != state.device): |
| raise TypeError('locked source-padding tensors must align with sampler state') |
| restored = state.clone() |
| restored[locked_mask] = locked_codes[locked_mask] |
| return restored |
|
|
| def _assert_source_padding_vq_locked(state, locked_mask, locked_codes): |
| if locked_mask is None: |
| return |
| if not torch.is_tensor(locked_mask) or locked_mask.dtype is not torch.bool or locked_mask.shape != state.shape or (not torch.is_tensor(locked_codes)) or (locked_codes.shape != state.shape): |
| raise TypeError('locked source-padding tensors must align with sampler state') |
| if not torch.equal(state[locked_mask], locked_codes[locked_mask]): |
| raise RuntimeError('a locked source-padding VQ code changed during sampling') |
|
|
