lars
4b2e0ba98e
CI / Format (ruff format) (push) Successful in 33s
CI / Lint (ruff check) (push) Successful in 36s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 27s
CI / Lint (ruff check) (pull_request) Successful in 27s
CI / Format (ruff format) (pull_request) Successful in 29s
CI / Tests (push) Successful in 3m33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
CI / Type check (ty) (pull_request) Successful in 32s
CI / Tests (pull_request) Successful in 2m45s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CriticModel was the one stage-shaped class left out of the trunk-registry (gitea #33), block-conditioning-registry (gitea #34), and StageModel-base (gitea #39) refactors: it hand-rolled a plain ResBlock stack, so a routed/FiLM/AdaLN trunk was available to every generative stage model except the critic competing against them under WGAN-GP. CriticModel now subclasses StageModel (reusing its cond_enc construction, and a stage-2 context-fusion helper factored out of Stage2OneShot onto the base) and builds its body via build_trunk (output width 1) instead of a bespoke ResBlock loop, so trunk.type/trunk.block_conditioning now affect the critic too. Each stage's critic inherits its own generator's trunk config rather than a new critic_trunk config key, mirroring the existing critic_hidden_dim/critic_n_res_blocks "0 = inherit from generator" pattern. Router mixing (MoE) for the critic stays out of scope. Since CriticModel is training-only and never persisted for inference, and WGAN-GP is still unbenchmarked, its state_dict shape has no back-compat burden. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
235 lines
10 KiB
Python
235 lines
10 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.objectives import build_objective
|
|
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 = ConditioningConfig.from_dict(cfg["conditioning"])
|
|
particle_cfg = conditioning_cfg.particle
|
|
material_cfg = conditioning_cfg.material
|
|
particle_conditioning = particle_cfg.type
|
|
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
|
|
objective = build_objective(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 objective.needs_time 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,
|
|
trunk_type=s1_spec.trunk.type,
|
|
block_conditioning=s1_spec.trunk.block_conditioning,
|
|
n_sec_head_k_max=n_sec_head_k_max,
|
|
cond_enc=shared_cond_enc,
|
|
n_sec_head_cfg=s1_spec.heads.n_sec.to_dict(),
|
|
)
|
|
|
|
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
|
|
objective = build_objective(generator)
|
|
# wgan has no time_dim concept — see the matching comment in stage 1
|
|
# above.
|
|
time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64
|
|
n_sec_owner = s2_spec.n_sec.owner
|
|
stop_token = s2_spec.n_sec.mode == "stop_token"
|
|
k_max = s2_spec.k_max
|
|
particle_type_cfg = s2_spec.particle_type
|
|
|
|
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,
|
|
trunk_type=s2_spec.trunk.type,
|
|
block_conditioning=s2_spec.trunk.block_conditioning,
|
|
build_n_sec_head=n_sec_owner != "stage1" and not stop_token,
|
|
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,
|
|
n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
|
|
type_head_cfg=s2_spec.heads.type.to_dict(),
|
|
build_stop_head=stop_token,
|
|
stop_sampling=s2_spec.n_sec.stop_sampling,
|
|
stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
|
|
)
|
|
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,
|
|
trunk_type=s2_spec.trunk.type,
|
|
block_conditioning=s2_spec.trunk.block_conditioning,
|
|
build_n_sec_head=n_sec_owner != "stage1",
|
|
particle_type_cfg=particle_type_cfg,
|
|
cond_enc=shared_cond_enc,
|
|
n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
|
|
type_head_cfg=s2_spec.heads.type.to_dict(),
|
|
)
|
|
|
|
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 = ConditioningConfig.from_dict(cfg["conditioning"])
|
|
particle_cfg = conditioning_cfg.particle
|
|
material_cfg = conditioning_cfg.material
|
|
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 build_objective(s1_spec.generator).is_adversarial:
|
|
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",
|
|
trunk_type=s1_spec.trunk.type,
|
|
block_conditioning=s1_spec.trunk.block_conditioning,
|
|
)
|
|
|
|
if s2_spec.active and build_objective(s2_spec.generator).is_adversarial:
|
|
k_max = s2_spec.k_max
|
|
particle_type_cfg = s2_spec.particle_type
|
|
in_dim = stage2_trunk_sec_dim(
|
|
particle_type_cfg,
|
|
s2_spec.generator,
|
|
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,
|
|
trunk_type=s2_spec.trunk.type,
|
|
block_conditioning=s2_spec.trunk.block_conditioning,
|
|
)
|
|
|
|
return result
|