Logic/mamba_diffusion
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(as_mamba): Remove dt prediction and use fixed dt

Browse Source 
Removes the `dt_head` network and associated configuration parameters
(dt_min, dt_max, lambda_nfe, warmup_epochs). Replaces predicted time
steps with a fixed value derived from sequence length. Eliminates
the warmup phase and NFE loss calculation.
This commit is contained in:
gameloader
2026-01-21 13:07:36 +08:00
parent c58a73ae26
commit 1446f97459
2 changed files with 21 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
main.py
View File

@@ -6,7 +6,6 @@ from as_mamba import TrainConfig, run_training_and_plot
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Train AS-Mamba on sphere-to-sphere flow.")
parser.add_argument("--epochs", type=int, default=None)
parser.add_argument("--warmup-epochs", type=int, default=None)
parser.add_argument("--batch-size", type=int, default=None)
parser.add_argument("--steps-per-epoch", type=int, default=None)
parser.add_argument("--seq-len", type=int, default=None)