Repository

context-scaling

[Tech Report] Context Scaling: Scaling Properties of Text Conditioning in Visual Generation

Endpoint
Input
POST /api/generate_sp
{user_prompt, width, height}
POST /api/caption_image
{image_base64}
POST /api/generate_image
{prompt, width, height, num_steps, seed, cfg_scale}

Specs

  • EndpointInput
  • POST /api/generate_sp{user_prompt, width, height}
  • POST /api/caption_image{image_base64}
  • POST /api/generate_image{prompt, width, height, num_steps, seed, cfg_scale}
  • PathWhat it does
  • training/QwenImage DiT training + serving stack (FSDP, NaviT-packed batches, EMA, delta-merge / model-soup checkpoint tools).
  • evalkit/Multi-benchmark T2I eval (9 benchmarks), structured-prompt prompter pipeline, and the ED/GPG detailness metrics.
  • pe-rsft/Distributed rejection-sampling pipeline that generates SFT (and optionally DPO) data for the prompter.
  • demo/The in-process FastAPI + browser studio wrapping PE, captioner, and DiT.

README

Context Scaling: Scaling Properties of Text Conditioning in Visual Generation

Project Page HF Models GitHub License Python

Zilong Chen · Chaorui Deng · Kunchang Li · Hongyi Yuan · Haoqi Fan ByteDance Seed [email protected]

Text prompt scaling law overview

We study empirical scaling properties for text conditioning in visual
generation. Such properties have rarely been measured because diffusion
loss does not scale with the number of tokens in natural-language prompts.
Surprisingly, we find that the converged diffusion loss scales with the
amount of structured language in the prompt. To quantify structured
language, we adapt two complementary measures: a white-box likelihood
metric (GPG) and a black-box attribute metric (ED). Across
controlled training runs, the converged diffusion loss decreases
approximately linearly with GPG and follows a power law with ED.

Guided by these scaling properties, we improve diffusability by
constructing structured prompts with semantic and geometric annotations
derived from images, and improve promptability by training a prompter
through supervised fine-tuning, cold-start, and verifier-gated on-policy
distillation. The resulting system outperforms all evaluated open-weight
models on nearly every compositional, reasoning, and world-knowledge
benchmark, while matching or surpassing the strongest closed-weight models
on most evaluations.

Highlights

  • Scaling laws for text conditioning. A controlled fifteen-cell sweep
    showing that converged diffusion loss tracks two caption-information
    metrics (GPG, likelihood-based; ED, attribute-based) much more
    tightly than it tracks caption length.
  • Structured prompts (SP) as a diffusability upgrade — a JSON schema
    with semantic and geometric fields, annotated by expert models.
  • A trainable LLM prompter + captioner that turns a user request or
    an image into SP.
  • Zero-shot structured editing — SP exposes each factor as a named
    field, so targeted edits regenerate an image with the rest of the
    composition preserved.

Installation

git clone https://github.com/heheyas/context-scaling.git
cd context-scaling
pip install -r requirements.txt

CUDA 12, torch >= 2.4, flash-attn. The demo, ED/GPG scoring, and
serving pipelines run on a single GPU (2 for the demo); DiT training
expects multi-GPU FSDP.

Quickstart

Try the demo — PE → DiT + image captioner in the browser

Boots one FastAPI process on a 2-GPU box: the released prompter LLM
(heheyas/SP-PE-Qwen3.5-35B-A3B)
plus the fine-tuned DiT (heheyas/Qwen-Image-SP),
both loaded resident in GPU memory, fronted by an HTML studio at
http://localhost:7860/.

pip install -r demo/requirements.txt
python demo/build_index.py            # one-off: assemble the studio HTML
export CUDA_VISIBLE_DEVICES=0,1
python -m demo.app                    # → http://localhost:7860/

Both PE_CKPT and DIT_CKPT default to the released HF Hub repos;
first launch downloads ~110 GB and merges the DiT overlay shards into a
cached file under $HF_HOME/context-scaling/. Subsequent launches take
seconds.

The demo exposes three JSON endpoints:

Endpoint Input Output System prompt
POST /api/generate_sp {user_prompt, width, height} {sp: ...} rft_iter0_with_ratio_v6_synthesis_student.txt
POST /api/caption_image {image_base64} {sp: ...} image2json.txt — bboxes on a 1000×1000 grid
POST /api/generate_image {prompt, width, height, num_steps, seed, cfg_scale} {image_base64} —

