| """Поток картинок из TFRecord-шардов i1 (PNG 256×256 + 5 подписей) вместо parquet-потока FLUX. |
| |
| Отличие от HFImageStream: запись TFRecord = [uint64 len][crc][payload][crc], внутри PNG-байты |
| и `caption` — пять длинных перефразировок одной картинки. Поэтому: |
| * из пяти подписей берём случайно самую короткую или следующую по длине (бюджета на все нет, |
| а короткая лучше влезает в --text-len); |
| * картинки уже 256×256 — ресайз только если сторона не совпала с --size; |
| * источник: локальная папка с шардами (быстро) или прямо репозиторий HF через range-запросы |
| (block_size 16 МБ даёт ~90 записей/с, ~9.7 МБ/с — на уровне нашего темпа обучения). |
| """ |
| from __future__ import annotations |
|
|
| import io |
| import queue |
| import random |
| import struct |
| import threading |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| SHARD_GLOB = "*.tfrecord-*" |
|
|
|
|
| |
| def _varint(buf: bytes, i: int) -> tuple[int, int]: |
| shift = val = 0 |
| while True: |
| b = buf[i]; i += 1 |
| val |= (b & 0x7F) << shift |
| if not b & 0x80: |
| return val, i |
| shift += 7 |
|
|
|
|
| def _fields(buf: bytes): |
| i = 0 |
| while i < len(buf): |
| key, i = _varint(buf, i) |
| num, wire = key >> 3, key & 7 |
| if wire == 2: |
| ln, i = _varint(buf, i) |
| yield num, buf[i:i + ln] |
| i += ln |
| elif wire == 0: |
| v, i = _varint(buf, i) |
| yield num, v |
| elif wire == 1: |
| yield num, buf[i:i + 8]; i += 8 |
| elif wire == 5: |
| yield num, buf[i:i + 4]; i += 4 |
| else: |
| raise ValueError(f"wire type {wire}") |
|
|
|
|
| def _packed_i64(buf: bytes) -> list[int]: |
| out, i = [], 0 |
| while i < len(buf): |
| v, i = _varint(buf, i) |
| out.append(v) |
| return out |
|
|
|
|
| def parse_example(payload: bytes) -> dict: |
| """tf.train.Example → {поле: список bytes | список int}""" |
| out = {} |
| for num, features in _fields(payload): |
| if num != 1: |
| continue |
| for fn, entry in _fields(features): |
| if fn != 1: |
| continue |
| key, val = None, {} |
| for en, ev in _fields(entry): |
| if en == 1: |
| key = ev.decode() |
| elif en == 2: |
| for fnum, fval in _fields(ev): |
| if fnum == 1: |
| val["bytes"] = [v for n, v in _fields(fval) if n == 1] |
| elif fnum == 3: |
| val["int64"] = [x for n, v in _fields(fval) if n == 1 for x in _packed_i64(v)] |
| out[key] = val.get("bytes", val.get("int64")) |
| return out |
|
|
|
|
| def count_records(paths) -> int: |
| """Сколько записей в шардах — читаем только заголовки (seek через payload), быстро.""" |
| total = 0 |
| for path in paths: |
| with open(path, "rb") as fh: |
| size = fh.seek(0, 2) |
| pos = 0 |
| while pos + 12 <= size: |
| fh.seek(pos) |
| head = fh.read(8) |
| if len(head) < 8: |
| break |
| (n,) = struct.unpack("<Q", head) |
| pos += 8 + 4 + n + 4 |
| total += 1 |
| return total |
|
|
|
|
| def count_shard_records(source: str) -> int | None: |
| """Точное число записей для локальной папки; None, если источник удалённый.""" |
| if not Path(source).is_dir(): |
| return None |
| paths = sorted(p.as_posix() for p in Path(source).glob(SHARD_GLOB)) |
| return count_records(paths) if paths else None |
|
|
|
|
| def read_records(fh): |
| """файловый объект (локальный или HfFileSystem) → поток payload'ов""" |
| while True: |
| head = fh.read(8) |
| if len(head) < 8: |
| return |
| (n,) = struct.unpack("<Q", head) |
| fh.read(4) |
| payload = fh.read(n) |
| fh.read(4) |
| yield payload |
|
|
|
|
| |
| class TFRecordStream: |
| def __init__(self, source: str, size: int = 256, workers: int = 8, prefetch: int = 256, |
| block_mb: int = 16, seed: int = 0, caption_short2: bool = True): |
| """source — локальная папка с шардами или id датасета на HF (repo_type=dataset).""" |
| self.q: queue.Queue = queue.Queue(maxsize=prefetch) |
| self.stop = False |
| self.size, self.workers, self.seed = size, workers, seed |
| self.caption_short2 = caption_short2 |
| self.local = Path(source).is_dir() |
| self.source = str(source) |
| self.block = block_mb * 2 ** 20 |
| threading.Thread(target=self._loop, daemon=True).start() |
|
|
| |
| def _shards(self) -> list[str]: |
| if self.local: |
| files = sorted(p.as_posix() for p in Path(self.source).glob(SHARD_GLOB)) |
| else: |
| from huggingface_hub import HfFileSystem |
| fs = HfFileSystem() |
| files = [f"datasets/{self.source}/{p}" for p in sorted(fs.ls(f"datasets/{self.source}")) |
| if ".tfrecord-" in p] |
| if not files: |
| raise FileNotFoundError(f"нет шардов в {self.source}") |
| random.Random(self.seed).shuffle(files) |
| return files |
|
|
| def _open(self, path: str): |
| if self.local: |
| return open(path, "rb") |
| from huggingface_hub import HfFileSystem |
| return HfFileSystem().open(path, "rb", block_size=self.block) |
|
|
| |
| def _prepare(self, payload: bytes): |
| d = parse_example(payload) |
| im = Image.open(io.BytesIO(d["image"][0])).convert("RGB") |
| w, h = im.size |
| if min(w, h) != self.size: |
| sc = self.size / min(w, h) |
| im = im.resize((max(self.size, round(w * sc)), max(self.size, round(h * sc))), Image.LANCZOS) |
| w, h = im.size |
| im = im.crop(((w - self.size) // 2, (h - self.size) // 2, |
| (w - self.size) // 2 + self.size, (h - self.size) // 2 + self.size)) |
| caps = [c.decode() if isinstance(c, bytes) else str(c) for c in d.get("caption", [])] |
| caps = sorted(caps, key=len) |
| if not caps: |
| caps = [""] |
| cap = random.choice(caps[:2]) if self.caption_short2 else random.choice(caps) |
| arr = torch.from_numpy(np.asarray(im, dtype=np.float32)).permute(2, 0, 1) / 127.5 - 1.0 |
| return arr, cap |
|
|
| def _loop(self): |
| import collections |
| from concurrent.futures import ThreadPoolExecutor |
| with ThreadPoolExecutor(max_workers=self.workers) as pool: |
| while not self.stop: |
| try: |
| for shard in self._shards(): |
| fh = self._open(shard) |
| try: |
| it = read_records(fh) |
| flying = collections.deque() |
| while not self.stop: |
| while len(flying) < self.workers: |
| flying.append(pool.submit(self._prepare, next(it))) |
| self.q.put(flying.popleft().result()) |
| finally: |
| fh.close() |
| except Exception as e: |
| print(f"[tfrecord] сбой ({type(e).__name__}: {e}), переподключение через 5 с", flush=True) |
| import time |
| time.sleep(5) |
|
|
| def batch(self, n: int): |
| imgs, caps = [], [] |
| while len(imgs) < n: |
| arr, cap = self.q.get() |
| imgs.append(arr); caps.append(cap) |
| return torch.stack(imgs), caps |
|
|
