{"id":50504322,"url":"https://github.com/cobanov/bucketsampler","last_synced_at":"2026-06-02T14:30:43.847Z","repository":{"id":358114287,"uuid":"1240058304","full_name":"cobanov/bucketsampler","owner":"cobanov","description":"Aspect ratio bucketing for diffusion training. PyTorch native, DDP correct, no training-framework lock-in.","archived":false,"fork":false,"pushed_at":"2026-05-15T21:10:03.000Z","size":134,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-05-15T21:14:58.624Z","etag":null,"topics":["aspect-ratio","bucketing","data-pipeline","dataloader","diffusion","hf-datasets","huggingface","machine-learning","pytorch","sdxl","stable-diffusion","training","vae"],"latest_commit_sha":null,"homepage":"https://pypi.org/project/bucketsampler/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cobanov.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-05-15T18:13:50.000Z","updated_at":"2026-05-15T21:10:17.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/cobanov/bucketsampler","commit_stats":null,"previous_names":["cobanov/bucketsampler"],"tags_count":3,"template":false,"template_full_name":null,"purl":"pkg:github/cobanov/bucketsampler","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cobanov%2Fbucketsampler","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cobanov%2Fbucketsampler/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cobanov%2Fbucketsampler/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cobanov%2Fbucketsampler/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cobanov","download_url":"https://codeload.github.com/cobanov/bucketsampler/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cobanov%2Fbucketsampler/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":33827062,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-02T02:00:07.132Z","response_time":109,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["aspect-ratio","bucketing","data-pipeline","dataloader","diffusion","hf-datasets","huggingface","machine-learning","pytorch","sdxl","stable-diffusion","training","vae"],"created_at":"2026-06-02T14:30:42.758Z","updated_at":"2026-06-02T14:30:43.839Z","avatar_url":"https://github.com/cobanov.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# bucketsampler\n\n\u003e Aspect ratio bucketing for diffusion model training (SDXL-style multi-AR\n\u003e batches). PyTorch native, DDP correct, zero training-framework lock-in.\n\n\u003e Heads up: \"bucket sampler\" also names a length-bucketing pattern in NLP.\n\u003e This is the image / diffusion variant, not the sequence one.\n\n## Why\n\nDiffusion U-Nets want a fixed `(H, W)` per batch. Real datasets do not. The\nnaive options either distort (squeeze every image to a square) or throw data\naway (center-crop to the smallest common size). Bucketing splits images into a\nsmall set of `(W, H)` targets and draws each batch from a single bucket, so\nnothing gets squished and nothing gets dropped.\n\n`bucketsampler` ships the plumbing: assignment, dataset wrapper, DDP-correct\nsampler, presets for SDXL / SD1.5 / NovelAI, and a CLI to inspect your data\nbefore you start training.\n\n## Install\n\n```bash\npip install bucketsampler             # core (no torch)\npip install \"bucketsampler[torch]\"    # + PyTorch integration\npip install \"bucketsampler[hf]\"       # + HuggingFace datasets adapter\npip install \"bucketsampler[cache]\"    # + parquet metadata cache\npip install \"bucketsampler[vae]\"      # + VAE latent precomputation\npip install \"bucketsampler[analyze]\"  # + HTML reports from the analyzer\n```\n\n## 30-second quickstart\n\n```python\nfrom pathlib import Path\nfrom torch.utils.data import DataLoader\nfrom bucketsampler import (\n    BucketBatchSampler,\n    BucketedDataset,\n    FixedBuckets,\n    load_preset,\n)\n\npaths = sorted(Path(\"data/\").glob(\"*.jpg\"))\nstrategy = FixedBuckets(load_preset(\"sdxl\"))\n\ndataset = BucketedDataset(paths=paths, strategy=strategy)\nsampler = BucketBatchSampler(dataset, batch_size=4)\nloader = DataLoader(dataset, batch_sampler=sampler)\n\nfor batch in loader:\n    images = batch[\"image\"]   # [4, 3, H, W], same (H, W) within a batch\n    buckets = batch[\"bucket\"] # list[Bucket], one per sample\n    # ... feed images to your VAE / U-Net / etc.\n```\n\n## Inspect your dataset\n\nBefore you commit to a bucket set, see how your images actually distribute:\n\n```bash\nbucketsampler analyze data/ --preset sdxl\nbucketsampler analyze data/ --preset sdxl --json \u003e report.json\nbucketsampler analyze data/ --preset sdxl --html report.html\n```\n\nThe report shows readable / broken counts, AR distribution, per-bucket counts,\nunderutilized buckets (so you know what to drop), and outliers (extreme ARs\nthat match no bucket well, often a sign of bad data).\n\n## HuggingFace datasets\n\nAlready have a `datasets.Dataset` of images and captions? Skip the\nfile-list step:\n\n```python\nfrom datasets import load_dataset\nfrom bucketsampler import BucketedDataset, FixedBuckets, load_preset\n\nhf = load_dataset(\"lambdalabs/pokemon-blip-captions\", split=\"train\")\n\ndataset = BucketedDataset.from_hf(\n    hf,\n    FixedBuckets(load_preset(\"sdxl\")),\n    image_column=\"image\",\n    caption_column=\"text\",\n)\n```\n\nThe adapter accepts PIL columns (the common case), raw bytes\n(`datasets.Image(decode=False)`), and numpy / torch tensor columns. CHW\nvs HWC is auto-detected by the small-channel axis. Streaming\n(`IterableDataset`) is not yet supported, pass a map-style dataset.\n\n## DDP\n\nSame sampler, two extra kwargs:\n\n```python\nsampler = BucketBatchSampler(\n    dataset,\n    batch_size=4,\n    num_replicas=world_size,\n    rank=rank,\n)\nfor epoch in range(num_epochs):\n    sampler.set_epoch(epoch)   # required, reseeds the per-bucket shuffle\n    for batch in DataLoader(dataset, batch_sampler=sampler):\n        ...\n```\n\nAll ranks yield the same number of batches per epoch and see disjoint\nindices, so gradient sync stays happy.\n\n## Auto-generate buckets from your dataset\n\nPresets are reasonable defaults, but a bucket set picked from *your*\ndistribution always crops less. Let `bucketsampler` derive one for you:\n\n```bash\nbucketsampler buckets-from-dataset data/ \\\n    --num 8 --target 1024 \\\n    --compare-to sdxl \\\n    --output my_buckets.toml\n```\n\nSample output:\n\n```\nGenerated 8 buckets (requested 8, k-means converged in 6 iterations):\n  640x1536      AR= 0.417  cluster log-AR=-0.875  size=120\n  768x1280      AR= 0.600  cluster log-AR=-0.511  size=210\n  ...\nMean crop loss (auto):       3.42%\nMean crop loss (sdxl):       7.18%  (auto is -3.76 pts vs sdxl)\n\nWrote bucket set to my_buckets.toml\n```\n\nThen feed the TOML straight back into your training script:\n\n```python\nfrom bucketsampler import FixedBuckets, load_from_toml\n\nstrategy = FixedBuckets(load_from_toml(\"my_buckets.toml\"))\n```\n\nOr do it inline (no file involved) with the `AutoBuckets` convenience:\n\n```python\nimport numpy as np\nfrom bucketsampler import AutoBuckets\n\n# dims = (N, 2) int array of (width, height); analyzer's scan returns this\nstrategy = AutoBuckets.from_dims(dims, num_buckets=8, target=1024)\n```\n\nThe algorithm is 1-D k-means on `log(width / height)`, then snaps each\ncluster center to a `(w, h)` whose product is close to `target^2` and\nwhose dims are multiples of `vae_factor` (default 64). Fully deterministic\nfor a given seed.\n\n## Custom buckets\n\nDrop a TOML file anywhere on disk:\n\n```toml\n# my_buckets.toml\nname = \"my-budget\"\nvae_factor = 8\n\n[[buckets]]\nwidth  = 768\nheight = 768\n\n[[buckets]]\nwidth  = 896\nheight = 640\n\n[[buckets]]\nwidth  = 640\nheight = 896\n```\n\n```python\nfrom bucketsampler import FixedBuckets, load_from_toml\n\nstrategy = FixedBuckets(load_from_toml(\"my_buckets.toml\"))\n```\n\nJSON is also supported via `load_from_json`. The bundled presets\n(`sdxl`, `sd15`, `novelai`) live in the same format.\n\n## Metadata cache\n\nHeader reads are cheap individually but expensive at 100K+ images.\nCache them once:\n\n```bash\nbucketsampler build-cache data/ --output data.cache.parquet\n```\n\nThen reuse on every subsequent run:\n\n```python\nfrom bucketsampler import BucketedDataset, FixedBuckets, MetadataCache, load_preset\n\ncache = MetadataCache.load(\"data.cache.parquet\")\ndataset = BucketedDataset(\n    paths=image_paths,\n    strategy=FixedBuckets(load_preset(\"sdxl\")),\n    metadata_cache=cache,\n)\n```\n\nCache invalidation is automatic per row: files whose `mtime` has\nchanged get re-read, new files get added, removed files are dropped.\nRe-run `build-cache --refresh` to refresh in place.\n\n## Precompute VAE latents\n\nMove the VAE forward pass off the training hot path:\n\n```bash\nbucketsampler precompute \\\n    data/ \\\n    --vae stabilityai/sdxl-vae \\\n    --output latents/ \\\n    --preset sdxl \\\n    --batch-size 8 \\\n    --dtype bfloat16\n```\n\nThen train against the precomputed latents:\n\n```python\nfrom bucketsampler import BucketBatchSampler, BucketedLatentDataset\nfrom torch.utils.data import DataLoader\n\ndataset = BucketedLatentDataset(\"latents/\")\nsampler = BucketBatchSampler(dataset, batch_size=8)\nloader = DataLoader(dataset, batch_sampler=sampler)\n\nfor batch in loader:\n    latents = batch[\"latents\"]   # [B, C, H/8, W/8]\n    captions = batch.get(\"caption\")\n    # ... feed latents straight to your U-Net\n```\n\nCustom VAE? Implement the tiny `VAEEncoder` protocol\n(`downsample_factor`, `latent_channels`, `scale_factor`, `encode`)\nand call `precompute_latents()` directly.\n\n## CLI cheatsheet\n\n```bash\nbucketsampler --help\nbucketsampler version\nbucketsampler presets [--json]\nbucketsampler analyze \u003cpath\u003e --preset sdxl [--json | --html report.html]\nbucketsampler buckets-from-dataset \u003cpath\u003e --num 8 --target 1024 [--output buckets.toml] [--compare-to sdxl]\nbucketsampler build-cache \u003cpath\u003e --output cache.parquet [--refresh]\nbucketsampler precompute \u003cpath\u003e --vae stabilityai/sdxl-vae --output latents/ --preset sdxl\n```\n\n## Examples\n\nSee [`examples/`](examples/) for runnable scripts:\n\n- [`minimal_training_loop.py`](examples/minimal_training_loop.py) — full end-to-end loop\n- [`auto_buckets_inline.py`](examples/auto_buckets_inline.py) — data-driven bucket sets\n- [`hf_dataset.py`](examples/hf_dataset.py) — train from a HuggingFace dataset\n- [`precompute_latents.py`](examples/precompute_latents.py) — VAE precompute + latent training\n- [`ddp_training.py`](examples/ddp_training.py) — distributed training sketch\n\n## Development\n\n```bash\ngit clone https://github.com/cobanov/bucketsampler\ncd bucketsampler\nuv venv \u0026\u0026 uv pip install -e \".[dev]\"\n\npytest                              # 313 tests, ~15s on a laptop CPU\nruff check bucketsampler tests      # lint\nruff format bucketsampler tests     # autoformat\nmypy bucketsampler                  # strict type check\n```\n\nContributions are welcome via pull request. The project follows\n[Conventional Commits](https://www.conventionalcommits.org/) and aims\nfor 100% line coverage on `core/`.\n\n## License\n\nMIT, see [`LICENSE`](LICENSE).\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcobanov%2Fbucketsampler","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcobanov%2Fbucketsampler","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcobanov%2Fbucketsampler/lists"}