lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Bump ruff line-length to 120 and reformat
CI / Lint (ruff check) (push) Successful in 31s
Details
CI / Format (ruff format) (push) Successful in 32s
Details
CI / Sync project version with tag (push) Has been skipped
Details
CI / Lint (ruff check) (pull_request) Successful in 26s
Details
CI / Type check (ty) (push) Successful in 29s
Details
CI / Format (ruff format) (pull_request) Successful in 33s
Details
CI / Sync project version with tag (pull_request) Has been skipped
Details
CI / Type check (ty) (pull_request) Successful in 35s
Details
CI / Tests (pull_request) Successful in 3m47s
Details
CI / Tests (push) Successful in 3m55s
Details

Browse Source 
Rejoins lines that only wrapped because they exceeded the old 88-char
limit; ruff check and the full test suite (725 passed) are unaffected.
This commit is contained in:
lars
2026-08-12 13:33:09 +02:00
parent 9ce7b32324
commit 55332db67a
66 changed files with 757 additions and 2413 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/sample.py
+12 -36
View File
@@ -67,9 +67,7 @@ def sample_ddpm(
alpha = schedule.alphas[i]
alpha_bar = schedule.alpha_bars[i]
z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
x = (1.0 / alpha.sqrt()) * (
x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
) + beta.sqrt() * z
x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
@@ -173,9 +171,7 @@ def _decode_stage2_flat(
else:
sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
1
)
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
return sec_cont, sec_type, sec_valid
@@ -206,9 +202,7 @@ def sample_secondaries(
v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
x = x + v * dt
return _decode_stage2_flat(
sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
)
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
@torch.no_grad()
@@ -226,9 +220,7 @@ def sample_secondaries_wgan(
B = cond_cont.size(0)
z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
return _decode_stage2_flat(
sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
)
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
@torch.no_grad()
@@ -306,12 +298,8 @@ def sample_secondaries_ar(
has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
history_feat = prev_repr.unsqueeze(1) # (B, 1, CONT_SLOT_DIM + type_dim)
remaining_frac = remaining.unsqueeze(1) # (B, 1)
slot_idx = torch.full(
(B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
)
hist, history_cache = sec_decoder.history_step(
history_feat, has_prev, history_cache
)
slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
if generator == "wgan":
z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
@@ -366,21 +354,15 @@ def sample_secondaries_ar(
sec_type[:, k] = type_k
if target == "onehot":
type_for_history = F.one_hot(
type_k.argmax(dim=-1), num_classes=type_dim
).float()
type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
else:
type_for_history = type_k
stick_fraction = torch.sigmoid(cont_k[:, 0])
prev_repr = torch.cat(
[stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1
)
prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
1
)
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
return sec_cont, sec_type, sec_valid
@@ -426,16 +408,10 @@ def sample_stage2(
case, so there's nothing to dispatch to here.
"""
if isinstance(sec_decoder, Stage2Autoregressive):
return sample_secondaries_ar(
sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
)
return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
if sec_decoder.generator_kind == "wgan":
return sample_secondaries_wgan(
sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
)
return sample_secondaries(
sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
)
return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
def resolve_n_sec(