Spaces:
Running on Zero
Running on Zero
File size: 3,007 Bytes
7405e85 | 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 | 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
@torch.compiler.disable
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
|