generate_sp and caption_image share the same VLM weights (loaded
once); the difference is text-only vs multimodal chat template + the
corresponding system prompt. caption_image accepts either raw base64
or a full data:image/png;base64,… URL:

IMG=$(base64 -w0 /path/to/some_image.png)
curl -s -X POST http://localhost:7860/api/caption_image \
  -H "Content-Type: application/json" \
  -d "{\"image_base64\": \"$IMG\"}" | jq .sp

Full env-var reference (GPU pinning, HF token for gated repos, PE memory
caps) and Hugging Face Docker Space deployment recipe live in
demo/README.md.

Measure caption quality of a dataset — ED + GPG

Score one or both metrics over an (image, caption) JSONL. Use this to
compare captions from different sources (human, GPT-4o, Gemini, your
prompter) on the same image pool.

{"image": "/abs/path/to/img.png", "caption": "A short caption.", "uid": "row-0001"}
# ED only — no GPU, LLM-bound (Gemini-3-pro + GPT-4o).
export GEMINI_API_KEY_FILE=<path>       # one Gemini key per line
export GEMINI_API_CONFIG=<path>         # GPT-4o key JSON (matcher)
python -m evalkit.detailness.evaluator \
    --jsonl /path/to/dataset.jsonl \
    --out   /tmp/scored_ed.jsonl \
    --metrics ed

# ED + GPG — needs one GPU for Qwen3-VL-8B-Instruct forward passes.
python -m evalkit.detailness.evaluator \
    --jsonl /path/to/dataset.jsonl \
    --out   /tmp/scored_both.jsonl \
    --metrics ed,gpg \
    --gpg-model  <WEIGHTS_ROOT>/Qwen3-VL-8B-Instruct \
    --gpg-device cuda:0

Each output row is the input row plus:

{"ed": 0.27, "ed_F05_A": 0.27, "ed_P_A": 1.0, "ed_R_A": 0.07,
 "gpg": 0.382, "gpg_total_nats": 289.4, "gpg_n_tokens": 758, ...}

Higher is more informative. See
evalkit/detailness/README.md for
metric definitions, cache options, and per-metric CLIs.

Repo layout

Everything past the quickstart lives in per-package docs.

Path What it does Doc
training/ QwenImage DiT training + serving stack (FSDP, NaviT-packed batches, EMA, delta-merge / model-soup checkpoint tools). docs/training.md
evalkit/ Multi-benchmark T2I eval (9 benchmarks), structured-prompt prompter pipeline, and the ED/GPG detailness metrics. docs/evalkit.md
pe-rsft/ Distributed rejection-sampling pipeline that generates SFT (and optionally DPO) data for the prompter. docs/pe-rsft.md
demo/ The in-process FastAPI + browser studio wrapping PE, captioner, and DiT. demo/README.md

Citation

@article{chen2026contextscaling,
  title   = {Context Scaling: Scaling Properties of Text Conditioning in Visual Generation},
  author  = {Chen, Zilong and Deng, Chaorui and Li, Kunchang and Yuan, Hongyi and Fan, Haoqi},
  journal = {Technical Report, ByteDance Seed},
  year    = {2026}
}

Acknowledgements

This release stands on the shoulders of several open-source projects.
We are grateful to their authors and maintainers.

  • Qwen-Image (Alibaba) — DiT
    architecture, tokenizer, and reference pipeline that
    training/modeling/qwenimage/ and the demo's QwenImagePipeline are
    built on. Also Qwen2.5-VL for
    training/modeling/qwen2/.
  • Bagel — provided the
    Bagel training runs used as the ground-truth MSE side of the ED/GPG
    scaling-law measurements in evalkit/detailness/.
  • ms-swift (ModelScope) —
    used to serve and fine-tune the prompt-expansion LLM (PE) throughout
    the RFT loop in pe-rsft/.
  • huggingface/diffusers
    and huggingface/transformers
    — modeling primitives underpinning training/modeling/.
  • Evaluation benchmarks (vendored under evalkit/context-scaling/benchmarks/):
    GenEval, GenEval2, GenEval++ (Echo-4o), DPG-Bench (ELLA), OneIG-Bench,
    TIIF-Bench, T2I-CoReBench, GenExam, WISE — each used under its own
    upstream license (see per-benchmark LICENSE files and
    NOTICE for the full attribution list).

License

Apache License 2.0 — see LICENSE and NOTICE.