Blame for scripts/package_upload.py

"""Validate completed training and upload artifacts without training images."""
import hashlib
import json
import os
import shutil
import subprocess
import time
import zipfile
from pathlib import Path

import torch
import toml
from huggingface_hub import HfApi
from safetensors import safe_open

ROOT = Path('/data/realstyle2')
OUT = ROOT / 'output'
UPLOAD = ROOT / 'upload'
REPO = 'ACCC1380/realstyle2_anima'
config = toml.load(ROOT / 'configs/train.toml')
final = OUT / 'realstyle2_anima.safetensors'
assert final.is_file(), 'Final checkpoint missing'
torch.set_num_threads(4)
with safe_open(final, framework='pt') as model:
    metadata = model.metadata()
    epochs = int(metadata['ss_epoch'])
    steps = int(metadata['ss_steps'])
    for key in model.keys():
        assert torch.isfinite(model.get_tensor(key)).all(), f'Non-finite tensor: {key}'
assert epochs in {30, 35, 40}, f'Unexpected epoch count: {epochs}'
numbered_final = OUT / f'realstyle2_anima-{epochs:06d}.safetensors'
if not numbered_final.exists():
    shutil.copyfile(final, numbered_final)
checks = []
for epoch in range(1, epochs + 1):
    path = OUT / f'realstyle2_anima-{epoch:06d}.safetensors'
    assert path.is_file(), f'Missing epoch {epoch}'
    with safe_open(path, framework='pt') as checkpoint:
        meta = checkpoint.metadata()
        assert int(meta['ss_epoch']) == epoch
        assert checkpoint.keys()
        for key in checkpoint.keys():
            assert torch.isfinite(checkpoint.get_tensor(key)).all(), f'Non-finite tensor at epoch {epoch}: {key}'
    checks.append({'file': path.name, 'epoch': epoch, 'steps': int(meta['ss_steps']),
                   'bytes': path.stat().st_size, 'sha256': hashlib.file_digest(path.open('rb'), 'sha256').hexdigest()})
summary = {'status': 'training_complete', 'epochs': epochs, 'steps': steps,
           'batch_size': config['train_batch_size'], 'gradient_checkpointing': config['gradient_checkpointing'],
           'optimizer': config['optimizer_type'], 'learning_rate': config['learning_rate'],
           'scheduler': config['lr_scheduler'], 'rank': config['network_dim'], 'alpha': config['network_alpha'],
           'source_requested': 'ACCC1380/realstyle2', 'source_actual': 'ACCC1380/realskin2',
           'destination': REPO, 'checkpoints': checks, 'final_metadata': metadata,
           'trainer_commit': subprocess.check_output(['git', '-C', '/data/lora-scripts-next', 'rev-parse', 'HEAD'], text=True).strip()}
(ROOT / 'metadata/training_summary.json').write_text(json.dumps(summary, indent=2))
(ROOT / 'metadata/environment.txt').write_text(subprocess.check_output(['/data/venv-training/bin/pip', 'freeze'], text=True))
for folder in ['configs', 'scripts', 'metadata']:
    shutil.copytree(ROOT / folder, UPLOAD / folder, dirs_exist_ok=True)
shutil.copytree(ROOT / 'logs', UPLOAD / 'logs', dirs_exist_ok=True)
for path in OUT.glob('*.safetensors'):
    dest = UPLOAD / path.name
    if not dest.exists():
        shutil.copyfile(path, dest)
for state in OUT.glob('*state'):
    if state.is_dir():
        shutil.copytree(state, UPLOAD / 'resume_state' / state.name, dirs_exist_ok=True)
