| |
| """Микро-датасет для обучения MOSS-Image: латенты VAE + условие адаптера + маска. |
| |
| Формат строки (HF arrow dataset, `save_to_disk`): |
| vae float16 [128, H/16, W/16] — латент, уже упакованный и bn-нормированный, |
| ровно то, что ждёт MOSSImageTransformer |
| text_embeds float16 [L, 4096] — выход адаптерного ТЕ, без паддинга (как в пайплайне) |
| (только с --te) |
| text_mask bool [L] — маска условия (для нашей нарезки вся True; |
| нужна, когда батч склеивают паддингом) (только с --te) |
| text str — исходный капшен (для глаз и отладки) |
| width/height int — размер картинки, из которой сделан латент |
| aspect str — '1:1' / '3:4' / '4:3' |
| source str — путь исходного файла |
| |
| С `--te` рядом с датасетом кладётся `empty_condition.npz` (embeds float16 [L,4096] + mask bool [L]) — |
| условие пустого промпта, то есть негатив CFG. Он считается тем же адаптером один раз здесь, |
| чтобы обучению (`src/train.py`) текстовый энкодер был не нужен вообще. Без `--te` текстовый |
| энкодер не загружается вообще: в датасете только латенты и мета (текст, размеры, источник). |
| |
| Что делает с картинкой: |
| * EXIF-поворот применяется до всего остального; |
| * выбирается ближайшее по |log(w/h / target)| из набора соотношений (см. ASPECTS ниже); |
| * центр-кроп до выбранного соотношения, затем LANCZOS-ресайз до целевого размера. Апскейл |
| разрешён до --max-upscale (1.3 по умолчанию): мельче 1/1.3 от цели — отбрасывается. |
| |
| Латент: vae.encode(x).latent_dist.mode() -> [32,H/8,W/8]; упаковка 2x2 -> [128,H/16,W/16]; |
| нормировка (z - bn.running_mean) / sqrt(bn.running_var + eps). Обратная операция — в |
| `MOSSImagePipeline.__call__`, round-trip проверен (MSE 3e-4). |
| |
| Текст (только с --te): тот же путь, что `MOSSImagePipeline.encode_prompt` — chat-шаблон с |
| add_generation_prompt=True/enable_thinking=True, обрезка по --text-len, отдельный forward на |
| промпт без паддинга (паддинг меняет цель на ~0.5 %, см. adapter/src/cache_targets.py). |
| |
| Инфраструктура — как в sdxs-micro/dataset.py: accelerate делит файлы между GPU, каждый процесс |
| пишет чанки-поддатасеты <save>/chunk_*/sub_* с маркером .done (перезапуск продолжает с места). |
| Склейки в один датасет нет: src/train.py сам находит все чанки в папке и конкатенирует их при |
| чтении (лениво, arrow-файлы не переписываются). |
| |
| Запуск: |
| python3 src/make_dataset.py --dry-run --limit 20 # только отчёт по размерам, без GPU |
| python3 src/make_dataset.py --te --limit 20 # 20 картинок + условие ТЕ |
| python3 src/make_dataset.py --folder ~/dataset/all # только латенты (ТЕ не грузится) |
| accelerate launch --multi_gpu --num_processes 2 src/make_dataset.py --te --folder ~/dataset/allphoto_train |
| |
| Как читать в обучении: |
| row = ds[i] # np.array() даёт float64: приводим явно |
| lat = torch.from_numpy(np.asarray(row["vae"], dtype=np.float16))[None] # [1,128,H/16,W/16] |
| cond = torch.from_numpy(np.asarray(row["text_embeds"], dtype=np.float16))[None] # [1,L,4096] |
| mask = torch.from_numpy(np.asarray(row["text_mask"]))[None] # [1,L] bool |
| """ |
| import argparse |
| import gc |
| import math |
| import os |
| import random |
| import shutil |
| import sys |
| import time |
|
|
| import numpy as np |
| import torch |
| from PIL import Image, ImageOps |
| from accelerate import Accelerator |
| from datasets import Dataset, load_from_disk |
|
|
| MODEL = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| if MODEL not in sys.path: |
| sys.path.insert(0, MODEL) |
|
|
| accelerator = Accelerator() |
| device = accelerator.device |
| process_index = accelerator.process_index |
| num_processes = accelerator.num_processes |
| is_main = accelerator.is_main_process |
|
|
| |
| |
| |
| ASPECTS = { |
| "1:1": (1152, 1152), |
| "4:3": (1152, 864), |
| "3:4": (864, 1152), |
| "3:2": (1152, 768), |
| "2:3": (768, 1152), |
| "16:9": (1280, 720), |
| "9:16": (720, 1280), |
| } |
| DEFAULT_FOLDER = os.path.expanduser("~/dataset") |
| DEFAULT_SAVE = os.path.join(MODEL, "datasets", "dataset") |
| |
| |
| VAE_DTYPE = torch.float16 |
| |
| _MAX_ELEMS_PER_DATASET = int(1_500_000_000) |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser(description="MOSS-Image: картинки -> латенты + условие + маска") |
| p.add_argument("--folder", default=DEFAULT_FOLDER, help="папка с *.jpg/*.png и парными *.txt") |
| p.add_argument("--save", default=DEFAULT_SAVE, help="куда сохранить датасет") |
| p.add_argument("--limit", type=int, default=0, |
| help="взять только N случайных пар (сид --seed); 0 — все") |
| p.add_argument("--seed", type=int, default=0, help="сид случайного отбора для --limit") |
| p.add_argument("--batch-size", type=int, default=4, help="картинок на один forward VAE") |
| p.add_argument("--chunk-size", type=int, default=10000, help="файлов на чанк (на процесс)") |
| p.add_argument("--text-len", type=int, default=1024, help="макс. токенов условия (как в пайплайне)") |
| p.add_argument("--max-upscale", type=float, default=1.3, |
| help="во сколько раз разрешено увеличивать картинку до целевой; 1.0 — только " |
| "даунскейл (всё мельче целевого размера отбрасывается)") |
| p.add_argument("--strip-words", default="*tags*", |
| help="через запятую: подстроки, вырезаемые из капшена (по умолчанию '*tags*', " |
| "то есть и '*tags* , 1girl', и '*tags* 1girl'); " |
| "пустая строка — ничего не вырезать. После вырезки капшен сжимается по " |
| "пробелам и strip-ится по ' ,;|'; пустой — строка отбрасывается") |
| p.add_argument("--dry-run", action="store_true", |
| help="только отчёт: что выбрано, сколько отброшено; без GPU и без записи") |
| p.add_argument("--te", action="store_true", |
| help="считать условие адаптерным ТЕ (text_embeds/text_mask + empty_condition.npz); " |
| "без флага текстовый энкодер не грузится, в датасете только латенты и мета") |
| return p.parse_args() |
|
|
|
|
| |
| def load_vae(): |
| from diffusers import AutoencoderKLFlux2 |
| vae = AutoencoderKLFlux2.from_pretrained(os.path.join(MODEL, "vae"), torch_dtype=VAE_DTYPE) |
| return vae.to(device).eval() |
|
|
|
|
| def load_text_encoder(): |
| from transformers import AutoTokenizer |
| from pipeline_moss_image import MOSSAdapterTextEncoder |
| path = os.path.join(MODEL, "text_encoder") |
| tok = AutoTokenizer.from_pretrained(path) |
| |
| te = MOSSAdapterTextEncoder.from_pretrained(path).to(device).eval() |
| return te, tok |
|
|
|
|
| def render(tok, text): |
| """Ровно то, что делает MOSSImagePipeline.encode_prompt (сырая строка, не id).""" |
| return tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=True) |
|
|
|
|
| @torch.no_grad() |
| def encode_text(te, tok, text, text_len): |
| """Промпт -> (embeds float16 [L,4096], mask bool [L]). Без паддинга, как в пайплайне.""" |
| ids = tok(render(tok, text), truncation=True, max_length=text_len)["input_ids"] |
| if not ids: |
| raise ValueError("пустой промпт после токенизации") |
| ids = torch.tensor([ids], dtype=torch.long, device=device) |
| out = te(input_ids=ids, attention_mask=torch.ones_like(ids, dtype=torch.bool), use_cache=False) |
| emb = out.last_hidden_state[0].to(torch.float16) |
| if not torch.isfinite(emb).all(): |
| raise RuntimeError(f"условие не влезло в float16 (промпт: {text[:60]!r})") |
| mask = np.ones(emb.shape[0], dtype=bool) |
| return emb.cpu().numpy(), mask |
|
|
|
|
| @torch.no_grad() |
| def encode_latents(vae, images): |
| """list[PIL] (одинакового размера) -> np.float16 [B,128,H/16,W/16].""" |
| x = torch.stack([torch.from_numpy(np.array(im)).permute(2, 0, 1).float().div(127.5).sub(1) |
| for im in images]).to(device, VAE_DTYPE) |
| lat = vae.encode(x).latent_dist.mode() |
| b, c, h, w = lat.shape |
| if h % 2 or w % 2: |
| raise ValueError(f"латент {h}x{w} не делится на 2 — картинка не кратна 16?") |
| |
| z = lat.reshape(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape( |
| b, c * 4, h // 2, w // 2) |
| bn = vae.bn |
| mean = bn.running_mean.to(z.dtype).view(1, -1, 1, 1) |
| std = (bn.running_var + vae.config.batch_norm_eps).sqrt().to(z.dtype).view(1, -1, 1, 1) |
| return ((z - mean) / std).half().cpu().numpy() |
|
|
|
|
| |
| def plan_size(width, height, max_upscale=1.0): |
| """(w,h) исходника -> (label, tw, th, cw, ch) кропа или None, если картинка слишком мала.""" |
| ratio = width / height |
| label, (tw, th) = min(ASPECTS.items(), |
| key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1])))) |
| if ratio > tw / th: |
| cw, ch = round(height * tw / th), height |
| else: |
| cw, ch = width, round(width * th / tw) |
| cw, ch = min(cw, width), min(ch, height) |
| if cw * max_upscale < tw or ch * max_upscale < th: |
| return None |
| return label, tw, th, cw, ch |
|
|
|
|
| def header_size(path): |
| """(w,h) с учётом EXIF-поворота, без декодирования пикселей.""" |
| with Image.open(path) as im: |
| w, h = im.size |
| if im.getexif().get(0x0112, 1) in (5, 6, 7, 8): |
| w, h = h, w |
| return w, h |
|
|
|
|
| def load_image(path, plan): |
| label, tw, th, cw, ch = plan |
| im = ImageOps.exif_transpose(Image.open(path)).convert("RGB") |
| left, top = (im.width - cw) // 2, (im.height - ch) // 2 |
| return im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS) |
|
|
|
|
| |
| def process_folder(folder, max_upscale=1.0): |
| """Пары картинка+капшен, детерминированный порядок (для resume и шардинга).""" |
| rows, skipped = [], {"unpaired": 0, "broken": 0, "small": 0} |
| seen = 0 |
| for root, dirs, files in os.walk(folder): |
| dirs.sort() |
| for fn in sorted(files): |
| if not fn.lower().endswith((".jpg", ".jpeg", ".png", ".webp")): |
| continue |
| seen += 1 |
| if seen % 2000 == 0: |
| print(f"process_folder: просмотрено {seen} картинок, пар {len(rows)}", flush=True) |
| path = os.path.join(root, fn) |
| txt = os.path.splitext(path)[0] + ".txt" |
| if not os.path.exists(txt): |
| skipped["unpaired"] += 1 |
| continue |
| try: |
| w, h = header_size(path) |
| except Exception: |
| skipped["broken"] += 1 |
| continue |
| plan = plan_size(w, h, max_upscale) |
| if plan is None: |
| skipped["small"] += 1 |
| continue |
| rows.append((path, txt, plan)) |
| print(f"process_folder: {len(rows)} пар, отброшено {skipped['unpaired']} без капшена + " |
| f"{skipped['broken']} битых + {skipped['small']} мелких", flush=True) |
| return rows |
|
|
|
|
| def pick_limit(rows, limit, seed): |
| if limit <= 0 or limit >= len(rows): |
| return rows |
| chosen = sorted(random.Random(seed).sample(range(len(rows)), limit)) |
| return [rows[i] for i in chosen] |
|
|
|
|
| def read_caption(path, strip_words=()): |
| try: |
| text = open(path, encoding="utf-8", errors="replace").read() |
| except Exception: |
| return None |
| for word in strip_words: |
| text = text.replace(word, " ") |
| text = " ".join(text.split()) |
| text = text.strip(" ,;|\t") |
| return text or None |
|
|
|
|
| def clear_cuda_memory(): |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| gc.collect() |
|
|
|
|
| def dry_run(rows): |
| print(f"\n{'aspect':6} {'crop':>12} {'target':>11} файл") |
| for path, _, (label, tw, th, cw, ch) in rows[:20]: |
| print(f"{label:6} {cw:5}x{ch:<6} {tw:5}x{th:<5} {os.path.basename(path)}") |
| hist = {} |
| for _, _, (label, *_rest) in rows: |
| hist[label] = hist.get(label, 0) + 1 |
| print("\nвсего", len(rows), "по соотношениям:", dict(sorted(hist.items())), |
| f"(из них показано {min(20, len(rows))})") |
|
|
|
|
| |
| def process_group(vae, te, tok, args, rows, group_dir, lat_shape, strip_words): |
| """Одна size-группа (в батч можно класть только одинаковые кадры) -> под-датасеты. |
| |
| С `--te` в строку добавляются text_embeds/text_mask; без него эти ключи не создаются вовсе. |
| """ |
| keys = ["vae", "text", "width", "height", "aspect", "source"] |
| if args.te: |
| keys += ["text_embeds", "text_mask"] |
| marker = os.path.join(group_dir, ".done") |
| if os.path.exists(marker): |
| print(f"[GPU {process_index}] пропуск готовой группы {os.path.basename(group_dir)}", |
| flush=True) |
| return 0 |
| if os.path.isdir(group_dir): |
| shutil.rmtree(group_dir) |
| os.makedirs(group_dir, exist_ok=True) |
|
|
| |
| |
| elems = lat_shape[0] * lat_shape[1] * lat_shape[2] |
| if args.te: |
| elems += 512 * 4096 |
| save_batch = max(64, _MAX_ELEMS_PER_DATASET // max(elems, 1)) |
|
|
| acc = {k: [] for k in keys} |
| saved = sub_n = skipped = 0 |
|
|
| def flush(): |
| nonlocal saved, sub_n |
| if not acc["vae"]: |
| return |
| ds = Dataset.from_dict(dict(acc)) |
| path = os.path.join(group_dir, f"sub_{sub_n:05d}") |
| ds.save_to_disk(path) |
| print(f"[GPU {process_index}] сохранено {len(acc['vae'])} -> {path}", flush=True) |
| saved += len(acc["vae"]) |
| sub_n += 1 |
| for v in acc.values(): |
| v.clear() |
|
|
| for start in range(0, len(rows), args.batch_size): |
| batch = [] |
| for path, txt_path, plan in rows[start:start + args.batch_size]: |
| try: |
| img = load_image(path, plan) |
| except Exception as e: |
| print(f"[GPU {process_index}] битая картинка {path}: {e}", flush=True) |
| skipped += 1 |
| continue |
| text = read_caption(txt_path, strip_words) |
| if text is None: |
| print(f"[GPU {process_index}] пустой/нечитаемый капшен {txt_path}", flush=True) |
| skipped += 1 |
| continue |
| batch.append((img, text, plan, path)) |
| if not batch: |
| continue |
|
|
| lat = encode_latents(vae, [im for im, _, _, _ in batch]) |
| embeds = None |
| if args.te: |
| embeds = [encode_text(te, tok, text, args.text_len) |
| for _, text, _, _ in batch] |
| for i, (row_lat, (_, text, (label, tw, th, _, _), path)) in enumerate(zip(lat, batch)): |
| acc["vae"].append(row_lat) |
| if embeds is not None: |
| emb, mask = embeds[i] |
| acc["text_embeds"].append(emb) |
| acc["text_mask"].append(mask) |
| acc["text"].append(text) |
| acc["width"].append(tw) |
| acc["height"].append(th) |
| acc["aspect"].append(label) |
| acc["source"].append(path) |
| del lat, batch, embeds |
| if len(acc["vae"]) >= save_batch: |
| flush() |
|
|
| flush() |
| open(marker, "w").close() |
| print(f"[GPU {process_index}] группа {group_dir.rsplit('_', 1)[-1]}: " |
| f"{saved} сэмплов (skip {skipped})", flush=True) |
| return saved |
|
|
|
|
| def process_in_chunks(vae, te, tok, args, rows, root, strip_words): |
| total = len(rows) |
| t0 = time.time() |
| if is_main and args.te: |
| |
| emb, mask = encode_text(te, tok, "", args.text_len) |
| np.savez(os.path.join(root, "empty_condition.npz"), embeds=emb, mask=mask) |
| print(f"пустой промпт: {emb.shape} -> {root}/empty_condition.npz", flush=True) |
| saved = 0 |
| for chunk_idx, start in enumerate(range(0, total, args.chunk_size), 1): |
| end = min(start + args.chunk_size, total) |
| groups = {} |
| for path, txt_path, plan in rows[start:end]: |
| key = (plan[1], plan[2]) |
| groups.setdefault(key, []).append((path, txt_path, plan)) |
| for (tw, th), group_rows in sorted(groups.items()): |
| lat_shape = (128, th // 16, tw // 16) |
| group_dir = os.path.join( |
| root, f"chunk_{chunk_idx:05d}_{tw}x{th}_proc_{process_index}") |
| saved += process_group(vae, te, tok, args, group_rows, group_dir, lat_shape, |
| strip_words) |
| clear_cuda_memory() |
| print(f"[GPU {process_index}] чанк {chunk_idx} готов ({end - start} файлов, " |
| f"{time.time() - t0:.0f} c всего)", flush=True) |
| return saved |
|
|
|
|
| |
| def count_chunks(root): |
| """Сколько под-датасетов и строк уже на месте (то, что прочтёт src/train.py).""" |
| n_ds = n_rows = 0 |
| for name in sorted(os.listdir(root)): |
| group_dir = os.path.join(root, name) |
| if not (os.path.isdir(group_dir) and name.startswith("chunk_")): |
| continue |
| if not os.path.exists(os.path.join(group_dir, ".done")): |
| print(f"[main] ВНИМАНИЕ: {name} без маркера .done — обучение увидит его " |
| f"неполным; запустите скрипт ещё раз, чтобы досчитать", flush=True) |
| for sub in sorted(os.listdir(group_dir)): |
| if sub.startswith("sub_"): |
| n_ds += 1 |
| n_rows += len(load_from_disk(os.path.join(group_dir, sub))) |
| return n_ds, n_rows |
|
|
|
|
| def main(): |
| args = parse_args() |
| strip_words = [w for w in (x.strip() for x in args.strip_words.split(",")) if w] |
| rows = pick_limit(process_folder(args.folder, args.max_upscale), args.limit, args.seed) |
| by_aspect = {} |
| for _, _, (label, *_rest) in rows: |
| by_aspect[label] = by_aspect.get(label, 0) + 1 |
| if is_main: |
| print(f"к обработке: {len(rows)}, по соотношениям: {dict(sorted(by_aspect.items()))}", |
| flush=True) |
| if not rows: |
| raise SystemExit("нечего кодировать: ни одной подходящей пары картинка+капшен") |
| if args.dry_run: |
| if is_main: |
| dry_run(rows) |
| return |
|
|
| |
| if not torch.cuda.is_available(): |
| raise SystemExit("CUDA недоступна (torch.cuda.is_available() == False)") |
| world_size = int(os.environ.get("WORLD_SIZE", num_processes)) |
| if world_size != num_processes: |
| raise SystemExit(f"accelerate не распределил процессы: WORLD_SIZE={world_size}, " |
| f"num_processes={num_processes}") |
|
|
| if is_main: |
| print(f"processes: {num_processes}, folder: {args.folder}, save: {args.save}", flush=True) |
|
|
| |
| rows.sort(key=lambda r: (r[2][1], r[2][2], r[0])) |
| mine = list(range(process_index, len(rows), num_processes)) |
| rows = [rows[i] for i in mine] |
| print(f"[GPU {process_index}] обрабатываю {len(rows)} файлов", flush=True) |
|
|
| os.makedirs(args.save, exist_ok=True) |
|
|
| vae = load_vae() |
| te, tok = (load_text_encoder() if args.te else (None, None)) |
| print(f"[GPU {process_index}] VAE{'' if args.te else ' (без ТЕ: --te не указан)'} на " |
| f"{device}, VRAM {torch.cuda.memory_allocated() / 2**30:.1f} ГиБ", flush=True) |
| t0 = time.time() |
| try: |
| process_in_chunks(vae, te, tok, args, rows, args.save, strip_words) |
| finally: |
| del vae, te, tok |
| clear_cuda_memory() |
| accelerator.wait_for_everyone() |
|
|
| if is_main: |
| n_ds, n_rows = count_chunks(args.save) |
| print(f"[main] готово за {time.time() - t0:.0f} c: {n_rows} сэмплов в {n_ds} чанках " |
| f"-> {args.save} (склейки нет, src/train.py читает чанки сам)", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
