Spaces:
Running on Zero
Running on Zero
| import io | |
| import math | |
| from typing import List, Optional | |
| import torch | |
| import torch.nn.functional as F | |
| def concat_bytes_list(bytes_list: List[bytes], bs_io: io.BytesIO = None) -> Optional[bytes]: | |
| """ | |
| assert 0 <= bytes_list[i] <= 16777216 and len(bytes_list) > 1 | |
| """ | |
| assert len(bytes_list) > 1 | |
| head_bits_num = 1 | |
| bytes_len_bytes_list = [] | |
| bytes_len_bytes_len_list = [] | |
| for _ in bytes_list: | |
| bytes_len = len(_) | |
| bytes_len_bytes_len = math.ceil(bytes_len.bit_length() / 8) or 1 | |
| head_bits_num = max(head_bits_num, (bytes_len_bytes_len - 1).bit_length()) | |
| bytes_len_bytes_list.append(bytes_len.to_bytes(bytes_len_bytes_len, 'little', signed=False)) | |
| bytes_len_bytes_len_list.append(bytes_len_bytes_len) | |
| assert head_bits_num == 1 or head_bits_num == 2, head_bits_num | |
| head_bytes = int( | |
| '1' + ''.join((format(_ - 1, f'0{head_bits_num}b') for _ in bytes_len_bytes_len_list)), 2 | |
| ).to_bytes(math.ceil(len(bytes_list) / (8 // head_bits_num) + 0.25), 'big', signed=False) | |
| if head_bits_num == 2: | |
| head_bytes = bytes([head_bytes[0] | 0x80]) + head_bytes[1:] | |
| if bs_io is None: | |
| return_bytes = True | |
| bs_io = io.BytesIO() | |
| else: | |
| return_bytes = False | |
| bs_io.write(head_bytes) | |
| for _ in bytes_len_bytes_list: | |
| bs_io.write(_) | |
| for _ in bytes_list: | |
| bs_io.write(_) | |
| if return_bytes: | |
| ret = bs_io.getvalue() | |
| bs_io.close() | |
| return ret | |
| def split_bytes_list(concat_bytes: Optional[bytes], bytes_list_len: int, | |
| bs_io: io.BytesIO = None) -> List[bytes]: | |
| if bs_io is None: | |
| bytes_given = True | |
| else: | |
| bytes_given = False | |
| assert concat_bytes is None | |
| if bytes_given: | |
| bs_io = io.BytesIO(concat_bytes) | |
| first_byte = bs_io.read(1)[0] | |
| head_bits_num = 2 if bool(first_byte & 0x80) else 1 | |
| head_bytes_len = math.ceil(bytes_list_len / (8 // head_bits_num) + 0.25) | |
| head_bits = f"{int.from_bytes(bytes([first_byte & 0x7f]) + bs_io.read(head_bytes_len - 1), 'big'):b}"[1:] | |
| bytes_len_bytes_len_list = [] | |
| for idx in range(0, bytes_list_len * head_bits_num, head_bits_num): | |
| bytes_len_bytes_len_list.append(int(head_bits[idx: idx + head_bits_num], 2) + 1) | |
| bytes_len_list = [] | |
| for bytes_len_bytes_len in bytes_len_bytes_len_list: | |
| bytes_len_list.append(int.from_bytes(bs_io.read(bytes_len_bytes_len), 'little')) | |
| bytes_list = [] | |
| for bytes_len in bytes_len_list: | |
| bytes_list.append(bs_io.read(bytes_len)) | |
| if bytes_given: | |
| bs_io.close() | |
| return bytes_list | |
| def pad_tensor_divisible_by(img, div=64): | |
| h_old, w_old = img.shape[2:] | |
| if (h_old % div == 0) and (w_old % div == 0): | |
| return img | |
| h_tgt = div * math.ceil(h_old / div) | |
| w_tgt = div * math.ceil(w_old / div) | |
| # left, top, right, bottom | |
| padded = F.pad(img, (0, w_tgt - w_old, 0, h_tgt - h_old), 'replicate') | |
| return padded | |