lars
1a3c907571
CI / Lint (ruff check) (push) Successful in 29s
CI / Format (ruff format) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 28s
CI / Type check (ty) (push) Successful in 34s
CI / Format (ruff format) (pull_request) Successful in 34s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 36s
CI / Tests (push) Successful in 3m41s
CI / Tests (pull_request) Successful in 3m39s
Stage1Model, Stage2OneShot and Stage2Autoregressive each independently implemented ~90 near-identical lines of __init__ scaffolding: build-or-share cond_enc, particle_type_cfg normalisation, objective -> time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head classifier heads (plus their identical RuntimeError guards). Now unblocked by #33 (trunk registry), #34 (block-conditioning registry) and #36 (build_mlp_head), which settled what belongs in the shared base. Adds StageModel(nn.Module) owning all of that: __init__ builds/shares cond_enc and normalises particle_type_cfg; _build_trunk_and_heads, called by each subclass after it sets up its own conditioning-assembly modules (cond_enc alone for Stage1Model, a context-fusion path for the two Stage2 classes), builds the objective/time embedding/trunk and the n_sec_head/type_head guarded by the shared _require_n_sec_head/ _require_type_head (Stage1Model overrides the n_sec guard since its message points at stage 2, not stage 1). Public __init__ signatures, attribute names, and forward/predict_* behaviour are unchanged. Verified with a pre/post state_dict-key-set diff against the pre-refactor classes (bit-identical) before writing this commit, plus new parametrized tests pinning each class's state_dict key set and the generator -> time_emb contract the base now owns. tests/test_migration_ v02_v03.py's existing bit-identical old-vs-new forward comparison and the rest of tests/test_network.py's per-class coverage pass unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
137 lines
3.1 KiB
Python
137 lines
3.1 KiB
Python
"""Re-export shim.
|
|
|
|
`network.py` used to hold every network-related class in one 1742-line file.
|
|
It's now split by concern into `giant/model/{layers,encoders,routers,trunks,
|
|
history,models,_legacy,builders}.py` (issues.md Issue 8); this module just
|
|
re-exports the public surface so every existing `from giant.model.network
|
|
import X` call site keeps working unchanged.
|
|
"""
|
|
|
|
from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_state_dict
|
|
from giant.model.builders import build_critics, build_models
|
|
from giant.model.encoders import ConditionEncoder
|
|
from giant.model.history import (
|
|
HISTORY_REGISTRY,
|
|
AttentionHistory,
|
|
HistoryEncoder,
|
|
MarkovHistory,
|
|
_CausalAttnBlock,
|
|
build_history,
|
|
register_history,
|
|
)
|
|
from giant.model.layers import (
|
|
BLOCK_REGISTRY,
|
|
AdaLNResBlock,
|
|
ContextAdapter,
|
|
FilmResBlock,
|
|
ResBlock,
|
|
SinusoidalEmbedding,
|
|
_make_axis_mlp,
|
|
build_block,
|
|
build_mlp_head,
|
|
register_block,
|
|
)
|
|
from giant.model.models import (
|
|
CriticModel,
|
|
Stage1Model,
|
|
Stage2Autoregressive,
|
|
Stage2OneShot,
|
|
StageModel,
|
|
resolve_type_n_classes,
|
|
stage2_trunk_sec_dim,
|
|
stage2_type_dim,
|
|
)
|
|
from giant.model.objectives import (
|
|
OBJECTIVE_REGISTRY,
|
|
DdpmObjective,
|
|
FlowObjective,
|
|
Objective,
|
|
WganObjective,
|
|
build_objective,
|
|
register_objective,
|
|
)
|
|
from giant.model.routers import (
|
|
ROUTER_REGISTRY,
|
|
ComposedRouter,
|
|
EnergyRouter,
|
|
PdgRouter,
|
|
ProcessRouter,
|
|
Router,
|
|
_build_router_from_cfg,
|
|
_check_router_conditioning_compat,
|
|
_parse_composed_axes,
|
|
build_composed_router,
|
|
build_router,
|
|
register_router,
|
|
)
|
|
from giant.model.trunks import (
|
|
TRUNK_REGISTRY,
|
|
ExpertTrunk,
|
|
RoutedTrunk,
|
|
Trunk,
|
|
_route_forward,
|
|
build_expert_body,
|
|
build_trunk,
|
|
register_trunk,
|
|
)
|
|
|
|
__all__ = [
|
|
"AdaLNResBlock",
|
|
"AttentionHistory",
|
|
"BLOCK_REGISTRY",
|
|
"ComposedRouter",
|
|
"ConditionEncoder",
|
|
"ContextAdapter",
|
|
"CriticModel",
|
|
"DdpmObjective",
|
|
"EnergyRouter",
|
|
"ExpertTrunk",
|
|
"FilmResBlock",
|
|
"FlowObjective",
|
|
"HISTORY_REGISTRY",
|
|
"HistoryEncoder",
|
|
"MarkovHistory",
|
|
"OBJECTIVE_REGISTRY",
|
|
"Objective",
|
|
"PdgRouter",
|
|
"ProcessRouter",
|
|
"ROUTER_REGISTRY",
|
|
"ResBlock",
|
|
"RoutedTrunk",
|
|
"Router",
|
|
"SinusoidalEmbedding",
|
|
"Stage1Model",
|
|
"Stage2Autoregressive",
|
|
"Stage2OneShot",
|
|
"StageModel",
|
|
"TRUNK_REGISTRY",
|
|
"Trunk",
|
|
"WganObjective",
|
|
"_CausalAttnBlock",
|
|
"_build_router_from_cfg",
|
|
"_check_router_conditioning_compat",
|
|
"_make_axis_mlp",
|
|
"_migrate_legacy_model_config",
|
|
"_parse_composed_axes",
|
|
"_route_forward",
|
|
"build_block",
|
|
"build_composed_router",
|
|
"build_critics",
|
|
"build_expert_body",
|
|
"build_history",
|
|
"build_mlp_head",
|
|
"build_models",
|
|
"build_objective",
|
|
"build_router",
|
|
"build_trunk",
|
|
"migrate_legacy_state_dict",
|
|
"register_block",
|
|
"register_history",
|
|
"register_objective",
|
|
"register_router",
|
|
"register_trunk",
|
|
"resolve_type_n_classes",
|
|
"stage2_trunk_sec_dim",
|
|
"stage2_type_dim",
|
|
]
|