model-adapter-chitudiffusion

name

model-adapter-chitudiffusion

description

Adapt a new image diffusion model into the ChituDiffusion repository. Use when the user gives a model name and wants Codex to pull or inspect the upstream reference code, run the original demo, quantify non-API-wrapper code, add a runtime adapter/stage-level API integration, wire attention backends and CFG/context parallelism, adapt FlexCache strategies, benchmark with ChituBench, update result.md, clean artifacts, and split categorized commits.

ChituDiffusion Model Adapter

Operating Rules

• Work in the current ChituDiffusion checkout. Use ./.venv/bin/python and the real launcher path (./.venv/bin/chitu run or the repo's srun wrapper) for acceptance.
• Put upstream/reference code under refs/ and model weights under ~/WORK/models unless the user gives another location.
• Prefer direct integration with existing runtime, adapter, scheduler, parallel, and FlexCache APIs. Do not add compatibility shims for obsolete paths unless the user asks.
• Treat generated outputs as evidence, not code. Commit only source changes, reusable benchmark scripts, final summaries, and final plots; delete smoke/debug runs and local-only symlinks before committing.
• Keep user updates short while running long Slurm or benchmark jobs, and reuse completed results where possible.

Workflow

1. Reference intake

   • Identify the official repo, license, model card, demo command, required weights, and expected image size/step defaults.
   • Clone or copy the upstream implementation into refs/<model-name>/.
   • Download weights to ~/WORK/models/<model-name>/ or verify existing weights.
   • Run the upstream demo first in the current virtualenv. If GPU resources are needed, use srun/script/srun_direct.sh; do not fall back to CPU for diffusion demo validation.

2. Gap report before heavy adaptation

   • Read the upstream pipeline, transformer, scheduler, VAE/image processor, prompt encoder, and attention code.
   • Classify code into:
     • directly covered by ChituDiffusion APIs,
     • needs a runtime adapter or stage-level wrapper,
     • needs custom attention/parallel logic,
     • needs scheduler or latent-shape handling,
     • benchmark-only glue.
   • Give the user a short plan and estimate of the non-wrapper code surface before large edits.

3. Runtime adapter

   • Add or extend a model runtime adapter under chitu_diffusion/runtime/adapter/.
   • Implement stage-level methods for loading components, preparing buffers, denoising one scheduler step, decoding, and output handling.
   • Normalize model-specific latent/token shapes at adapter boundaries. Prefer one authoritative shape variable instead of duplicate image_size/img_shape state.
   • Register the model name and connect request parameters through DiffusionTask/Generator without hiding user-set options.

4. Attention and parallelism

   • Wire the requested attention backends through the existing Chitu attention abstraction.
   • For joint text/image attention, verify whether text tokens need sharding. If encoder states are short, prefer keeping text replicated and sharding image tokens only.
   • Bring up CFG parallel first when CFG semantics exist. Then bring up Ulysses/context parallel. Only implement Ring/USP if it fits the model sequence semantics without a large separate attention implementation.
   • Validate 1/2/4/8 GPU cases with minimal prompts before full benchmarks.

5. FlexCache and acceleration

   • First adapt strategies that do not require official model-specific coefficients or hidden-state semantics, such as PAB, BlockDance, Cubic, or step-level cache strategies.
   • For cache methods from reference code, isolate the train-free algorithm in chitu_diffusion/flexcache/strategy/ and keep model-specific shape handling in the adapter or a small module under flexcache/modules/.
   • Add params to flexcache/params.py, registration to Generator._build_flexcache_strategy, and compute metrics so ChituBench can report speedups.
   • Test at least one end-to-end image before sweeping.

6. Benchmark and visualization

   • Use ChituBench scripts or add a thin model-specific benchmark script under ChituBench/scripts/.
   • Reuse previous runs by collecting summaries into a consolidated result directory instead of re-running model loads unnecessarily.
   • For FlexCache trade-off plots and visual sheets, use the flexcache-bench-visualizer skill.
   • Update ChituBench/result.md with commands, notes, summary table, readout, and plot paths.

7. Cleanup and commits

   • Run git diff --check and a targeted py_compile/compileall on touched Python files.
   • Delete smoke/debug result directories, local-only symlinks, stale plots, and generated bytecode.
   • Split commits by domain. A typical split is:
     • model/runtime/FlexCache adaptation,
     • benchmark scripts and recorded results,
     • docs or skill updates if they live inside the repo.
   • Report commit hashes and validation commands.

Acceptance Checklist

• Upstream demo ran successfully on GPU.
• ChituDiffusion generated at least one end-to-end image through the real launcher.
• Attention backend and parallel cases requested by the user were either validated or explicitly scoped out with a concrete reason.
• FlexCache strategies requested by the user have at least one correctness run and benchmark evidence.
• result.md and final plots match the measured data.
• Worktree is clean after categorized commits, except intentionally untracked external artifacts outside the repo.