Ra fp16 (step 460000)

src/export.py

@@ -37,8 +37,7 @@ def main():
     save_file(tensors, str(out / "diffusion_pytorch_model.safetensors"))
 
     ae_cfg = json.loads((ROOT / "vae" / "config.json").read_text())
-    ta = ck.get("args") or {}          # чекпойнт помнит свои аргументы — не сочиняем заново
-    graft = dict((k.strip(), int(v)) for k, v in (kv.split("=") for kv in (ta.get("graft") or "").split(",") if "=" in kv))
+    ta = ck.get("args") or {}
     counts = {}                          # сколько блоков в каждом списке — по ключам тензоров
     for k in ck["body"]:
         g = k.split(".")[0]
@@ -47,76 +46,40 @@ def main():
             if idx.isdigit():
                 counts.setdefault(g, set()).add(idx)
     nblk = {g: len(v) for g, v in counts.items()}
-    training = {
-        "dataset": ta.get("hf_dataset") or ta.get("data") or "~/dataset/butterflies_256x64",
-        "parts": ta.get("parts"),
-        "size": ta.get("size"),
-        "steps": int(ck["step"]),
-        "batch": ta.get("batch", 8),
-        "accum": ta.get("accum", 1),
-        "lr": ta.get("lr", 1e-4),
-        "lr_tail": ta.get("lr_tail", 3e-5),
-        "drop_text": ta.get("drop_text", 0.1),
-        "init": ta.get("init"),
-        "metrics": args.metrics,
-    }
+    # конфиг описывает ТОЛЬКО тело; статистики латента — в vae/config.json, шаг — параметр запуска
     cfg = {
         "_class_name": "UDiT",
-        "name": "Ra",
         "format": "udm/v1",
-        "step": int(ck["step"]),
-        "latent": {
-            "channels": ae_cfg["latent_channels"],
-            "spatial_factor": ae_cfg["patch_size"],
-            "mean": ae_cfg["latent_mean"],
-            "std": ae_cfg["latent_std"],
-            "source": "vae/ (Muinez/f8c32-ae)",
-        },
-        "text": {
-            "encoder": "text_encoder/ (jhu-clsp/ettin-encoder-150m, 768-dim, 22 слоя)",
-            "max_len": 192,
-            "slices": [2, 5, 8, 11, 14, 17, 20],
-            "slice_blocks": 2,
-            "trainable": "text_mixer.*",
-        },
+        "storage_dtype": "float16" if args.half else "float32",
+        "latent_channels": ae_cfg["latent_channels"],
+        "text_dim": ta.get("text_dim", 768),
         "body": {
             "kind": "пирамидальный U: ширина растёт к низу",
-            "graft": graft or None,      # чем привито сверх базовой архитектуры (zero-init блоки)
+            "widths": [192, 512, 1024, 1536],
+            "heads": [6, 8, 8, 8],
             "blocks": nblk,
             "levels": [
-                {"level": 0, "tokens": "÷1", "width": 192, "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону",
-                 "text": "кросс-внимание в каждом conv-блоке и в self-attn"},
-                {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 5 enc / 4 dec", "text": "кросс-внимание"},
-                {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 5 enc / 4 dec", "text": "кросс-внимание"},
-                {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 4 (боттлнек)", "text": "кросс-внимание"},
+                {"level": 0, "tokens": "÷1", "width": 192,
+                 "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону"},
+                {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 5 enc / 4 dec"},
+                {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 5 enc / 4 dec"},
+                {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 4 (боттлнек)"},
             ],
             "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
-            "down": "mean-pool 2× + Linear",
-            "up": "Linear + nearest 2×",
-            "rope": "2D RoPE, base 100, фаза в латентных пикселях на каждом уровне",
+            "down": "mean-pool 2x + Linear",
+            "up": "Linear + nearest 2x",
+            "rope": {"kind": "2d", "base": 100, "phase": "в латентных пикселях на каждом уровне"},
+            "text": "кросс-внимание на всех уровнях",
             "parametrization": "x0 латента",
+            "params": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2),
         },
-        "tail": {
-            "source": "декодер AE, 6 блоков @512, PixelShuffle ×3",
-            "trained_as": "flow по пикселям, условие z0 + (ошибка тела).detach()",
-            "sample_steps": ta.get("tail_steps", 4),
-        },
-        "params": {
-            "body": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2),
-            "text_mixer": round(sum(v.numel() for k, v in tensors.items()
-                                    if k.startswith("text_mixer.")) / 1e6, 2),
-            "tail": round(sum(v.numel() for k, v in tensors.items() if k.startswith("tail.")) / 1e6, 2),
-        },
-        "training": {k: v for k, v in training.items() if v is not None},
-        "storage_dtype": "float16" if args.half else "float32",
     }
-    cfg["params"]["total"] = round(sum(cfg["params"].values()), 2)
     (out / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False))
 
     print(f"сохранено: {out}/diffusion_pytorch_model.safetensors "
           f"({sum(v.numel() for v in tensors.values())/1e6:.1f}M параметров, "
           f"{(out / 'diffusion_pytorch_model.safetensors').stat().st_size/1e6:.0f} МБ)")
-    print(f"части: {cfg['params']}")
+    print(f"тело {cfg['body']['params']}M, блоки {nblk}")
 
 
 def load_exported(model, enc, ae, path, device=None):

src/publish.py

@@ -50,6 +50,7 @@ def main():
     ap.add_argument("--out", default=str(ROOT / "tmp" / "publish"))
     ap.add_argument("--repo", default=REPO_ID)
     ap.add_argument("--push", action="store_true", help="залить собранное на Hub (публичный репозиторий)")
+    ap.add_argument("--step", type=int, default=None, help="шаг чекпойнта — только для сообщения коммита")
     args = ap.parse_args()
     out = Path(args.out)
     if out.exists():
@@ -90,7 +91,7 @@ def main():
         api = HfApi()
         api.create_repo(args.repo, repo_type="model", private=False, exist_ok=True)
         api.upload_folder(folder_path=str(out), repo_id=args.repo, repo_type="model",
-                          commit_message=f"Ra fp16 (step {cfg['step']})")
+                          commit_message=(f"Ra fp16 (step {args.step})" if args.step else "Ra fp16 (weights)"))
         print(f"[push] https://huggingface.co/{args.repo}", flush=True)
 
 

src/sample.py

@@ -51,7 +51,7 @@ def main():
             ae.decoder.load_state_dict(sd["decoder"])
         step = sd.get("step")
     m.eval()
-    print(f"[model] {src}, шаг {step}", flush=True)
+    print(f"[model] {src}" + (f", шаг {step}" if step else ""), flush=True)
 
     null_hs, null_mk = enc.slices([""], pad_to_len=args.text_len)
     null_t = enc.mix(null_hs.float())

src/train.py

@@ -174,6 +174,8 @@ def main():
     ap.add_argument("--steps", type=int, default=None, help="бюджет в шагах (по умолчанию — до --target)")
     ap.add_argument("--target", type=int, default=708000,
                     help="абсолютный шаг финиша, если не заданы --steps/--epochs")
+    ap.add_argument("--step", type=int, default=0,
+                    help="стартовый шаг: от него считается остаток и расписание lr")
     ap.add_argument("--batch", type=int, default=16)
     ap.add_argument("--accum", type=int, default=1, help="накопление: эффективный батч = batch*accum")
     ap.add_argument("--probe", action="store_true", help="подобрать максимальный батч под память и выйти")
@@ -273,9 +275,9 @@ def main():
         p = Path(args.init)
         if p.is_dir():
             from export import load_exported
-            cfg = load_exported(m, enc, ae, p, DEV)
-            start = int(cfg.get("step", 0) or 0)     # счётчик продолжается с шага экспорта
-            print(f"[init] папка {p}, шаг {start}", flush=True)
+            load_exported(m, enc, ae, p, DEV)
+            start = args.step                    # шаг — параметр запуска, не поле конфига
+            print(f"[init] папка {p}; старт шага {start} (--step)", flush=True)
         else:
             sd = torch.load(p, map_location=DEV)
             n_take = load_body(sd["body"]); enc.mixer.load_state_dict(sd["mixer"])

transformer/config.json

@@ -1,107 +1,23 @@
 {
   "_class_name": "UDiT",
-  "name": "Ra",
   "format": "udm/v1",
-  "step": 460000,
-  "latent": {
-    "channels": 32,
-    "spatial_factor": 8,
-    "mean": [
-      -4.137261,
-      2.725146,
-      -1.435803,
-      -2.187023,
-      -4.161275,
-      -1.58918,
-      -2.527523,
-      3.002302,
-      2.793705,
-      -2.091501,
-      -2.632896,
-      -3.59589,
-      -3.069669,
-      -4.568713,
-      3.26668,
-      -0.897648,
-      2.072541,
-      5.451158,
-      1.747871,
-      -2.113284,
-      2.03444,
-      -1.379274,
-      0.66369,
-      2.712183,
-      -0.111181,
-      6.054838,
-      1.235881,
-      -2.02777,
-      3.663243,
-      -2.454632,
-      1.141335,
-      1.786136
-    ],
-    "std": [
-      1.030738,
-      0.951682,
-      1.079687,
-      0.957871,
-      1.198266,
-      1.102266,
-      1.213605,
-      0.947544,
-      1.085991,
-      1.093223,
-      0.9027,
-      1.010205,
-      0.920466,
-      1.044107,
-      1.023179,
-      0.992978,
-      0.991627,
-      1.073585,
-      1.114629,
-      1.056403,
-      1.093012,
-      1.015109,
-      1.007069,
-      0.976026,
-      0.982324,
-      1.093805,
-      1.036086,
-      1.005727,
-      1.104078,
-      1.042298,
-      1.095522,
-      0.960255
+  "storage_dtype": "float16",
+  "latent_channels": 32,
+  "text_dim": 768,
+  "body": {
+    "kind": "пирамидальный U: ширина растёт к низу",
+    "widths": [
+      192,
+      512,
+      1024,
+      1536
     ],
-    "source": "vae/ (Muinez/f8c32-ae)"
-  },
-  "text": {
-    "encoder": "text_encoder/ (jhu-clsp/ettin-encoder-150m, 768-dim, 22 слоя)",
-    "max_len": 192,
-    "slices": [
-      2,
-      5,
+    "heads": [
+      6,
       8,
-      11,
-      14,
-      17,
-      20
+      8,
+      8
     ],
-    "slice_blocks": 2,
-    "trainable": "text_mixer.*"
-  },
-  "body": {
-    "kind": "пирамидальный U: ширина растёт к низу",
-    "graft": {
-      "e0x": 1,
-      "d0x": 1,
-      "e1": 1,
-      "d1": 1,
-      "e2": 2,
-      "d2": 2,
-      "e3": 2
-    },
     "blocks": {
       "e0": 6,
       "e0x": 1,
@@ -118,60 +34,37 @@
         "level": 0,
         "tokens": "÷1",
         "width": 192,
-        "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону",
-        "text": "кросс-внимание в каждом conv-блоке и в self-attn"
+        "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону"
       },
       {
         "level": 1,
         "tokens": "÷2",
         "width": 512,
-        "blocks": "attn 5 enc / 4 dec",
-        "text": "кросс-внимание"
+        "blocks": "attn 5 enc / 4 dec"
       },
       {
         "level": 2,
         "tokens": "÷4",
         "width": 1024,
-        "blocks": "attn 5 enc / 4 dec",
-        "text": "кросс-внимание"
+        "blocks": "attn 5 enc / 4 dec"
       },
       {
         "level": 3,
         "tokens": "÷8",
         "width": 1536,
-        "blocks": "attn 4 (боттлнек)",
-        "text": "кросс-внимание"
+        "blocks": "attn 4 (боттлнек)"
       }
     ],
     "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
-    "down": "mean-pool 2× + Linear",
-    "up": "Linear + nearest 2×",
-    "rope": "2D RoPE, base 100, фаза в латентных пикселях на каждом уровне",
-    "parametrization": "x0 латента"
-  },
-  "tail": {
-    "source": "декодер AE, 6 блоков @512, PixelShuffle ×3",
-    "trained_as": "flow по пикселям, условие z0 + (ошибка тела).detach()",
-    "sample_steps": 8
-  },
-  "params": {
-    "body": 238.54,
-    "text_mixer": 10.04,
-    "tail": 23.87,
-    "total": 272.45
-  },
-  "training": {
-    "dataset": "LucasFang/FLUX-Reason-6M",
-    "parts": "Aesthetics-Part01,Aesthetics-Part02,Imaginative,Text",
-    "size": 256,
-    "steps": 460000,
-    "batch": 24,
-    "accum": 1,
-    "lr": 5e-05,
-    "lr_tail": 3e-05,
-    "drop_text": 0.1,
-    "init": "tmp/run/ckpt/step_0404000.pt",
-    "metrics": ""
-  },
-  "storage_dtype": "float16"
+    "down": "mean-pool 2x + Linear",
+    "up": "Linear + nearest 2x",
+    "rope": {
+      "kind": "2d",
+      "base": 100,
+      "phase": "в латентных пикселях на каждом уровне"
+    },
+    "text": "кросс-внимание на всех уровнях",
+    "parametrization": "x0 латента",
+    "params": 238.54
+  }
 }