lars
32aa5a5f92
conditioning.particle.emb_dim and stage2_model.particle_type.target="onehot"'s class count were silently the same number everywhere (pipeline.py's PDG top-N map build, Stage2OneShot/Stage2Autoregressive's type head, StageSpec's training loss width, the checkpoint's shared pdg_topn_map), fixing the secondary-species vocabulary at whatever width the unrelated physical-conditioning MLP happened to use — the exact vocabulary the v0.3.0 pivot exists to fix. Adds stage2_model.particle_type.n_classes (default 0 = inherit conditioning.particle.emb_dim, preserving today's behavior and every existing checkpoint) and a single resolve_type_n_classes helper used everywhere the coupling used to be implicit. Splits the checkpoint's shared pdg_topn_map into a conditioning-only pdg_topn_map and a new sec_type_topn_map, built independently through the existing (axis, n_classes)-keyed setup cache (no extra scan when they still resolve to the same N) and threaded through giant predict/giant rollout's decode path. A checkpoint with no sec_type_topn_map key (pre-#29) falls back to reusing pdg_topn_map, reproducing the old shared behavior exactly. Decided with the user during planning: commit directly on this branch; represent the split as an additive sec_type_topn_map checkpoint key rather than conditionally reusing pdg_topn_map; build the two top-N maps independently rather than the issue's proposed build-at-max-and-slice, since the setup cache already avoids redundant scans across runs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
212 lines
9.1 KiB
Python
212 lines
9.1 KiB
Python
"""Factories: `build_models`/`build_critics` assemble the top-level stage
|
|
models from a config dict (issues.md Issue 8)."""
|
|
|
|
import torch.nn as nn
|
|
|
|
from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
|
|
from giant.constants import X_DIM
|
|
from giant.model._legacy import _migrate_legacy_model_config
|
|
from giant.model.encoders import ConditionEncoder
|
|
from giant.model.models import (
|
|
CriticModel,
|
|
Stage1Model,
|
|
Stage2Autoregressive,
|
|
Stage2OneShot,
|
|
resolve_type_n_classes,
|
|
stage2_trunk_sec_dim,
|
|
)
|
|
from giant.model.routers import Router, _build_router_from_cfg
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factories
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
|
"""Construct `{"stage1": ..., "stage2": ...}` from a config dict — either
|
|
the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/
|
|
`"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's
|
|
flat `model_config`, auto-migrated via `_migrate_legacy_model_config`.
|
|
|
|
A stage is `None` in the result when that stage's `active = False`.
|
|
`stage2_model.router.tie_to_stage1` shares stage 1's literal `Router`
|
|
instance rather than building a second, independently-parameterized one
|
|
(v0.2's actual — probably accidental — behaviour: two routers built from
|
|
one config with no semantic relationship between them).
|
|
|
|
`conditioning.share_stages = true` builds one `ConditionEncoder`
|
|
instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
|
|
`Stage2Autoregressive`'s `cond_enc` param), instead of each stage
|
|
building its own — halving the conditioning parameter count and forcing a
|
|
common representation. `false` (default) keeps v0.2 behaviour:
|
|
independent instances with identical config but independent weights.
|
|
"""
|
|
cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
|
|
pdg_vocab = cfg["pdg_vocab"]
|
|
mat_vocab = cfg["mat_vocab"]
|
|
conditioning = cfg["conditioning"]
|
|
particle_cfg = conditioning["particle"]
|
|
material_cfg = conditioning["material"]
|
|
particle_conditioning = particle_cfg["type"]
|
|
conditioning_cfg = ConditioningConfig.from_dict(conditioning)
|
|
s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
|
|
s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
|
|
cond_out_dim = conditioning_cfg.out_dim
|
|
shared_cond_enc: ConditionEncoder | None = None
|
|
if conditioning_cfg.share_stages:
|
|
shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
|
|
|
|
result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
|
|
|
|
stage1_router: Router | None = None
|
|
if s1_spec.active:
|
|
router_cfg = cfg["stage1_model"].get("router") or {}
|
|
if s1_spec.router.enabled:
|
|
stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
|
|
generator = s1_spec.generator
|
|
# wgan has no time_dim concept (no diffusion/flow time variable) —
|
|
# matches the pre-dataclass .get("time_dim", 64) fallback, which
|
|
# always hit its default for a wgan sub-block too.
|
|
time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64
|
|
n_sec_owner = s2_spec.n_sec.owner
|
|
n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
|
|
result["stage1"] = Stage1Model(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=material_cfg,
|
|
hidden_dim=s1_spec.hidden_dim,
|
|
n_res_blocks=s1_spec.n_res_blocks,
|
|
cond_out_dim=cond_out_dim,
|
|
dropout=s1_spec.dropout,
|
|
generator=generator,
|
|
time_dim=time_dim,
|
|
noise_dim=s1_spec.wgan.noise_dim,
|
|
router=stage1_router,
|
|
n_sec_head_k_max=n_sec_head_k_max,
|
|
cond_enc=shared_cond_enc,
|
|
)
|
|
|
|
if s2_spec.active:
|
|
decoder = s2_spec.decoder
|
|
router_cfg = cfg["stage2_model"].get("router") or {}
|
|
stage2_router: Router | None = None
|
|
if s2_spec.router.enabled:
|
|
if s2_spec.router.tie_to_stage1 and stage1_router is not None:
|
|
stage2_router = stage1_router
|
|
else:
|
|
stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
|
|
generator = s2_spec.generator
|
|
# wgan has no time_dim concept — see the matching comment in stage 1
|
|
# above.
|
|
time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
|
|
n_sec_owner = s2_spec.n_sec.owner
|
|
k_max = s2_spec.k_max
|
|
particle_type_cfg = s2_spec.particle_type.to_dict()
|
|
|
|
if decoder == "autoregressive":
|
|
ar_cfg = s2_spec.autoregressive
|
|
result["stage2"] = Stage2Autoregressive(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=material_cfg,
|
|
hidden_dim=s2_spec.hidden_dim,
|
|
n_res_blocks=s2_spec.n_res_blocks,
|
|
cond_out_dim=cond_out_dim,
|
|
context_dim=s2_spec.context_dim,
|
|
dropout=s2_spec.dropout,
|
|
generator=generator,
|
|
time_dim=time_dim,
|
|
noise_dim=s2_spec.wgan.noise_dim,
|
|
k_max=k_max,
|
|
router=stage2_router,
|
|
build_n_sec_head=n_sec_owner != "stage1",
|
|
particle_type_cfg=particle_type_cfg,
|
|
history=ar_cfg.history,
|
|
attn_n_heads=ar_cfg.attn_n_heads,
|
|
attn_n_layers=ar_cfg.attn_n_layers,
|
|
cond_enc=shared_cond_enc,
|
|
)
|
|
else:
|
|
sec_dim = stage2_trunk_sec_dim(
|
|
particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
|
|
)
|
|
result["stage2"] = Stage2OneShot(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=material_cfg,
|
|
hidden_dim=s2_spec.hidden_dim,
|
|
n_res_blocks=s2_spec.n_res_blocks,
|
|
cond_out_dim=cond_out_dim,
|
|
context_dim=s2_spec.context_dim,
|
|
sec_dim=sec_dim,
|
|
dropout=s2_spec.dropout,
|
|
generator=generator,
|
|
time_dim=time_dim,
|
|
noise_dim=s2_spec.wgan.noise_dim,
|
|
k_max=k_max,
|
|
router=stage2_router,
|
|
build_n_sec_head=n_sec_owner != "stage1",
|
|
particle_type_cfg=particle_type_cfg,
|
|
cond_enc=shared_cond_enc,
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
|
|
"""Construct `{"stage1": ..., "stage2": ...}` critics for `generator =
|
|
"wgan"` training. Training-only — never persisted for inference the way
|
|
`build_models`'s pair is. `None` for a stage that's inactive or not
|
|
WGAN."""
|
|
cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
|
|
pdg_vocab = cfg["pdg_vocab"]
|
|
mat_vocab = cfg["mat_vocab"]
|
|
conditioning = cfg["conditioning"]
|
|
particle_cfg = conditioning["particle"]
|
|
material_cfg = conditioning["material"]
|
|
conditioning_cfg = ConditioningConfig.from_dict(conditioning)
|
|
cond_out_dim = conditioning_cfg.out_dim
|
|
s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
|
|
s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
|
|
|
|
result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
|
|
|
|
if s1_spec.active and s1_spec.generator == "wgan":
|
|
result["stage1"] = CriticModel(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=material_cfg,
|
|
in_dim=X_DIM,
|
|
hidden_dim=s1_spec.wgan.critic_hidden_dim or s1_spec.hidden_dim,
|
|
n_res_blocks=s1_spec.wgan.critic_n_res_blocks or s1_spec.n_res_blocks,
|
|
cond_out_dim=cond_out_dim,
|
|
dropout=s1_spec.dropout,
|
|
stage="stage1",
|
|
)
|
|
|
|
if s2_spec.active and s2_spec.generator == "wgan":
|
|
k_max = s2_spec.k_max
|
|
particle_type_cfg = s2_spec.particle_type.to_dict()
|
|
in_dim = stage2_trunk_sec_dim(
|
|
particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
|
|
)
|
|
result["stage2"] = CriticModel(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
particle_cfg=particle_cfg,
|
|
material_cfg=material_cfg,
|
|
in_dim=in_dim,
|
|
hidden_dim=s2_spec.wgan.critic_hidden_dim or s2_spec.hidden_dim,
|
|
n_res_blocks=s2_spec.wgan.critic_n_res_blocks or s2_spec.n_res_blocks,
|
|
cond_out_dim=cond_out_dim,
|
|
dropout=s2_spec.dropout,
|
|
stage="stage2",
|
|
context_dim=s2_spec.context_dim,
|
|
)
|
|
|
|
return result
|