| from pathlib import Path |
| import shutil,json,hashlib |
| from PIL import Image |
| root=Path('/data/nyako-training');src=root/'source_dataset';dst=root/'artifacts/dataset';dst.mkdir(exist_ok=True) |
| prefix=b'@nyako, '; rows=[] |
| for p in sorted(src.rglob('*.jpg')): |
| c=p.with_suffix('.txt'); assert c.exists(), 'Missing caption' |
| rel=p.relative_to(src); target=dst/rel;target.parent.mkdir(parents=True,exist_ok=True);shutil.copy2(p,target) |
| |
| original=c.read_bytes(); amended=prefix+original |
| target.with_suffix('.txt').write_bytes(amended) |
| assert target.with_suffix('.txt').read_bytes()[len(prefix):]==original |
| with Image.open(p) as im: size=im.size |
| rows.append({'image':str(rel),'size':size,'original_caption_sha256':hashlib.sha256(original).hexdigest(),'prefixed_caption_sha256':hashlib.sha256(amended).hexdigest()}) |
| assert rows |
| (root/'artifacts/config/dataset_manifest.json').write_text(json.dumps({'repo':'ACCC1380/Nyako','revision':'f1048ee38473bc253ebdc4b4eba3b0c2dcebacb3','prefix':'@nyako, ','count':len(rows),'files':rows},indent=2)) |
| print('Prepared image/caption pairs:',len(rows)); print('Image dimensions range:',[min(r['size'][i] for r in rows) for i in (0,1)],[max(r['size'][i] for r in rows) for i in (0,1)]) |
| folders=sorted(set(str((dst/r['image']).parent) for r in rows)) |
| text='[general]\ncaption_extension = ".txt"\nshuffle_caption = false\n\n[[datasets]]\nresolution = 1024\nbatch_size = 2\nenable_bucket = true\nmin_bucket_reso = 256\nmax_bucket_reso = 2048\nbucket_reso_steps = 64\nbucket_no_upscale = false\n' |
| for folder in folders:text+='\n[[datasets.subsets]]\nimage_dir = '+json.dumps(folder)+'\nnum_repeats = 1\n' |
| (root/'artifacts/config/dataset.toml').write_text(text) |
|
|
