Add an Objective registry for the flow/ddpm/wgan generator choice (gitea #32) #51
Assignees
lars (Lars Bogner)
Clear assignees
lars (Lars Bogner)
No Assignees
Notifications
Due Date
No due date set.
Dependencies
No dependencies set.
Reference: lars/giant#51
Reference in New Issue
Block a user
Delete Branch "fix/issue-32"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
generator ∈ {"flow", "ddpm", "wgan"} was tested as a bare string in ~45
sites across models.py, sample.py, builders.py, trainers.py, and
stage2_inputs.py, each independently re-deriving one of five consequences
of the choice (needs a time embedding? what does the trunk take as input?
is the type slice folded into the trunk output? which sampler? which
loss?). giant/model/objectives.py adds an Objective ABC + OBJECTIVE_REGISTRY
bare-string site now goes through it (needs_time, is_adversarial,
folds_type_slice, trunk_in_dim, build_schedule, stage1_loss/stage2_loss).
Per discussion: FlowDDPMStageTrainer and WGANStageTrainer stay separate
classes rather than merging into one StageTrainer as the issue's sketch
proposed — their training loops are genuinely different shapes (single loss
vs. dual G/D step with gradient penalty/n_critic/ST-Gumbel), and trainers.py
is the least-covered-by-fast-tests part of the codebase, so a full merge
was judged out of proportion to this issue's risk budget.
FlowDDPMStageTrainer's own loss dispatch (flow vs ddpm, one-shot vs AR) does
move onto the objective, so a future non-adversarial objective (rectified
flow, consistency distillation) is still a one-file, zero-trainer-edits
addition.
No config-schema change — stage{1,2}_model.generator stays the persisted
string, just looked up in the registry instead of string-compared. An
unrecognized generator value now fails fast with a clear ValueError instead
of silently falling through some bare-string checks and not others (same
behavior build_router/build_trunk already have for their own type keys).
Co-Authored-By: Claude Opus 5 noreply@anthropic.com
generator ∈ {"flow", "ddpm", "wgan"} was tested as a bare string in ~45 sites across models.py, sample.py, builders.py, trainers.py, and stage2_inputs.py, each independently re-deriving one of five consequences of the choice (needs a time embedding? what does the trunk take as input? is the type slice folded into the trunk output? which sampler? which loss?). giant/model/objectives.py adds an Objective ABC + OBJECTIVE_REGISTRY + build_objective factory, mirroring routers.py's Router pattern, and every bare-string site now goes through it (needs_time, is_adversarial, folds_type_slice, trunk_in_dim, build_schedule, stage1_loss/stage2_loss). Per discussion: FlowDDPMStageTrainer and WGANStageTrainer stay separate classes rather than merging into one StageTrainer as the issue's sketch proposed — their training loops are genuinely different shapes (single loss vs. dual G/D step with gradient penalty/n_critic/ST-Gumbel), and trainers.py is the least-covered-by-fast-tests part of the codebase, so a full merge was judged out of proportion to this issue's risk budget. FlowDDPMStageTrainer's own loss dispatch (flow vs ddpm, one-shot vs AR) does move onto the objective, so a future non-adversarial objective (rectified flow, consistency distillation) is still a one-file, zero-trainer-edits addition. No config-schema change — stage{1,2}_model.generator stays the persisted string, just looked up in the registry instead of string-compared. An unrecognized generator value now fails fast with a clear ValueError instead of silently falling through some bare-string checks and not others (same behavior build_router/build_trunk already have for their own type keys). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>