File size: 5,910 Bytes
300c478 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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')
|