with zipfile.ZipFile(UPLOAD / 'training_data_without_images.zip', 'w', zipfile.ZIP_DEFLATED, compresslevel=1) as archive:
    for folder in ['dataset', 'configs', 'scripts', 'metadata']:
        for path in sorted((ROOT / folder).rglob('*')):
            if path.is_file() and path.suffix.lower() not in {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif', '.pyc'}:
                archive.write(path, path.relative_to(ROOT))
with zipfile.ZipFile(UPLOAD / 'lora_scripts_next_source.zip', 'w', zipfile.ZIP_DEFLATED, compresslevel=1) as archive:
    source = Path('/data/lora-scripts-next')
    for path in sorted(source.rglob('*')):
        relative = path.relative_to(source)
        if path.is_file() and not any(p in {'.git', '__pycache__', 'node_modules', 'huggingface', 'logs'} for p in relative.parts):
            archive.write(path, relative)
stats = json.loads((ROOT / 'metadata/tag_summary.json').read_text())
card = f'''---
base_model: circlestone-labs/Anima
library_name: diffusers
tags:
  - lora
  - anima
  - text-to-image
  - style
---
# realstyle2 Anima LoRA

基于本地 `anima-base-v1.0.safetensors`，使用 lora-scripts-next 标准 Anima 后端训练。

**数据来源说明：** 请求中的 `ACCC1380/realstyle2` 在执行时不存在（model/dataset 均返回 404）。根据同账号仓库，将其视为名称笔误，实际使用 `ACCC1380/realskin2` 的 `realskin2.zip`（revision `fda57fa984e9cf9cbb964929a6a8431b59976a28`）。未经用户再次确认，因为用户明确要求无人值守完成。

- 数据：{stats['images']} 张图片，1 repeat；训练集图片未上传至本仓库。
- 打标：PixAI Tagger v1.0，平均 {stats['mean_tags']:.1f} 个标签/图；阈值详见 `metadata/tagger_settings.json`。
- 触发词：`realstyle2`。
- 完成：{epochs} epochs，{steps} steps，每轮均保存。`realstyle2_anima.safetensors` 为最终权重，与第 {epochs} 轮相同。
- Batch size：{config['train_batch_size']}；梯度检查点关闭；gradient accumulation 1。
- AdamW，learning rate 1e-4，constant，无 warmup。
- LoRA rank 32 / alpha 16；BF16；1024 分桶；冻结文本编码器；缓存 latents 和文本编码。
- timestep sampling sigmoid，flow shift 1.0，uniform loss weighting。
- 所有 epoch 权重均检查 tensor 有限值和 epoch 元数据。未进行人工生成效果评测。

## 文件

`realstyle2_tags.zip`：与原图同名的 caption TXT。

`training_data_without_images.zip`：caption、latent/text encoder 缓存、配置、脚本、元数据，排除训练图片。

`lora_scripts_next_source.zip`：所用训练器源码快照；依赖见 `metadata/environment.txt`。

`resume_state/`：训练恢复状态；`logs/`：完整日志及 TensorBoard 数据；`metadata/training_summary.json`：训练参数、基础模型元数据和各轮 SHA256。

所有本地任务文件位于 `/data/realstyle2`；训练器位于 `/data/lora-scripts-next`。原有基础模型从 `/marimo` 读取。
'''
(UPLOAD / 'README.md').write_text(card)
for path in UPLOAD.rglob('*'):
    assert not (path.is_file() and path.suffix.lower() in {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif'}), f'Unexpected image: {path}'
api = HfApi()
for attempt in range(1, 7):
    try:
        print(f'Upload attempt {attempt}: {REPO}', flush=True)
        commit = api.upload_folder(repo_id=REPO, folder_path=UPLOAD,
            commit_message=f'Add {epochs}-epoch Anima LoRA, PixAI captions and training artifacts (no dataset images)',
            ignore_patterns=['.cache/**', '*.pyc'])
        expected = {p.relative_to(UPLOAD).as_posix(): p.stat().st_size for p in UPLOAD.rglob('*') if p.is_file() and '.cache' not in p.parts}
        info = api.repo_info(REPO, files_metadata=True)
        remote = {p.rfilename: p for p in info.siblings}
        for name, size in expected.items():
            assert name in remote, f'Missing upload: {name}'
            assert remote[name].size == size, f'Size mismatch: {name}'
        for check in checks:
            lfs = remote[check['file']].lfs
            if lfs:
                assert lfs.sha256 == check['sha256'], f'Hash mismatch: {check["file"]}'
        receipt = {'status': 'complete', 'repo': REPO, 'commit': commit.oid,
                   'url': commit.commit_url, 'verified_files': len(expected),
                   'bytes': sum(expected.values()), 'epochs': epochs, 'steps': steps,
                   'verified_at': time.time()}
        (ROOT / 'upload_receipt.json').write_text(json.dumps(receipt, indent=2))
        api.upload_file(repo_id=REPO, path_or_fileobj=ROOT / 'upload_receipt.json', path_in_repo='upload_receipt.json', commit_message='Record verified upload receipt')
        print(json.dumps(receipt), flush=True)
        break
    except Exception as exc:
        print(f'Upload failed {type(exc).__name__}: {exc}', flush=True)
        if attempt == 6:
            raise
        time.sleep(30)