diff --git a/CONTRIBUTING_ch.md b/CONTRIBUTING_ch.md
new file mode 100644
index 00000000..1563c310
--- /dev/null
+++ b/CONTRIBUTING_ch.md
@@ -0,0 +1,407 @@
+# Contributing to PaddleMaterials
+
+感谢你对 PaddleMaterials 的关注与贡献！
+
+本文面向线上仓库 [PaddlePaddle/PaddleMaterials](https://github.com/PaddlePaddle/PaddleMaterials) 的贡献。PaddleMaterials 当前是基于 PaddlePaddle 的 AI4Materials 工具库，核心 Python 包为 `ppmat`，模型权重通常使用 `.pdparams`。
+
+## 1. 开发原则
+
+提交到 PaddleMaterials 的代码不仅需要能运行，还需要满足以下项目要求：
+
+- 符合仓库现有任务划分与目录结构。
+- 符合 repo 已有代码逻辑、配置风格和调用方式。
+- 按 README 可以完成主要训练、评测、推理或采样流程。
+- 文档、配置、代码实现和复现结果保持一致。
+- 优先保证可维护、可复现、可扩展。
+- 不明显破坏已有模型、数据集和任务入口。
+
+对于模型复现类 PR，本项目更关注：
+
+- 是否接入到了正确的任务目录。
+- 是否沿用了 PaddleMaterials 现有的数据流、配置流和入口脚本。
+- 是否复用了 `ppmat` 中已有的 model、dataset、trainer、metric、scheduler、predictor、sampler 等公共能力。
+- 是否补齐了 README、配置、测试或最小验证说明。
+- 是否与现有模块风格一致，而不是简单搬运外部仓库代码。
+
+## 2. 仓库结构与放置规范
+
+请严格按照仓库现有结构放置代码，不要随意新建平行体系。
+
+当前主要任务目录包括：
+
+- `property_prediction/`：性质预测任务（PP）。
+- `structure_generation/`：结构生成任务（SG）。
+- `interatomic_potentials/`：机器学习原子间势函数任务（MLIP）。
+- `electronic_structure/`：电子结构任务（ES）。
+- `spectrum_elucidation/`：谱图解析任务（SE）。
+- `spectrum_enhancement/`：谱图或显微图像增强任务（SPEN）。
+- `ppmatSim/`：基于模型或势函数的材料模拟流程。
+- `research/`：研究性流程或实验性工作。
+- `jointContribution/`：联合贡献或外部协同代码区域。该目录当前被 pre-commit 排除，新主线能力不建议默认放在这里，除非维护者明确要求。
+
+公共实现主要放在：
+
+- `ppmat/models/`：模型、图构建器和模型构建入口。
+- `ppmat/datasets/`：数据集、数据构建、collate、transform 和 dataloader 构建入口。
+- `ppmat/trainer/`：训练流程和训练状态。
+- `ppmat/losses/`、`ppmat/metrics/`、`ppmat/optimizer/`、`ppmat/schedulers/`：损失、指标、优化器和调度器。
+- `ppmat/predictor/`、`ppmat/sampler/`：公共推理和采样封装。
+- `ppmat/calculator/`：ASE 等外部接口。
+- `ppmat/utils/`：通用工具。
+- `test/`：项目测试或验证脚本。
+- `docs/`：项目文档。
+
+## 3. 新模型接入
+
+新增模型时，一般应至少涉及以下位置中的一部分：
+
+- `ppmat/models/<model_name>/`：模型主体、图构建器或模型内部工具。
+- `ppmat/models/__init__.py`：导入模型类，补充 `__all__`，有预训练权重时再补充 `MODEL_REGISTRY`。
+- 对应任务目录下的 `configs/<model_name>/`：模型配置与 README。
+- 对应任务目录下的 `train.py`、`predict.py`、`sample.py` 或相关入口适配。
+- 必要时补充 `test/` 下测试或最小验证脚本。
+
+请不要把完整模型逻辑直接堆在任务脚本里，也不要绕开 `ppmat` 现有抽象单独再造一套训练框架。
+
+新增模型时请注意：
+
+- 模型主体应使用 PaddlePaddle API，通常继承 `paddle.nn.Layer`。
+- 配置中继续使用 `__class_name__` 和 `__init_params__` 约定。
+- 优先复用现有 `Trainer`、`Dataset`、`Optimizer`、`Metric`、`Predict`、`Sample` 等配置结构。
+- 预训练权重命名与现有习惯保持一致，常见后缀为 `.pdparams`，例如 `best.pdparams`、`latest.pdparams`。
+- 当前仓库已经存在 `ppmat/predictor/` 和 `ppmat/sampler/` 公共封装。新增推理器或采样器时，应优先复用或扩展公共实现；确需任务特化时，应对齐对应任务目录已有入口风格。
+- 如果修改公共 predictor 或 sampler，需要验证相关已有模型不被破坏。
+
+来自其他框架或其他仓库的代码不能简单复制后直接提交，必须完成 PaddleMaterials 风格适配，包括：
+
+- PaddlePaddle API 适配。
+- 参数组织方式统一。
+- 配置命名方式统一。
+- 训练、推理、采样接口统一。
+- 文档风格统一。
+- 依赖和许可证说明清楚。
+
+## 4. 新数据集接入
+
+新增数据集一般应放在 `ppmat/datasets/`。如果数据处理逻辑较复杂，可以在该目录下新建子目录组织实现。
+
+数据集实现应尽量复用现有 dataset 风格，保持：
+
+- 初始化参数风格一致。
+- 数据划分逻辑清晰。
+- 字段命名清楚。
+- 返回值结构稳定。
+- 与现有 dataloader、collate、transform 兼容。
+- 与下游模型所需字段匹配。
+- 不在 dataset 内写过多任务特判逻辑。
+- 不把预处理、训练逻辑和 dataset 强耦合。
+
+数据集接入通常需要：
+
+- 在 `ppmat/datasets/` 下新增数据集类。
+- 在 `ppmat/datasets/__init__.py` 中导入数据集类，使现有 `build_dataloader` 能通过配置中的 `__class_name__` 找到它。
+- 必要时补充 dataset info、transform、collate 或 graph build 逻辑。
+- 如已有工厂函数、工具类或数据处理逻辑可以复用，请优先使用现有封装。
+
+数据集适配类 PR 应优先参考仓库已有数据集实现，而不是直接照搬外部仓库写法。
+
+数据集 README 至少说明：
+
+- 数据集名称。
+- 来源链接和许可限制。
+- 原始格式。
+- 下载方式。
+- 预处理方式。
+- 划分方式。
+- 标签含义、单位和字段名。
+- 与配置中路径、缓存路径的对应关系。
+
+数据集适配 PR 至少应满足：
+
+- 数据集类可实例化。
+- 数据能够被正常读取。
+- 样本字段与下游模型匹配。
+- 按 README 能跑通基本流程。
+- 文档与代码一致。
+
+## 5. 工具脚本与转换脚本
+
+`test/`、`tools/` 或任务目录下的工具脚本，只应放置：
+
+- 用户需要明确调用的正式工具。
+- 权重转换、数据转换、评测、结果对齐等确有长期维护意义的脚本。
+- 有 README 或命令说明的迁移工具。
+
+不应提交：
+
+- 临时调试脚本。
+- 只在作者本地环境下可用的脚本。
+- 含硬编码个人路径、私有机器名或私有凭据的实验脚本。
+- 一次性对齐脚本但没有说明文档的文件。
+- 大日志、缓存、训练输出和无说明的中间文件。
+
+## 6. PaddleMaterials 特定要求
+
+### 6.1 必须沿用仓库已有入口与组织方式
+
+新增内容应尽量接入现有流程，而不是另起炉灶。包括但不限于：
+
+- 训练入口。
+- 推理入口。
+- 采样入口。
+- 配置组织方式。
+- 数据集注册和调用方式。
+- 模型调用接口。
+- 权重下载和加载方式。
+
+如果现有任务目录已经存在同类模型或同类任务，请优先对齐该目录下已有实现。
+
+例如：
+
+- 晶体或分子性质预测任务，优先参考 `property_prediction/` 下已有模型与配置组织。
+- 机器学习原子间势函数任务，优先参考 `interatomic_potentials/` 下已有模型与 README 组织。
+- 结构生成任务，优先参考 `structure_generation/` 下 MatterGen、DiffCSP 的配置和采样入口。
+- 谱图解析任务，优先参考 `spectrum_elucidation/` 下 DiffNMR 的训练和采样流程。
+- 谱图或显微图像增强任务，优先参考 `spectrum_enhancement/` 下 SFIN 的配置和预测入口。
+- 数据集适配，优先参考 `ppmat/datasets/` 中已有实现。
+
+### 6.2 不接受仓库外风格强行迁入
+
+来自其他框架或其他仓库的代码，不能简单复制后直接提交。
+
+必须完成 PaddleMaterials 风格适配，包括：
+
+- PaddlePaddle API 适配，不在主线功能中强依赖其他深度学习框架。
+- 参数组织方式统一。
+- 配置命名方式统一。
+- 训练、推理、采样接口统一。
+- 文档风格统一。
+- 许可证、引用来源和第三方代码边界说明清楚。
+
+### 6.3 使用 Paddle 官方正式版本
+
+- 必须使用 PaddlePaddle 官方正式发布版本。
+- 不得依赖 Paddle develop 版本作为唯一可运行环境。
+- 若对 Paddle 版本有最低要求，需在 README 中明确写明。
+- 若依赖特定 CUDA、Python、paddle_scatter、pgl、pymatgen、ase 等版本，也应一并说明。
+
+对于 PR 验收，若代码只能在 Paddle develop 版本或私有环境下运行，一般不视为满足合入条件。
+
+## 7. README 与文档要求
+
+PaddleMaterials 的贡献，尤其是模型复现类和数据集适配类 PR，必须保证 README 可用。
+
+README 至少应包含：
+
+- 任务简介。
+- 模型或数据集简介。
+- 数据准备方式。
+- 环境依赖和已验证版本。
+- 训练命令。
+- 评测命令。
+- 推理或采样命令，如适用。
+- 关键配置说明。
+- 参考结果。
+- 预训练权重或数据下载链接，如适用。
+- 参考论文、官方实现或数据来源链接。
+
+README 必须满足：
+
+- 按照 README 操作，可以完成主要流程。
+- 文档与代码一致。
+- 命令能直接从仓库根目录复制运行，避免隐式前置条件。
+- 不省略关键步骤。
+- 不允许 README 写一套、代码实现另一套。
+
+对于模型复现类 PR，建议在 README 中写清：
+
+- 复现目标。
+- 所用数据集版本。
+- 关键超参数。
+- 训练轮数、batch size、学习率。
+- 硬件环境。
+- 评测指标。
+- 与原论文或参考实现的结果对比。
+- 偏差范围与原因说明，如有。
+
+若暂时无法完整对齐原论文结果，也应明确说明当前完成到什么程度，而不是模糊表述“支持该模型”。
+
+## 8. 配置文件要求
+
+配置文件应与 PaddleMaterials 现有 YAML 风格保持一致。
+
+当前常见顶层字段包括：
+
+- `Global`
+- `Trainer`
+- `Model`
+- `Optimizer`
+- `Metric`
+- `Dataset`
+- `Predict`
+- `Sample` 或 `Sampler`
+- `Loss`，如任务需要
+
+配置内容应清晰完整，至少建议包含：
+
+- 模型构建参数。
+- 数据集路径、字段、transform、graph converter 或预处理参数。
+- dataloader、sampler、batch size、num workers。
+- optimizer 和 lr scheduler。
+- loss 和 metric。
+- trainer、predictor 或 sampler 相关设置。
+- save、log、eval 相关设置。
+- pretrained model 或 checkpoint 设置，如适用。
+
+配置默认应尽量可运行。提交的配置不要求都是大规模训练配置，但至少应保证：
+
+- 能正常解析。
+- 能构建模型。
+- 能构建数据集，或明确说明数据下载与准备方式。
+- 能启动最小训练、评测、推理或采样流程。
+
+命名建议体现任务、模型名、数据集名和训练目标，例如：
+
+- `megnet_mp2018_train_60k_e_form.yaml`
+- `chgnet_oc20_s2ef_energy.yaml`
+- `dimenet++_mp2018_train_60k_band_gap.yaml`
+- `diffcsp_mp20.yaml`
+- `mattergen_alex_mp20_dft_band_gap.yaml`
+- `infgcn_qm9.yaml`
+- `sfin_haadf_enhance.yaml`
+
+避免使用难以理解的命名，例如：
+
+- `final_new.yaml`
+- `test2.yaml`
+- `run_ok.yaml`
+- `debug.yaml`
+
+## 9. 数据集文件和预训练模型权重
+
+不建议把大体积数据集、预训练权重、缓存文件或训练输出直接提交到 git。
+
+数据集或模型适配 PR 应提供验收所需的数据集文件或预训练模型权重。可以先将文件上传到网盘，并在 PR 中提供链接，或按项目约定 @ reviewer 协助上传到 BCE。
+
+维护者或 reviewer 将数据集文件和预训练模型权重上传到 BCE 后，贡献者需要将正式链接添加到对应位置，例如：
+
+- README 的下载说明。
+- 配置中的默认路径或说明。
+- `ppmat/models/__init__.py` 中的 `MODEL_REGISTRY`，如属于预训练模型。
+
+请同时说明：
+
+- 文件名和大小。
+- 解压后的目录结构。
+- checksum，建议提供。
+- 数据或权重来源。
+- 许可限制。
+
+## 10. 测试与验收要求
+
+PR 合入前，至少应满足核心功能的最小验证：
+
+- 代码可正常导入。
+- 关键配置可解析。
+- 模型能完成最小前向。
+- 数据集能正常加载。
+- 主要训练、评测、推理或采样流程能跑通。
+- 不明显破坏现有功能。
+
+建议为新增内容补充以下测试之一或多项：
+
+- dataset smoke test。
+- model forward test。
+- config load test。
+- train/eval/predict/sample smoke test。
+- 关键模块单测。
+- 公共 builder、transform、collate、metric 或 scheduler 的回归测试。
+
+模型复现类 PR 不强求一次性达到完全工程化，但至少要做到：
+
+- 复现目标明确。
+- 结果有基本支撑。
+- README 可跑通。
+- 配置与代码匹配。
+- 主要指标可复现或可解释。
+
+若仅完成“代码移植”而没有基本跑通与说明，不建议直接合入。
+
+## 11. 代码风格与格式检查
+
+本仓库使用 pre-commit 进行基础格式检查，当前包括 isort、black、ruff、Markdown CRLF/tab 检查、YAML/JSON 检查和 C/CUDA clang-format。
+
+提交前建议执行：
+
+```bash
+pre-commit run --all-files
+```
+
+如新增或修改了测试，也建议执行对应测试命令，例如：
+
+```bash
+python -m pytest test
+```
+
+如果无法完整运行测试，请在 PR 描述中说明未运行项、原因和替代验证方式。
+
+基本格式要求：
+
+- Python 代码风格与现有仓库保持一致。
+- YAML 和 JSON 格式合法。
+- Markdown 不应包含 CRLF 和 tab。
+- C/CUDA/C++/proto 代码需通过 clang-format。
+- 不保留无关 debug 输出。
+- 不提交冲突标记、私钥、token、无关临时文件。
+- 不提交 `__pycache__`、`.pytest_cache`、训练输出、大日志或本地缓存。
+
+以下内容通常不建议直接合入 PaddleMaterials 主仓库：
+
+- 与仓库主任务方向无关的代码。
+- 无 README、无说明、无最小验证的模型移植。
+- 含个人本地绝对路径的代码。
+- 强依赖私有环境的脚本。
+- 临时实验文件。
+- 大量未说明来源的权重或数据。
+- 文档与代码明显不一致的实现。
+- 只能在 Paddle develop 版本运行的功能。
+- 完全绕开现有项目结构、另写一套训练框架的实现。
+
+## 12. 提交前检查清单
+
+提交 PR 前，请至少确认以下内容：
+
+- 修改内容属于 PaddleMaterials 现有任务方向，或已与维护者讨论。
+- 目录位置正确。
+- 沿用了现有代码组织、配置方式和调用入口。
+- 使用 PaddlePaddle 官方正式发布版本。
+- 新增依赖合理，且版本和环境要求已说明。
+- README 已补充且可用。
+- 训练、评测、推理或采样主要流程能说明白。
+- 配置文件能正常解析。
+- 数据集或模型最小流程已验证。
+- 已执行格式检查，或说明无法执行的原因。
+- 未提交无关临时文件、大日志、缓存文件。
+- 未提交无说明的大体积权重文件或数据文件。
+- PR 描述完整，包含动机、修改内容、验证方式和已知限制。
+- 如贡献属于指定活动，请按活动要求在对应 issue 或活动页同步进展，例如 `#194`。
+
+## 13. Review 关注重点
+
+Review 时，维护者通常会重点关注：
+
+- 是否符合 PaddleMaterials 当前任务方向。
+- 是否符合 repo 已有结构与逻辑。
+- 是否复用了公共抽象，而不是绕开主流程。
+- README 是否可支撑完整使用。
+- 文档、配置和代码是否一致。
+- 是否具备最小可复现性。
+- 是否容易长期维护。
+- 是否对现有功能造成副作用。
+- 数据、权重和外部代码来源是否清晰合规。
+
+请理解，review 的目标不是增加门槛，而是保证 PaddleMaterials 作为统一 AI4Materials 工具库的整体质量。
+
+## 14. 致谢
+
+感谢每一位贡献者对 PaddleMaterials 的支持。期待大家共同完善 PaddleMaterials 的模型、数据集、训练评测、推理采样和材料模拟能力。
diff --git a/CONTRIBUTING_en.md b/CONTRIBUTING_en.md
new file mode 100644
index 00000000..0bc46607
--- /dev/null
+++ b/CONTRIBUTING_en.md
@@ -0,0 +1,407 @@
+# Contributing to PaddleMaterials
+
+Thank you for your interest in contributing to PaddleMaterials!
+
+This guide applies to the online repository [PaddlePaddle/PaddleMaterials](https://github.com/PaddlePaddle/PaddleMaterials). PaddleMaterials is an AI4Materials toolkit built on PaddlePaddle. Its core Python package is `ppmat`, and pretrained model weights commonly use the `.pdparams` suffix.
+
+## 1. Development Principles
+
+Code submitted to PaddleMaterials should do more than run successfully. It should also meet the following project expectations:
+
+- Follow the existing task split and repository structure.
+- Match the existing code logic, configuration style, and invocation patterns.
+- Allow the main training, evaluation, inference, or sampling workflow to run by following the README.
+- Keep documentation, configuration files, implementation, and reproduced results consistent.
+- Prioritize maintainability, reproducibility, and extensibility.
+- Avoid breaking existing models, datasets, and task entrypoints.
+
+For model reproduction PRs, maintainers will pay particular attention to:
+
+- Whether the contribution is placed under the correct task directory.
+- Whether it reuses the existing PaddleMaterials data flow, configuration flow, and entry scripts.
+- Whether it reuses existing `ppmat` abstractions such as models, datasets, trainers, metrics, schedulers, predictors, and samplers.
+- Whether the README, configuration files, tests, or minimum verification instructions are complete.
+- Whether the implementation follows the style of existing modules instead of directly copying code from an external repository.
+
+## 2. Repository Structure
+
+Please place code according to the existing repository structure. Do not introduce a separate parallel system unless maintainers have agreed to it.
+
+Current task directories include:
+
+- `property_prediction/`: property prediction tasks (PP).
+- `structure_generation/`: structure generation tasks (SG).
+- `interatomic_potentials/`: machine learning interatomic potential tasks (MLIP).
+- `electronic_structure/`: electronic structure tasks (ES).
+- `spectrum_elucidation/`: spectrum elucidation tasks (SE).
+- `spectrum_enhancement/`: spectrum or microscopy image enhancement tasks (SPEN).
+- `ppmatSim/`: material simulation workflows based on trained models or interatomic potentials.
+- `research/`: research-oriented or experimental workflows.
+- `jointContribution/`: joint contribution or external collaboration code. This directory is currently excluded from pre-commit checks. New mainline functionality should not be placed here by default unless maintainers explicitly request it.
+
+Shared implementation should generally live in:
+
+- `ppmat/models/`: model definitions, graph converters, and model builders.
+- `ppmat/datasets/`: datasets, data builders, collate functions, transforms, and dataloader builders.
+- `ppmat/trainer/`: training logic and trainer state.
+- `ppmat/losses/`, `ppmat/metrics/`, `ppmat/optimizer/`, `ppmat/schedulers/`: losses, metrics, optimizers, and schedulers.
+- `ppmat/predictor/`, `ppmat/sampler/`: shared inference and sampling wrappers.
+- `ppmat/calculator/`: external interfaces such as ASE integration.
+- `ppmat/utils/`: common utilities.
+- `test/`: project tests or validation scripts.
+- `docs/`: project documentation.
+
+## 3. Adding a New Model
+
+Adding a new model usually involves some of the following locations:
+
+- `ppmat/models/<model_name>/`: the model implementation, graph converter, or model-specific utilities.
+- `ppmat/models/__init__.py`: import the model class, update `__all__`, and add `MODEL_REGISTRY` entries if pretrained weights are provided.
+- `configs/<model_name>/` under the corresponding task directory: model configurations and README.
+- `train.py`, `predict.py`, `sample.py`, or related entrypoints under the corresponding task directory when adaptation is necessary.
+- `test/`: tests or minimum verification scripts when needed.
+
+Do not put the full model implementation directly into task scripts. Do not bypass existing `ppmat` abstractions by creating a separate training framework.
+
+When adding a model:
+
+- Use PaddlePaddle APIs. Model classes usually inherit from `paddle.nn.Layer`.
+- Continue using the `__class_name__` and `__init_params__` convention in configuration files.
+- Reuse existing `Trainer`, `Dataset`, `Optimizer`, `Metric`, `Predict`, and `Sample` configuration structures whenever possible.
+- Follow existing pretrained weight naming conventions. Common names include `best.pdparams` and `latest.pdparams`.
+- The repository already contains shared `ppmat/predictor/` and `ppmat/sampler/` wrappers. Prefer reusing or extending these shared implementations. If task-specific inference or sampling code is necessary, keep the task entrypoint lightweight and consistent with nearby examples.
+- If you modify shared predictor or sampler logic, verify that related existing models still work.
+
+Code from other frameworks or repositories must not be copied into PaddleMaterials without adaptation. It must be converted to the PaddleMaterials style, including:
+
+- PaddlePaddle API usage.
+- Consistent parameter organization.
+- Consistent configuration naming.
+- Consistent training, inference, and sampling interfaces.
+- Consistent documentation style.
+- Clear dependency and license information.
+
+## 4. Adding a New Dataset
+
+New datasets should generally be added under `ppmat/datasets/`. If the processing logic is complex, create a subdirectory under `ppmat/datasets/` to organize it.
+
+Dataset implementations should follow the style of existing datasets:
+
+- Keep initialization parameters consistent with existing datasets.
+- Make train, validation, and test split logic clear.
+- Use clear field names.
+- Return a stable sample structure.
+- Stay compatible with existing dataloaders, collate functions, and transforms.
+- Match the fields expected by downstream models.
+- Avoid excessive task-specific branching inside the dataset class.
+- Avoid tightly coupling preprocessing, training logic, and dataset logic.
+
+Dataset integration usually requires:
+
+- Adding the dataset class under `ppmat/datasets/`.
+- Importing the dataset class in `ppmat/datasets/__init__.py` so that `build_dataloader` can find it through `__class_name__` in the config.
+- Adding dataset info, transforms, collate logic, or graph building logic when necessary.
+- Reusing existing factory functions, utility classes, and data processing logic whenever possible.
+
+Dataset PRs should use existing PaddleMaterials datasets as references instead of directly copying the style of an external repository.
+
+The dataset README should describe at least:
+
+- Dataset name.
+- Source link and license restrictions.
+- Original format.
+- Download method.
+- Preprocessing steps.
+- Split strategy.
+- Label meanings, units, and field names.
+- How paths and cache paths in the configuration correspond to the prepared files.
+
+At minimum, a dataset PR should show that:
+
+- The dataset class can be instantiated.
+- Data can be read successfully.
+- Sample fields match downstream model requirements.
+- The basic workflow runs by following the README.
+- Documentation and code are consistent.
+
+## 5. Utility and Conversion Scripts
+
+Scripts under `test/`, `tools/`, or task directories should be limited to:
+
+- Official tools that users are expected to invoke.
+- Long-term-maintained scripts for weight conversion, data conversion, evaluation, or result alignment.
+- Migration tools with README instructions or command examples.
+
+Do not submit:
+
+- Temporary debugging scripts.
+- Scripts that only work in the author's local environment.
+- Experimental scripts with hard-coded personal paths, private machine names, or credentials.
+- One-off alignment scripts without documentation.
+- Large logs, caches, training outputs, or unexplained intermediate files.
+
+## 6. PaddleMaterials-Specific Requirements
+
+### 6.1 Reuse Existing Entrypoints and Organization
+
+New functionality should be integrated into existing workflows instead of introducing a new standalone system. This includes:
+
+- Training entrypoints.
+- Inference entrypoints.
+- Sampling entrypoints.
+- Configuration organization.
+- Dataset registration and invocation.
+- Model invocation interfaces.
+- Weight download and loading logic.
+
+If a similar model or task already exists, align with the implementation in the corresponding task directory.
+
+For example:
+
+- For crystal or molecular property prediction, refer to model and config organization under `property_prediction/`.
+- For machine learning interatomic potentials, refer to model and README organization under `interatomic_potentials/`.
+- For structure generation, refer to MatterGen and DiffCSP configurations and sampling entrypoints under `structure_generation/`.
+- For spectrum elucidation, refer to DiffNMR training and sampling workflows under `spectrum_elucidation/`.
+- For spectrum or microscopy image enhancement, refer to SFIN configurations and prediction entrypoints under `spectrum_enhancement/`.
+- For dataset adaptation, refer to existing implementations under `ppmat/datasets/`.
+
+### 6.2 Do Not Force External Repository Style Into PaddleMaterials
+
+Code from other frameworks or repositories cannot be submitted as a direct copy.
+
+It must be adapted to the PaddleMaterials style, including:
+
+- PaddlePaddle API adaptation. Mainline functionality should not depend on another deep learning framework.
+- Consistent parameter organization.
+- Consistent configuration naming.
+- Consistent training, inference, and sampling interfaces.
+- Consistent documentation style.
+- Clear license, citation, and third-party code boundary information.
+
+### 6.3 Use Official Paddle Releases
+
+- Use officially released PaddlePaddle versions.
+- Do not make the Paddle develop branch the only supported runtime.
+- If a minimum PaddlePaddle version is required, state it clearly in the README.
+- If specific CUDA, Python, `paddle_scatter`, `pgl`, `pymatgen`, `ase`, or other dependency versions are required, document them as well.
+
+For PR acceptance, functionality that only runs on the Paddle develop branch or a private environment is generally not considered ready to merge.
+
+## 7. README and Documentation Requirements
+
+Contributions to PaddleMaterials, especially model reproduction and dataset adaptation PRs, must include usable documentation.
+
+The README should include at least:
+
+- Task introduction.
+- Model or dataset introduction.
+- Data preparation instructions.
+- Environment dependencies and verified versions.
+- Training command.
+- Evaluation command.
+- Inference or sampling command, if applicable.
+- Key configuration explanation.
+- Reference results.
+- Pretrained weight or dataset download links, if applicable.
+- Reference paper, official implementation, or data source links.
+
+The README must satisfy the following:
+
+- Following the README should allow users to complete the main workflow.
+- Documentation and code must be consistent.
+- Commands should be copyable from the repository root, with no hidden prerequisites.
+- Key steps must not be omitted.
+- The README must not describe one workflow while the code implements another.
+
+For model reproduction PRs, we recommend documenting:
+
+- Reproduction target.
+- Dataset version.
+- Key hyperparameters.
+- Number of epochs, batch size, and learning rate.
+- Hardware environment.
+- Evaluation metrics.
+- Comparison with the original paper or reference implementation.
+- Deviation range and possible reasons, if any.
+
+If the reproduced result does not fully match the original paper yet, clearly state the current status instead of vaguely saying that the model is supported.
+
+## 8. Configuration File Requirements
+
+Configuration files should follow the existing PaddleMaterials YAML style.
+
+Common top-level fields include:
+
+- `Global`
+- `Trainer`
+- `Model`
+- `Optimizer`
+- `Metric`
+- `Dataset`
+- `Predict`
+- `Sample` or `Sampler`
+- `Loss`, when required by the task
+
+Configurations should be clear and complete. They should generally include:
+
+- Model construction parameters.
+- Dataset paths, fields, transforms, graph converters, or preprocessing parameters.
+- Dataloader, sampler, batch size, and number of workers.
+- Optimizer and learning rate scheduler.
+- Loss and metrics.
+- Trainer, predictor, or sampler settings.
+- Save, log, and evaluation settings.
+- Pretrained model or checkpoint settings, if applicable.
+
+Default configurations should be runnable as much as possible. Submitted configurations do not all need to be large-scale training configs, but they should at least:
+
+- Parse successfully.
+- Build the model.
+- Build the dataset, or clearly describe how to download and prepare the data.
+- Start a minimal training, evaluation, inference, or sampling workflow.
+
+Configuration names should reflect the task, model name, dataset name, and training target. For example:
+
+- `megnet_mp2018_train_60k_e_form.yaml`
+- `chgnet_oc20_s2ef_energy.yaml`
+- `dimenet++_mp2018_train_60k_band_gap.yaml`
+- `diffcsp_mp20.yaml`
+- `mattergen_alex_mp20_dft_band_gap.yaml`
+- `infgcn_qm9.yaml`
+- `sfin_haadf_enhance.yaml`
+
+Avoid unclear or temporary names such as:
+
+- `final_new.yaml`
+- `test2.yaml`
+- `run_ok.yaml`
+- `debug.yaml`
+
+## 9. Dataset Files and Pretrained Weights
+
+Large datasets, pretrained weights, cache files, and training outputs should not be committed directly to git.
+
+Dataset or model adaptation PRs should provide the files required for review, such as dataset files or pretrained model weights. You may upload them to cloud storage first and provide the link in the PR, or follow the project convention and ask a reviewer to help upload them to BCE.
+
+After maintainers or reviewers upload dataset files or pretrained weights to BCE, contributors should update the official links in the corresponding locations, such as:
+
+- Download instructions in the README.
+- Default paths or notes in configuration files.
+- `MODEL_REGISTRY` in `ppmat/models/__init__.py`, if the file is a pretrained model.
+
+Please also provide:
+
+- File name and size.
+- Directory structure after extraction.
+- Checksum, recommended when available.
+- Data or weight source.
+- License restrictions.
+
+## 10. Testing and Acceptance Requirements
+
+Before a PR is merged, the core functionality should satisfy the minimum verification requirements:
+
+- Code imports successfully.
+- Key configurations parse successfully.
+- The model can complete a minimal forward pass.
+- The dataset can be loaded.
+- The main training, evaluation, inference, or sampling workflow can run.
+- Existing functionality is not obviously broken.
+
+We recommend adding one or more of the following tests for new functionality:
+
+- Dataset smoke test.
+- Model forward test.
+- Config load test.
+- Train, evaluation, prediction, or sampling smoke test.
+- Unit tests for key modules.
+- Regression tests for shared builders, transforms, collate functions, metrics, or schedulers.
+
+Model reproduction PRs do not need to be fully productionized in one step, but they should at least:
+
+- Have a clear reproduction target.
+- Provide basic result evidence.
+- Include a runnable README.
+- Keep configurations and code aligned.
+- Make the main metrics reproducible or explainable.
+
+Code migration alone, without a basic runnable workflow and explanation, is not recommended for merge.
+
+## 11. Code Style and Formatting
+
+This repository uses pre-commit for basic formatting checks. Current checks include isort, black, ruff, Markdown CRLF/tab checks, YAML/JSON checks, and clang-format for C/CUDA files.
+
+Before submitting, we recommend running:
+
+```bash
+pre-commit run --all-files
+```
+
+If you add or modify tests, also run the corresponding test command, for example:
+
+```bash
+python -m pytest test
+```
+
+If you cannot run the full test suite, explain in the PR description what was not run, why it was not run, and what alternative verification was performed.
+
+Basic formatting requirements:
+
+- Python code should follow the existing repository style.
+- YAML and JSON files must be valid.
+- Markdown files should not contain CRLF line endings or tab characters.
+- C/CUDA/C++/proto files should pass clang-format.
+- Do not keep unrelated debug output.
+- Do not commit conflict markers, private keys, tokens, or unrelated temporary files.
+- Do not commit `__pycache__`, `.pytest_cache`, training outputs, large logs, or local caches.
+
+The following content is generally not recommended for direct merge into the PaddleMaterials main repository:
+
+- Code unrelated to the main task directions of the repository.
+- Model migrations without README, explanation, or minimum verification.
+- Code containing personal absolute paths.
+- Scripts that strongly depend on private environments.
+- Temporary experiment files.
+- Large weights or datasets without source explanation.
+- Implementations where documentation and code are clearly inconsistent.
+- Features that only run on the Paddle develop branch.
+- Implementations that bypass the existing project structure and introduce a separate training framework.
+
+## 12. Pre-Submission Checklist
+
+Before submitting a PR, please confirm that:
+
+- The change belongs to an existing PaddleMaterials task direction, or has been discussed with maintainers.
+- Files are placed in the correct directories.
+- Existing code organization, configuration style, and invocation entrypoints are reused.
+- An officially released PaddlePaddle version is used.
+- New dependencies are reasonable, and version or environment requirements are documented.
+- The README has been updated and is usable.
+- The main training, evaluation, inference, or sampling workflow is clearly described.
+- Configuration files can be parsed successfully.
+- The minimum dataset or model workflow has been verified.
+- Formatting checks have been run, or the reason for not running them is explained.
+- No unrelated temporary files, large logs, or cache files are committed.
+- No large weight or data files are committed without explanation.
+- The PR description is complete, including motivation, changes, verification, and known limitations.
+- If the contribution belongs to a specific campaign or activity, update the corresponding issue or activity page as required, for example `#194`.
+
+## 13. Review Focus
+
+During review, maintainers usually focus on:
+
+- Whether the contribution fits the current PaddleMaterials task directions.
+- Whether it follows the existing repository structure and logic.
+- Whether it reuses shared abstractions instead of bypassing the main workflow.
+- Whether the README supports complete usage.
+- Whether documentation, configuration, and code are consistent.
+- Whether minimum reproducibility is provided.
+- Whether the implementation is maintainable long term.
+- Whether existing functionality is affected.
+- Whether sources of data, weights, and external code are clear and compliant.
+
+Review is not meant to raise unnecessary barriers. Its purpose is to preserve the quality of PaddleMaterials as a unified AI4Materials toolkit.
+
+## 14. Acknowledgements
+
+Thank you to every contributor who supports PaddleMaterials. We look forward to improving PaddleMaterials together across models, datasets, training and evaluation, inference and sampling, and material simulation workflows.
diff --git a/Install.md b/Install.md
index 8efb4733..c55cf10d 100644
--- a/Install.md
+++ b/Install.md
@@ -90,7 +90,7 @@ Generate crystal structures with four atoms using a pretrained MatterGen model:
 python structure_generation/sample.py \
     --model_name='mattergen_mp20' \
     --weights_name='latest.pdparams' \
-    --save_path='result_mattergen_mp20/' \
+    --output_dir='result_mattergen_mp20/' \
     --mode='by_num_atoms' \
     --num_atoms=4
 ```
@@ -113,11 +113,12 @@ Predict electron density using a trained InfGCN checkpoint:
 
 ```bash
 python electronic_structure/predict.py \
-    --config='electronic_structure/configs/infgcn/infgcn_qm9.yaml' \
-    --checkpoint='path/to/infgcn_qm9.pdparams' \
-    --split='validation' \
-    --index=0 \
-    --output_dir='output/infgcn_qm9/validation_0' \
+    --model_name='infgcn_qm9' \
+    --weights_name='best.pdparams' \
+    --mol_input='electronic_structure/configs/infgcn/example/methane.mol' \
+    --atom_file='electronic_structure/configs/qm9.json' \
+    --mol_grid_shape=8 \
+    --output_dir='output/infgcn_qm9/methane' \
     --save_pred_cube
 ```
 
@@ -130,9 +131,9 @@ Run NMR spectrum elucidation using a trained DiffNMR checkpoint:
 
 ```bash
 python spectrum_elucidation/sample.py \
-    --config_path='spectrum_elucidation/configs/diffnmr/DiffNMR.yaml' \
-    --checkpoint_path='path/to/DiffNMR_nless15_best.pdparams' \
-    --save_path='result_diffnmr_nless15/'
+    --model_name='diffnmr_msdnmr_nless15' \
+    --weights_name='best.pdparams' \
+    --output_dir='result_diffnmr_nless15/'
 ```
 
 ### 2.6 Spectrum Enhancement
@@ -142,7 +143,9 @@ Enhance STEM images using a pretrained SFIN model:
 ```bash
 python spectrum_enhancement/predict.py \
     --model_name='sfin_haadf_enhance' \
-    --split='val'
+    --weights_name='best.pdparams' \
+    --input_path='path/to/noisy_image.png' \
+    --output_dir='result_sfin/'
 ```
 
 For more usage instructions, refer to the task-specific README files or the
diff --git a/Install_cn.md b/Install_cn.md
index d12097a3..0c7594f9 100644
--- a/Install_cn.md
+++ b/Install_cn.md
@@ -89,7 +89,7 @@ python property_prediction/predict.py \
 python structure_generation/sample.py \
     --model_name='mattergen_mp20' \
     --weights_name='latest.pdparams' \
-    --save_path='result_mattergen_mp20/' \
+    --output_dir='result_mattergen_mp20/' \
     --mode='by_num_atoms' \
     --num_atoms=4
 ```
@@ -112,11 +112,12 @@ python interatomic_potentials/predict.py \
 
 ```bash
 python electronic_structure/predict.py \
-    --config='electronic_structure/configs/infgcn/infgcn_qm9.yaml' \
-    --checkpoint='path/to/infgcn_qm9.pdparams' \
-    --split='validation' \
-    --index=0 \
-    --output_dir='output/infgcn_qm9/validation_0' \
+    --model_name='infgcn_qm9' \
+    --weights_name='best.pdparams' \
+    --mol_input='electronic_structure/configs/infgcn/example/methane.mol' \
+    --atom_file='electronic_structure/configs/qm9.json' \
+    --mol_grid_shape=8 \
+    --output_dir='output/infgcn_qm9/methane' \
     --save_pred_cube
 ```
 
@@ -129,9 +130,9 @@ python electronic_structure/predict.py \
 
 ```bash
 python spectrum_elucidation/sample.py \
-    --config_path='spectrum_elucidation/configs/diffnmr/DiffNMR.yaml' \
-    --checkpoint_path='path/to/DiffNMR_nless15_best.pdparams' \
-    --save_path='result_diffnmr_nless15/'
+    --model_name='diffnmr_msdnmr_nless15' \
+    --weights_name='best.pdparams' \
+    --output_dir='result_diffnmr_nless15/'
 ```
 
 ### 2.6 谱图增强
@@ -141,7 +142,9 @@ python spectrum_elucidation/sample.py \
 ```bash
 python spectrum_enhancement/predict.py \
     --model_name='sfin_haadf_enhance' \
-    --split='val'
+    --weights_name='best.pdparams' \
+    --input_path='path/to/noisy_image.png' \
+    --output_dir='result_sfin/'
 ```
 
 更多使用说明请参考各任务 README 或 [Get Started](./get_started.md)。
diff --git a/README.md b/README.md
index 3f34cd0c..d20c43d1 100755
--- a/README.md
+++ b/README.md
@@ -19,6 +19,12 @@
   </a>
 </p>
 
+<p align="center">
+  <a href="README.md"><img src="https://img.shields.io/badge/English-blue" alt="English"></a>
+  <a href="README_zh.md"><img src="https://img.shields.io/badge/简体中文-lightgrey" alt="简体中文"></a>
+  <a href="README_ja.md"><img src="https://img.shields.io/badge/日本語-lightgrey" alt="日本語"></a>
+</p>
+
 ## 🚀 Introduction
 
 **PaddleMaterials** is an end-to-end AI4Materials toolkit built on the **PaddlePaddle** deep learning framework. Designed as a data-mechanism dual-driven platform for developing and deploying foundation models in materials science, **PPMat** enables researchers to efficiently build AI models and accelerate material discovery using pretrained models.
@@ -106,7 +112,7 @@ Generate novel crystal structures using a pretrained MatterGen model:
 python structure_generation/sample.py \
     --model_name='mattergen_mp20' \
     --weights_name='latest.pdparams' \
-    --save_path='result_mattergen_mp20/' \
+    --output_dir='result_mattergen_mp20/' \
     --mode='by_num_atoms' \
     --num_atoms=4
 ```
@@ -125,29 +131,33 @@ python interatomic_potentials/predict.py \
 
 #### Electronic Structure
 
-Predict electron density using a pretrained InfGCN checkpoint:
+Predict electron density from the bundled methane example using a pretrained InfGCN
+model:
 
 ```bash
 python electronic_structure/predict.py \
-    --config='electronic_structure/configs/infgcn/infgcn_qm9.yaml' \
-    --checkpoint='path/to/infgcn_qm9.pdparams' \
-    --split='validation' \
-    --index=0 \
-    --output_dir='output/infgcn_qm9/validation_0' \
+    --model_name='infgcn_qm9' \
+    --weights_name='best.pdparams' \
+    --mol_input='electronic_structure/configs/infgcn/example/methane.mol' \
+    --atom_file='electronic_structure/configs/qm9.json' \
+    --mol_grid_shape=8 \
+    --output_dir='output/infgcn_qm9/methane' \
     --save_pred_cube
 ```
 
-See the [InfGCN prediction guide](electronic_structure/configs/infgcn/README.md#prediction) for dataset and checkpoint preparation.
+See the [InfGCN prediction guide](electronic_structure/configs/infgcn/README.md#prediction)
+for dataset-based and local-checkpoint inference.
 
 #### Spectrum Elucidation
 
-Run NMR spectrum elucidation using a pretrained DiffNMR checkpoint:
+Run NMR spectrum elucidation using the bundled example and a pretrained DiffNMR
+model:
 
 ```bash
 python spectrum_elucidation/sample.py \
-    --config_path='spectrum_elucidation/configs/diffnmr/DiffNMR.yaml' \
-    --checkpoint_path='path/to/DiffNMR_nless15_best.pdparams' \
-    --save_path='result_diffnmr_nless15/'
+    --model_name='diffnmr_msdnmr_nless15' \
+    --weights_name='best.pdparams' \
+    --output_dir='result_diffnmr_nless15/'
 ```
 
 #### Spectrum Enhancement
@@ -157,7 +167,9 @@ Enhance STEM images using a pretrained SFIN model:
 ```bash
 python spectrum_enhancement/predict.py \
     --model_name='sfin_haadf_enhance' \
-    --split='val'
+    --weights_name='best.pdparams' \
+    --input_path='path/to/noisy_image.png' \
+    --output_dir='result_sfin/'
 ```
 
 ---
diff --git a/README_PYPI.md b/README_PYPI.md
index f6f4bb7b..5e26956d 100644
--- a/README_PYPI.md
+++ b/README_PYPI.md
@@ -108,7 +108,7 @@ Generate novel crystal structures using a pretrained MatterGen model:
 python structure_generation/sample.py \
     --model_name='mattergen_mp20' \
     --weights_name='latest.pdparams' \
-    --save_path='result_mattergen_mp20/' \
+    --output_dir='result_mattergen_mp20/' \
     --mode='by_num_atoms' \
     --num_atoms=4
 ```
@@ -127,29 +127,33 @@ python interatomic_potentials/predict.py \
 
 #### Electronic Structure
 
-Predict electron density using a pretrained InfGCN checkpoint:
+Predict electron density from the bundled methane example using a pretrained InfGCN
+model:
 
 ```bash
 python electronic_structure/predict.py \
-    --config='electronic_structure/configs/infgcn/infgcn_qm9.yaml' \
-    --checkpoint='path/to/infgcn_qm9.pdparams' \
-    --split='validation' \
-    --index=0 \
-    --output_dir='output/infgcn_qm9/validation_0' \
+    --model_name='infgcn_qm9' \
+    --weights_name='best.pdparams' \
+    --mol_input='electronic_structure/configs/infgcn/example/methane.mol' \
+    --atom_file='electronic_structure/configs/qm9.json' \
+    --mol_grid_shape=8 \
+    --output_dir='output/infgcn_qm9/methane' \
     --save_pred_cube
 ```
 
-See the [InfGCN prediction guide](https://github.com/PaddlePaddle/PaddleMaterials/blob/develop/electronic_structure/configs/infgcn/README.md#prediction) for dataset and checkpoint preparation.
+See the [InfGCN prediction guide](https://github.com/PaddlePaddle/PaddleMaterials/blob/develop/electronic_structure/configs/infgcn/README.md#prediction)
+for dataset-based and local-checkpoint inference.
 
 #### Spectrum Elucidation
 
-Run NMR spectrum elucidation using a pretrained DiffNMR checkpoint:
+Run NMR spectrum elucidation using the bundled example and a pretrained DiffNMR
+model:
 
 ```bash
 python spectrum_elucidation/sample.py \
-    --config_path='spectrum_elucidation/configs/diffnmr/DiffNMR.yaml' \
-    --checkpoint_path='path/to/DiffNMR_nless15_best.pdparams' \
-    --save_path='result_diffnmr_nless15/'
+    --model_name='diffnmr_msdnmr_nless15' \
+    --weights_name='best.pdparams' \
+    --output_dir='result_diffnmr_nless15/'
 ```
 
 #### Spectrum Enhancement
@@ -159,7 +163,9 @@ Enhance STEM images using a pretrained SFIN model:
 ```bash
 python spectrum_enhancement/predict.py \
     --model_name='sfin_haadf_enhance' \
-    --split='val'
+    --weights_name='best.pdparams' \
+    --input_path='path/to/noisy_image.png' \
+    --output_dir='result_sfin/'
 ```
 
 ---
diff --git a/README_ja.md b/README_ja.md
new file mode 100644
index 00000000..a8dbc9ce
--- /dev/null
+++ b/README_ja.md
@@ -0,0 +1,225 @@
+# PaddleMaterials
+
+<p align="center">
+  <a href="README.md"><img src="https://img.shields.io/badge/English-lightgrey" alt="English"></a>
+  <a href="README_zh.md"><img src="https://img.shields.io/badge/简体中文-lightgrey" alt="简体中文"></a>
+  <a href="README_ja.md"><img src="https://img.shields.io/badge/日本語-blue" alt="日本語"></a>
+</p>
+
+<p align="center"><img src="docs/ppmat_logo.png" alt="PaddleMaterials" width="400"></p>
+
+<p align="center">
+  <a href="Install.md"><img alt="Python 3.10+" src="https://img.shields.io/badge/Python-3.10%2B-3776AB?logo=python&amp;logoColor=white"></a>
+  <a href="https://pypi.org/project/ppmat/"><img alt="PyPI バージョン" src="https://img.shields.io/pypi/v/ppmat?logo=pypi&amp;logoColor=white"></a>
+  <a href="LICENSE"><img alt="Apache 2.0 ライセンス" src="https://img.shields.io/github/license/PaddlePaddle/PaddleMaterials"></a>
+  <a href="https://github.com/PaddlePaddle/PaddleMaterials/stargazers"><img alt="GitHub Stars" src="https://img.shields.io/github/stars/PaddlePaddle/PaddleMaterials?style=flat&amp;logo=github"></a>
+</p>
+
+## 🚀 はじめに
+
+**PaddleMaterials** は、**PaddlePaddle** ディープラーニングフレームワークを基盤とするエンドツーエンドの AI4Materials ツールキットです。材料科学における基盤モデルの開発とデプロイを目的とした、データとメカニズムの二重駆動型プラットフォームとして設計されています。**PPMat** により、研究者は AI モデルを効率的に構築し、事前学習済みモデルを利用して材料探索を加速できます。
+
+<p align="left"><img src="docs/overview_en.png" align="middle" width="1000"/></p>
+
+### 🧩 コア機能
+
+| タスク | 説明 | 主な用途 |
+|--------|------|----------|
+| **物性予測（PP）** | 構造から材料物性を予測 | 順設計、生成エネルギー・バンドギャップ・弾性率などの予測 |
+| **構造生成（SG）** | 新しい結晶構造を生成 | 逆設計または構造生成 |
+| **機械学習原子間ポテンシャル（MLIP）** | DFT の代理となる ML ポテンシャル | 分子動力学シミュレーション |
+| **電子構造（ES）** | DFT の代理モデルとして物理場を予測 | 電子密度の予測 |
+| **スペクトル解析（SE）** | スペクトルから構造を再構成 | NMR 構造解析 |
+| **スペクトル強調（SPEN）** | 顕微鏡画像およびスペクトル信号を強調 | STEM 画像の強調、ノイズ除去 |
+
+### 🧱 対応する材料
+
+- **無機結晶** — 複数のデータセットと事前学習済みモデルによる充実したサポート
+- **有機分子** — 低分子や一部のポリマーを含む複数のデータセットと事前学習済みモデルをサポート
+
+### ✨ PaddleMaterials を選ぶ理由
+
+- ✅ **豊富な事前学習済みモデルと AI-ready データセット** — 推論にすぐ利用できる 50 以上の事前学習済みモデルと、学習用に厳選された複数のデータセット
+- ✅ **マルチタスク統合** — PP、SG、MLIP、ES、SE、SPEN などのタスクを統一されたフレームワークでサポート
+- ✅ **マルチハードウェア対応** — NVIDIA GPU、MetaX GPU、Intel CPU を全面的にサポート
+- ✅ **プロダクション対応** — 標準化された設計により使いやすく、分散学習、混合精度、チェックポイントからの再開に対応
+
+### 📑 対応タスク
+
+| タスク | 説明 | リンク |
+|--------|------|--------|
+| **物性予測（PP）** | 生成エネルギー、バンドギャップ、弾性特性を予測 | [README](property_prediction/README.md) |
+| **構造生成（SG）** | 拡散モデルで新しい結晶構造を生成 | [README](structure_generation/README.md) |
+| **機械学習原子間ポテンシャル（MLIP）** | 分子動力学向けの DFT 精度ポテンシャル | [README](interatomic_potentials/README.md) |
+| **電子構造（ES）** | 電子構造特性を予測 | [README](electronic_structure/README.md) |
+| **スペクトル解析（SE）** | NMR スペクトルから分子構造を再構成 | [README](spectrum_elucidation/README.md) |
+| **スペクトル強調（SPEN）** | 顕微鏡画像とスペクトル信号を強調 | [README](spectrum_enhancement/README.md) |
+
+### 🤖 利用可能な事前学習済みモデル
+
+| タスク | モデル | データセット |
+|--------|--------|--------------|
+| **物性予測** | MEGNet、iComformer、DimeNet++、SphereNet | MP2018、MP2024、JARVIS、QM9 など |
+| **構造生成** | MatterGen、DiffCSP | MP20、ALEX など |
+| **機械学習原子間ポテンシャル** | CHGNet、MatterSim、SphereNet | MPTRJ、MD17 など |
+| **電子構造** | InfGCN | QM9_ES、MP_ES、OMol25_MC_ES など |
+| **スペクトル解析** | DiffNMR | MSD_NMR など |
+| **スペクトル強調** | SFIN | SFIN-HAADF/BF など |
+
+モデルの完全な一覧は [MODEL_REGISTRY](ppmat/models/__init__.py#L75) を参照してください。
+
+---
+
+## 🚀 クイックスタート
+
+### 🔧 インストール
+
+お使いのハードウェア環境に応じて、[インストールガイド](Install.md)を参照してください。マルチハードウェア対応の詳細については、[対応ハードウェア一覧](./docs/multi_device.md)を参照してください。
+
+---
+
+### ⚡ 簡単な推論
+
+#### 物性予測
+
+事前学習済み MEGNet モデルを使用して材料の生成エネルギーを予測します。
+
+```bash
+python property_prediction/predict.py \
+    --model_name='megnet_mp2018_train_60k_e_form' \
+    --weights_name='best.pdparams' \
+    --cif_file_path='./property_prediction/example_data/cifs/' \
+    --save_path='result.csv'
+```
+
+#### 構造生成
+
+事前学習済み MatterGen モデルを使用して新しい結晶構造を生成します。
+
+```bash
+python structure_generation/sample.py \
+    --model_name='mattergen_mp20' \
+    --weights_name='latest.pdparams' \
+    --output_dir='result_mattergen_mp20/' \
+    --mode='by_num_atoms' \
+    --num_atoms=4
+```
+
+#### 原子間ポテンシャル
+
+事前学習済み MatterSim モデルを使用してエネルギーと力を予測します。
+
+```bash
+python interatomic_potentials/predict.py \
+    --model_name='mattersim_1M' \
+    --weights_name='mattersim-v1.0.0-1M_model.pdparams' \
+    --cif_file_path='./interatomic_potentials/example_data/cifs/' \
+    --save_path='result.csv'
+```
+
+#### 電子構造
+
+事前学習済み InfGCN モデルを使用して、同梱のメタン例の電子密度を予測します。
+
+```bash
+python electronic_structure/predict.py \
+    --model_name='infgcn_qm9' \
+    --weights_name='best.pdparams' \
+    --mol_input='electronic_structure/configs/infgcn/example/methane.mol' \
+    --atom_file='electronic_structure/configs/qm9.json' \
+    --mol_grid_shape=8 \
+    --output_dir='output/infgcn_qm9/methane' \
+    --save_pred_cube
+```
+
+データセットまたはローカルチェックポイントを使用した推論については、
+[InfGCN 予測ガイド](electronic_structure/configs/infgcn/README.md#prediction)を参照してください。
+
+#### スペクトル解析
+
+事前学習済み DiffNMR モデルと同梱の例を使用して NMR スペクトル解析を実行します。
+
+```bash
+python spectrum_elucidation/sample.py \
+    --model_name='diffnmr_msdnmr_nless15' \
+    --weights_name='best.pdparams' \
+    --output_dir='result_diffnmr_nless15/'
+```
+
+#### スペクトル強調
+
+事前学習済み SFIN モデルを使用して STEM 画像を強調します。
+
+```bash
+python spectrum_enhancement/predict.py \
+    --model_name='sfin_haadf_enhance' \
+    --weights_name='best.pdparams' \
+    --input_path='path/to/noisy_image.png' \
+    --output_dir='result_sfin/'
+```
+
+---
+
+### 🏋️ 学習を始める
+
+学習およびファインチューニングについては、[ドキュメント](get_started.md)を参照してください。
+
+---
+
+## 🤝 コントリビューター・協力組織・コミュニティ
+
+[![Star History Chart](https://api.star-history.com/svg?repos=PaddlePaddle/PaddleMaterials&type=date&legend=top-left)](https://www.star-history.com/#PaddlePaddle/PaddleMaterials&type=date&legend=top-left)
+
+PaddleMaterials の構築に貢献してくださったすべての皆様に感謝します！
+
+<a href="https://github.com/PaddlePaddle/PaddleMaterials/graphs/contributors"><img src="https://contrib.rocks/image?repo=PaddlePaddle/PaddleMaterials" /></a>
+
+ご協力いただいている以下の組織に感謝します！
+
+<p align="left">
+  <img src="docs/logo_SZNL_2.jpeg" align="middle" width="180"/>
+  <img src="docs/logo_SinochemDI_2.jpeg" align="middle" width="180"/>
+  <img src="docs/logo_MetaX.png" align="middle" width="140"/>
+</p>
+
+PaddleMaterials の WeChat グループに参加して、ぜひ交流してください！
+
+<p align="left"><img src="docs/wechat_group.png" align="middle" width="100"/></p>
+
+## 🛠️ PaddleMaterials への貢献
+
+開発者の方は[アーキテクチャドキュメント](docs/ARCHITECTURE_ch.md)を参照してください。
+
+---
+
+## 📜 ライセンス
+
+PaddleMaterials は [Apache License 2.0](LICENSE) の下で提供されています。
+
+---
+
+## 🎓 引用
+
+```bibtex
+@misc{paddlematerials2025,
+  title={PaddleMaterials, a deep learning toolkit based on PaddlePaddle for material science.},
+  author={PaddleMaterials Contributors},
+  howpublished = {\url{https://github.com/PaddlePaddle/PaddleMaterials}},
+  year={2025}
+}
+```
+
+---
+
+## 🙏 謝辞
+
+本リポジトリは、以下のプロジェクトのコードを参考にしています。
+
+[PaddleScience](https://github.com/PaddlePaddle/PaddleScience) |
+[Matgl](https://github.com/materialsvirtuallab/matgl) |
+[CDVAE](https://github.com/txie-93/cdvae) |
+[DiffCSP](https://github.com/jiaor17/DiffCSP) |
+[MatterGen](https://github.com/microsoft/mattergen) |
+[MatterSim](https://github.com/microsoft/mattersim) |
+[CHGNet](https://github.com/CederGroupHub/chgnet) |
+[AIRS](https://github.com/divelab/AIRS)
diff --git a/README_zh.md b/README_zh.md
new file mode 100644
index 00000000..7ef17470
--- /dev/null
+++ b/README_zh.md
@@ -0,0 +1,225 @@
+# PaddleMaterials
+
+<p align="center">
+  <a href="README.md"><img src="https://img.shields.io/badge/English-lightgrey" alt="English"></a>
+  <a href="README_zh.md"><img src="https://img.shields.io/badge/简体中文-blue" alt="简体中文"></a>
+  <a href="README_ja.md"><img src="https://img.shields.io/badge/日本語-lightgrey" alt="日本語"></a>
+</p>
+
+<p align="center"><img src="docs/ppmat_logo.png" alt="PaddleMaterials" width="400"></p>
+
+<p align="center">
+  <a href="Install.md"><img alt="Python 3.10+" src="https://img.shields.io/badge/Python-3.10%2B-3776AB?logo=python&amp;logoColor=white"></a>
+  <a href="https://pypi.org/project/ppmat/"><img alt="PyPI 版本" src="https://img.shields.io/pypi/v/ppmat?logo=pypi&amp;logoColor=white"></a>
+  <a href="LICENSE"><img alt="Apache 2.0 许可证" src="https://img.shields.io/github/license/PaddlePaddle/PaddleMaterials"></a>
+  <a href="https://github.com/PaddlePaddle/PaddleMaterials/stargazers"><img alt="GitHub Stars" src="https://img.shields.io/github/stars/PaddlePaddle/PaddleMaterials?style=flat&amp;logo=github"></a>
+</p>
+
+## 🚀 简介
+
+**PaddleMaterials** 是一款基于 **PaddlePaddle** 深度学习框架构建的端到端 AI4Materials 工具包。作为面向材料科学基础模型开发与部署的数据—机理双驱动平台，**PPMat** 帮助研究人员高效构建 AI 模型，并利用预训练模型加速材料发现。
+
+<p align="left"><img src="docs/overview_en.png" align="middle" width="1000"/></p>
+
+### 🧩 核心能力
+
+| 任务 | 描述 | 典型应用 |
+|------|------|----------|
+| **性质预测（PP）** | 根据材料结构预测性质 | 正向设计，或预测形成能、带隙、弹性模量等 |
+| **结构生成（SG）** | 生成新型晶体结构 | 逆向设计或结构生成 |
+| **机器学习原子间势（MLIP）** | 用机器学习势作为 DFT 的代理模型 | 分子动力学模拟 |
+| **电子结构（ES）** | 用代理模型替代 DFT 预测物理场 | 电子密度预测 |
+| **谱图解析（SE）** | 根据谱图重建结构 | NMR 结构解析 |
+| **谱图增强（SPEN）** | 增强显微图像与谱图信号 | STEM 图像增强、去噪 |
+
+### 🧱 支持的材料体系
+
+- **无机晶体** — 支持完善，提供多个数据集和预训练模型
+- **有机分子** — 支持小分子和部分聚合物等多个数据集及预训练模型
+
+### ✨ 为什么选择 PaddleMaterials？
+
+- ✅ **丰富的预训练模型与 AI-ready 数据集** — 提供 50+ 个可直接推理的预训练模型，以及多个用于训练的精选数据集
+- ✅ **多任务集成** — 使用统一框架支持 PP、SG、MLIP、ES、SE、SPEN 等任务
+- ✅ **多硬件支持** — 全面支持 NVIDIA GPU、MetaX GPU 和 Intel CPU
+- ✅ **生产就绪** — 标准化设计，易于使用，并支持分布式训练、混合精度和断点恢复
+
+### 📑 支持的任务
+
+| 任务 | 描述 | 链接 |
+|------|------|------|
+| **性质预测（PP）** | 预测形成能、带隙和弹性性质 | [README](property_prediction/README.md) |
+| **结构生成（SG）** | 使用扩散模型生成新的晶体结构 | [README](structure_generation/README.md) |
+| **机器学习原子间势（MLIP）** | 用于分子动力学、具有 DFT 精度的势函数 | [README](interatomic_potentials/README.md) |
+| **电子结构（ES）** | 预测电子结构性质 | [README](electronic_structure/README.md) |
+| **谱图解析（SE）** | 根据 NMR 谱图重建分子结构 | [README](spectrum_elucidation/README.md) |
+| **谱图增强（SPEN）** | 增强显微图像和谱图信号 | [README](spectrum_enhancement/README.md) |
+
+### 🤖 可用的预训练模型
+
+| 任务 | 模型 | 数据集 |
+|------|------|--------|
+| **性质预测** | MEGNet、iComformer、DimeNet++、SphereNet | MP2018、MP2024、JARVIS、QM9 等 |
+| **结构生成** | MatterGen、DiffCSP | MP20、ALEX 等 |
+| **机器学习原子间势** | CHGNet、MatterSim、SphereNet | MPTRJ、MD17 等 |
+| **电子结构** | InfGCN | QM9_ES、MP_ES、OMol25_MC_ES 等 |
+| **谱图解析** | DiffNMR | MSD_NMR 等 |
+| **谱图增强** | SFIN | SFIN-HAADF/BF 等 |
+
+完整模型列表：参见 [MODEL_REGISTRY](ppmat/models/__init__.py#L75)。
+
+---
+
+## 🚀 快速开始
+
+### 🔧 安装
+
+请根据您的硬件环境参阅[安装文档](Install.md)。有关多硬件适配的更多信息，请参阅[支持的硬件列表](./docs/multi_device.md)。
+
+---
+
+### ⚡ 快速推理
+
+#### 性质预测
+
+使用预训练 MEGNet 模型预测材料形成能：
+
+```bash
+python property_prediction/predict.py \
+    --model_name='megnet_mp2018_train_60k_e_form' \
+    --weights_name='best.pdparams' \
+    --cif_file_path='./property_prediction/example_data/cifs/' \
+    --save_path='result.csv'
+```
+
+#### 结构生成
+
+使用预训练 MatterGen 模型生成新的晶体结构：
+
+```bash
+python structure_generation/sample.py \
+    --model_name='mattergen_mp20' \
+    --weights_name='latest.pdparams' \
+    --output_dir='result_mattergen_mp20/' \
+    --mode='by_num_atoms' \
+    --num_atoms=4
+```
+
+#### 原子间势
+
+使用预训练 MatterSim 模型预测能量和力：
+
+```bash
+python interatomic_potentials/predict.py \
+    --model_name='mattersim_1M' \
+    --weights_name='mattersim-v1.0.0-1M_model.pdparams' \
+    --cif_file_path='./interatomic_potentials/example_data/cifs/' \
+    --save_path='result.csv'
+```
+
+#### 电子结构
+
+使用预训练 InfGCN 模型对仓库自带的甲烷示例预测电子密度：
+
+```bash
+python electronic_structure/predict.py \
+    --model_name='infgcn_qm9' \
+    --weights_name='best.pdparams' \
+    --mol_input='electronic_structure/configs/infgcn/example/methane.mol' \
+    --atom_file='electronic_structure/configs/qm9.json' \
+    --mol_grid_shape=8 \
+    --output_dir='output/infgcn_qm9/methane' \
+    --save_pred_cube
+```
+
+有关基于数据集或本地检查点推理的方法，请参阅
+[InfGCN 预测指南](electronic_structure/configs/infgcn/README.md#prediction)。
+
+#### 谱图解析
+
+使用预训练 DiffNMR 模型和仓库自带示例运行 NMR 谱图解析：
+
+```bash
+python spectrum_elucidation/sample.py \
+    --model_name='diffnmr_msdnmr_nless15' \
+    --weights_name='best.pdparams' \
+    --output_dir='result_diffnmr_nless15/'
+```
+
+#### 谱图增强
+
+使用预训练 SFIN 模型增强 STEM 图像：
+
+```bash
+python spectrum_enhancement/predict.py \
+    --model_name='sfin_haadf_enhance' \
+    --weights_name='best.pdparams' \
+    --input_path='path/to/noisy_image.png' \
+    --output_dir='result_sfin/'
+```
+
+---
+
+### 🏋️ 开始训练
+
+有关训练和微调的方法，请参阅[相关文档](get_started.md)。
+
+---
+
+## 🤝 贡献者、合作伙伴与社区
+
+[![Star History Chart](https://api.star-history.com/svg?repos=PaddlePaddle/PaddleMaterials&type=date&legend=top-left)](https://www.star-history.com/#PaddlePaddle/PaddleMaterials&type=date&legend=top-left)
+
+感谢所有为 PaddleMaterials 建设作出贡献的开发者！
+
+<a href="https://github.com/PaddlePaddle/PaddleMaterials/graphs/contributors"><img src="https://contrib.rocks/image?repo=PaddlePaddle/PaddleMaterials" /></a>
+
+感谢以下组织提供合作支持！
+
+<p align="left">
+  <img src="docs/logo_SZNL_2.jpeg" align="middle" width="180"/>
+  <img src="docs/logo_SinochemDI_2.jpeg" align="middle" width="180"/>
+  <img src="docs/logo_MetaX.png" align="middle" width="140"/>
+</p>
+
+欢迎加入 PaddleMaterials 微信群与我们交流！
+
+<p align="left"><img src="docs/wechat_group.png" align="middle" width="100"/></p>
+
+## 🛠️ 参与 PaddleMaterials
+
+开发者请参阅[架构文档](docs/ARCHITECTURE_ch.md)。
+
+---
+
+## 📜 许可证
+
+PaddleMaterials 基于 [Apache License 2.0](LICENSE) 许可。
+
+---
+
+## 🎓 引用
+
+```bibtex
+@misc{paddlematerials2025,
+  title={PaddleMaterials, a deep learning toolkit based on PaddlePaddle for material science.},
+  author={PaddleMaterials Contributors},
+  howpublished = {\url{https://github.com/PaddlePaddle/PaddleMaterials}},
+  year={2025}
+}
+```
+
+---
+
+## 🙏 致谢
+
+本仓库参考了以下项目的代码：
+
+[PaddleScience](https://github.com/PaddlePaddle/PaddleScience) |
+[Matgl](https://github.com/materialsvirtuallab/matgl) |
+[CDVAE](https://github.com/txie-93/cdvae) |
+[DiffCSP](https://github.com/jiaor17/DiffCSP) |
+[MatterGen](https://github.com/microsoft/mattergen) |
+[MatterSim](https://github.com/microsoft/mattersim) |
+[CHGNet](https://github.com/CederGroupHub/chgnet) |
+[AIRS](https://github.com/divelab/AIRS)
diff --git a/docs/MetaX/PaddleMaterials_MetaX_README.md b/docs/MetaX/PaddleMaterials_MetaX_README.md
index dd509dc1..8cc7f565 100644
--- a/docs/MetaX/PaddleMaterials_MetaX_README.md
+++ b/docs/MetaX/PaddleMaterials_MetaX_README.md
@@ -31,7 +31,7 @@ python structure_generation/train.py -c structure_generation/configs/diffcsp/dif
 
 ### 2) Sample
 ```bash
-python structure_generation/sample.py --model_name='diffcsp_mp20' --weights_name='latest.pdparams' --save_path='result_diffcsp_mp20-1/' --chemical_formula='LiMnO2'
+python structure_generation/sample.py --model_name='diffcsp_mp20' --weights_name='latest.pdparams' --output_dir='result_diffcsp_mp20-1/' --chemical_formula='LiMnO2'
 ```
 ![](./pic5.png)
 
diff --git a/docs/model_package.md b/docs/model_package.md
new file mode 100644
index 00000000..fdd289c9
--- /dev/null
+++ b/docs/model_package.md
@@ -0,0 +1,40 @@
+# Pretrained model package contract
+
+Registered pretrained models must use a predictable archive layout. The archive
+name, top-level directory, and configuration stem must all match the
+`MODEL_REGISTRY` key:
+
+```text
+<model_name>.zip
+└── <model_name>/
+    ├── <model_name>.yaml
+    └── checkpoints/
+        └── best.pdparams
+```
+
+`best.pdparams` must contain the model state dictionary directly:
+
+```python
+paddle.save(model.state_dict(), "best.pdparams")
+```
+
+Do not package a trainer checkpoint such as `{"model": state_dict, "step": ...}`
+as `best.pdparams`. Training metadata belongs in a separate file. A model that
+needs auxiliary pretrained components may add clearly named files under
+`checkpoints/`, while its primary weight remains `best.pdparams`.
+
+The package may include inference-only examples or assets beside
+`checkpoints/`. Paths in the YAML must be relative to the package directory,
+and package consumers must resolve them explicitly without changing the process
+working directory.
+
+Before publishing a registry entry:
+
+1. Download and extract the final archive through
+   `ppmat.utils.download.get_weights_path_from_url`.
+2. Resolve it with `resolve_model_package_dir`.
+3. Build the model from its packaged YAML.
+4. Load `checkpoints/best.pdparams` without missing or unexpected parameters.
+5. Run the documented inference command from a clean cache.
+6. Record the final archive SHA256 in the release or pull-request validation
+   notes.
diff --git a/electronic_structure/configs/infgcn/README.md b/electronic_structure/configs/infgcn/README.md
index f5ec02dd..c0bdf418 100644
--- a/electronic_structure/configs/infgcn/README.md
+++ b/electronic_structure/configs/infgcn/README.md
@@ -88,7 +88,7 @@ $$
 ### Datasets
 - **QM9_EC**: Electron densities stored as `*.CHGCAR.lz4` in `dataset_ES/data_qm9` (train 123,835 / val 50 / test 10,000). [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/QM9_ES/qm9_es.tar), [Atom dictionary](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/QM9_ES/qm9.json), [Split file](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/QM9_ES/qm9_data_split.json).
 - **MP_EC (cubic)**: Materials Project-style crystals serialized as `.json.xz` under `dataset_ES/data_cubic` (train 14,421 / val 1,000 / test 1,000). [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MP_ES/mp_es.tar), [Atom dictionary](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MP_ES/crystal.json), [Split file](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MP_ES/crystal_data_split.json).
-- **OMol25_EC**: Organic molecule cubes expected under `/home/liuxuwei01/processed_output` (train 16 / val 2 / test 2). [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/OMol25_ES/MC_5k/omol25_mc_5k.tar), [Atom dictionary](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/OMol25_ES/MC_5k/omol25.json), [Split file](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/OMol25_ES/MC_5k/omol25_data_split.json).
+- **OMol25_EC**: Organic molecule cubes expected under `data/dataset_OMol25_MC_5k` (train 16 / val 2 / test 2). [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/OMol25_ES/MC_5k/omol25_mc_5k.tar), [Atom dictionary](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/OMol25_ES/MC_5k/omol25.json), [Split file](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/OMol25_ES/MC_5k/omol25_data_split.json).
 - **MD17_EC**: Small molecules (for example, ethanol, benzene, phenol, resorcinol) from the MD17 electron-density release in `dataset_ES/data_md`; default config trains on ethanol. [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MD17_ES/md17_es.tar.gz).
 
 ---
@@ -115,7 +115,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">59min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_md17_benzene.yaml">infgcn_md17_benzene</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_benzene_t_20260104_083255_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_benzene.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_md17_ethane</td>
@@ -124,7 +124,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">1hour17min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_md17_ethane.yaml">infgcn_md17_ethane</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_ethane_t_20260105_035402_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_ethane.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_md17_ethanol</td>
@@ -133,7 +133,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">7min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_md17_ethanol.yaml">infgcn_md17_ethanol</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_ethanol_t_20260105_051323_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_ethanol.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_md17_malonaldehyde</td>
@@ -142,7 +142,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">1hour29min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_md17_malonaldehyde.yaml">infgcn_md17_malonaldehyde</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_malonaldehyde_t_20260105_052204_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_malonaldehyde.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_md17_phenol</td>
@@ -151,7 +151,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">1hour17min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_md17_phenol.yaml">infgcn_md17_phenol</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_phenol_t_20260107_041611_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_phenol.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_md17_resorcinol</td>
@@ -160,7 +160,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">1hour23min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_md17_resorcinol.yaml">infgcn_md17_resorcinol</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_resorcinol_t_20260104_083255_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_resorcinol.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_qm9</td>
@@ -169,7 +169,7 @@ $$
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">75hour41min</td>
             <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_qm9.yaml">infgcn_qm9</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_qm9_t_20260107_113954_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_qm9.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_cubic</td>
@@ -177,8 +177,8 @@ $$
             <td nowrap="nowrap">47.3829%</td>
             <td nowrap="nowrap">1</td>
             <td nowrap="nowrap">12hour6min</td>
-            <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_cubic.yaml">infgcn_cubic</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_mp_t_20260108_024145_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_mp.yaml">infgcn_cubic</a></td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_mp.zip">checkpoint</td>
         </tr>
         <tr>
             <td nowrap="nowrap">infgcn_omol25_mc_5k</td>
@@ -186,13 +186,13 @@ $$
             <td nowrap="nowrap">12.6260%</td>
             <td nowrap="nowrap">4</td>
             <td nowrap="nowrap">66hour28min</td>
-            <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_omol25_trimmed.yaml">infgcn_omol25</a></td>
-            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42.zip">checkpoint | log</td>
+            <td nowrap="nowrap"><a href="../../../electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml">infgcn_omol25</a></td>
+            <td nowrap="nowrap"><a href="https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_omol25_mc_5k_trimmed.zip">checkpoint</td>
         </tr>
     </tbody>
 </table>
 
-**Note**: Benchmarks are being regenerated in Paddle; metrics and downloadable checkpoints will be published once validation completes. Pretrained QM9 weights: [infgcn_qm9](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_qm9.pdparams)
+**Note**: Benchmarks are being regenerated in Paddle; metrics and downloadable checkpoints will be published once validation completes. Pretrained packages store weights under `checkpoints/`.
 
 ---
 
@@ -220,50 +220,45 @@ python electronic_structure/train.py -c electronic_structure/configs/infgcn/infg
 
 ### Prediction
 ```bash
-# 1) Dataset-sample inference (uses dataset paths from the YAML unless overridden).
+# 1) One-click inference with a registered pretrained model and bundled molecule.
 python electronic_structure/predict.py \
-  --config electronic_structure/configs/infgcn/infgcn_qm9.yaml \
-  --checkpoint output/infgcn_qm9_best/infgcn_qm9.pdparams \
-  --split validation \
-  --index 0 \
-  --grid_batch_size 20000 \
-  --output_dir output/infgcn_qm9_best/vis_val0 \
-  --save_pred_cube \
-  --save_true_cube \
-  --cube_dir output/infgcn_qm9_best/cubes
-
-# 2) MOL-file inference (single file or directory).
+  --model_name infgcn_qm9 \
+  --weights_name best.pdparams \
+  --mol_input electronic_structure/configs/infgcn/example/methane.mol \
+  --atom_file electronic_structure/configs/qm9.json \
+  --mol_grid_shape 8 \
+  --grid_batch_size 128 \
+  --skip_vis
+
+# 2) Dataset-sample inference with a custom config and checkpoint.
+python electronic_structure/predict.py \
+  --config_path electronic_structure/configs/infgcn/infgcn_qm9.yaml \
+  --checkpoint_path path/to/infgcn_qm9.pdparams
+
+# 3) MOL-file inference (single file or directory).
 # This mode predicts electron density from molecular structure files (*.mol),
 # and can export predicted cube + html visualization.
-CUDA_VISIBLE_DEVICES=4 conda run -n ppmat python electronic_structure/predict.py \
-  --config output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/infgcn_omol25_trimmed.yaml \
-  --checkpoint output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/checkpoints/latest.pdparams \
-  --mol_input mols/Baidu_infGCN_Example_20260206 \
-  --atom_file /home/liuxuwei01/processed_output/omol25.json \
-  --output_dir output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/mol_predict_latest_gpu4 \
-  --cube_dir output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/mol_predict_latest_gpu4/cubes \
-  --save_pred_cube \
-  --save_html \
-  --grid_batch_size 4096
-
-# 3) MOL-file inference with reference (true) cube files.
+python electronic_structure/predict.py \
+  --config_path electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml \
+  --checkpoint_path path/to/infgcn_omol25.pdparams \
+  --mol_input path/to/mols_or_mol_file
+
+# 4) MOL-file inference with reference (true) cube files.
 # If --mol_true_cube_dir provides matching files (<name>.cube or <name>_true.cube),
 # the script additionally writes true cube and true/diff html.
-CUDA_VISIBLE_DEVICES=4 conda run -n ppmat python electronic_structure/predict.py \
-  --config output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/infgcn_omol25_trimmed.yaml \
-  --checkpoint output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/checkpoints/latest.pdparams \
-  --mol_input mols/Baidu_infGCN_Example_20260206 \
-  --mol_true_cube_dir /path/to/true_cubes \
-  --atom_file /home/liuxuwei01/processed_output/omol25.json \
-  --output_dir output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/mol_predict_latest_gpu4 \
-  --cube_dir output/infgcn_omol25_s1_trimmed_t_20260118_183549_s_42/mol_predict_latest_gpu4/cubes \
-  --save_true_cube \
-  --save_pred_cube \
-  --save_html \
-  --grid_batch_size 4096
+python electronic_structure/predict.py \
+  --config_path electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml \
+  --checkpoint_path path/to/infgcn_omol25.pdparams \
+  --mol_input path/to/mols_or_mol_file \
+  --mol_true_cube_dir path/to/true_cubes \
+  --save_true_cube
 ```
 
 Notes:
+- Replace `path/to/*.pdparams` with a downloaded pretrained checkpoint or a checkpoint produced by training.
+- Prediction defaults such as `split`, `index`, `output_dir`, `grid_batch_size`,
+  cube export, html export, and MOL grid settings are configured under `Predict`
+  in each YAML and can still be overridden from the command line.
 - `--mol_input` supports either one `.mol` file or a directory of `.mol` files.
 - Optional grid controls for MOL mode: `--mol_grid_shape` (default `80,80,80`) and `--mol_grid_padding` (default `6.0` Angstrom).
 - If true/reference cube is not provided, only predicted outputs are available (`*_pred.cube`, `*_pred_density.html`).
diff --git a/electronic_structure/configs/infgcn/example/methane.mol b/electronic_structure/configs/infgcn/example/methane.mol
new file mode 100644
index 00000000..6abf7ea2
--- /dev/null
+++ b/electronic_structure/configs/infgcn/example/methane.mol
@@ -0,0 +1,14 @@
+Methane
+  PaddleMaterials
+
+  5  4  0  0  0  0  0  0  0  0999 V2000
+    0.0000    0.0000    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
+    0.6291    0.6291    0.6291 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.6291   -0.6291    0.6291 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.6291    0.6291   -0.6291 H   0  0  0  0  0  0  0  0  0  0  0  0
+    0.6291   -0.6291   -0.6291 H   0  0  0  0  0  0  0  0  0  0  0  0
+  1  2  1  0  0  0  0
+  1  3  1  0  0  0  0
+  1  4  1  0  0  0  0
+  1  5  1  0  0  0  0
+M  END
diff --git a/electronic_structure/configs/infgcn/infgcn_md17_benzene.yaml b/electronic_structure/configs/infgcn/infgcn_md17_benzene.yaml
index 8b799efe..e5b5a987 100644
--- a/electronic_structure/configs/infgcn/infgcn_md17_benzene.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_md17_benzene.yaml
@@ -157,3 +157,17 @@ Predict:
   num_infer: null
   num_vis: 2
   inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_md17_benzene/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_md17_benzene/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_md17_ethane.yaml b/electronic_structure/configs/infgcn/infgcn_md17_ethane.yaml
index 4516b860..dddc85b6 100644
--- a/electronic_structure/configs/infgcn/infgcn_md17_ethane.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_md17_ethane.yaml
@@ -156,3 +156,17 @@ Predict:
   num_infer: null
   num_vis: 2
   inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_md17_ethane/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_md17_ethane/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_md17_ethanol.yaml b/electronic_structure/configs/infgcn/infgcn_md17_ethanol.yaml
index 7891453d..6fe57a59 100644
--- a/electronic_structure/configs/infgcn/infgcn_md17_ethanol.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_md17_ethanol.yaml
@@ -157,3 +157,17 @@ Predict:
   num_infer: null
   num_vis: 2
   inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_md17_ethanol/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_md17_ethanol/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_md17_malonaldehyde.yaml b/electronic_structure/configs/infgcn/infgcn_md17_malonaldehyde.yaml
index 46d1e658..caabc4b8 100644
--- a/electronic_structure/configs/infgcn/infgcn_md17_malonaldehyde.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_md17_malonaldehyde.yaml
@@ -157,3 +157,17 @@ Predict:
   num_infer: null
   num_vis: 2
   inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_md17_malonaldehyde/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_md17_malonaldehyde/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_md17_phenol.yaml b/electronic_structure/configs/infgcn/infgcn_md17_phenol.yaml
index bc308d1a..1521c88a 100644
--- a/electronic_structure/configs/infgcn/infgcn_md17_phenol.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_md17_phenol.yaml
@@ -157,3 +157,17 @@ Predict:
   num_infer: null
   num_vis: 2
   inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_md17_phenol/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_md17_phenol/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_md17_resorcinol.yaml b/electronic_structure/configs/infgcn/infgcn_md17_resorcinol.yaml
index e224fd97..bb4b9ee9 100644
--- a/electronic_structure/configs/infgcn/infgcn_md17_resorcinol.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_md17_resorcinol.yaml
@@ -157,3 +157,17 @@ Predict:
   num_infer: null
   num_vis: 2
   inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_md17_resorcinol/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_md17_resorcinol/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_mp.yaml b/electronic_structure/configs/infgcn/infgcn_mp.yaml
index 7202978a..27e17eb9 100644
--- a/electronic_structure/configs/infgcn/infgcn_mp.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_mp.yaml
@@ -170,3 +170,23 @@ Dataset:
       collate_fn: DensityCollator # DensityVoxelCollator # if use voxel
       collate_params:
         n_samples: 2048 # recommend set blank if have enough gpu memory
+
+Predict:
+  eval_with_no_grad: True
+  num_infer: null
+  num_vis: 2
+  inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: ./output/infgcn_mp/predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: False
+  cube_dir: ./output/infgcn_mp/predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml b/electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml
index 3f77e897..3f5b4e0f 100644
--- a/electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_omol25_MC_5k_trimmed.yaml
@@ -94,7 +94,7 @@ Optimizer:
     beta1: 0.9
     beta2: 0.999
 
-# remind: just subsets of Metal Complexes of OMol25 
+# remind: just subsets of Metal Complexes of OMol25
 Dataset:
   train:
     dataset:
@@ -160,6 +160,7 @@ Dataset:
         importance_ratio: 0.6
         extreme_threshold: 100.0
         extreme_ratio: 0.05
+
   test:
     dataset:
       __class_name__: DensityDataset
@@ -193,3 +194,23 @@ Dataset:
         importance_ratio: 0.6
         extreme_threshold: 100.0
         extreme_ratio: 0.05
+
+Predict:
+  eval_with_no_grad: True
+  num_infer: null
+  num_vis: 2
+  inf_samples: 4096
+  split: test
+  index: 0
+  output_dir: output/infgcn_omol25/mol_predict
+  grid_batch_size: 4096
+  skip_vis: False
+  save_true_cube: False
+  save_pred_cube: True
+  save_html: True
+  cube_dir: output/infgcn_omol25/mol_predict/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/configs/infgcn/infgcn_qm9.yaml b/electronic_structure/configs/infgcn/infgcn_qm9.yaml
index 79aaba21..43da515a 100644
--- a/electronic_structure/configs/infgcn/infgcn_qm9.yaml
+++ b/electronic_structure/configs/infgcn/infgcn_qm9.yaml
@@ -165,7 +165,7 @@ Dataset:
         batch_size: 1
     loader:
       num_workers: 4
-      use_shared_memory: Flase
+      use_shared_memory: False
       collate_fn: DensityCollator # DensityVoxelCollator # if use voxel
       collate_params:
         n_samples: 2048 # recommend set blank if have enough gpu memory
@@ -175,3 +175,17 @@ Predict:
   num_infer: 100
   num_vis: 2
   inf_samples: 4096
+  split: validation
+  index: 0
+  output_dir: output/infgcn_qm9/vis_val0
+  grid_batch_size: 20000
+  skip_vis: False
+  save_true_cube: True
+  save_pred_cube: True
+  save_html: False
+  cube_dir: output/infgcn_qm9/cubes
+  show_plot: False
+  mol_pattern: "*.mol"
+  mol_grid_shape: "80,80,80"
+  mol_grid_padding: 6.0
+  mol_true_cube_dir: null
diff --git a/electronic_structure/predict.py b/electronic_structure/predict.py
index d466f3ce..7c386a9b 100644
--- a/electronic_structure/predict.py
+++ b/electronic_structure/predict.py
@@ -12,692 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import argparse
-import copy
-import gzip
-import json
-import lzma
-import math
-from pathlib import Path
-import numpy as np
-import time
-
-import paddle
-import plotly.graph_objects as go
-from omegaconf import OmegaConf
-from tqdm import tqdm
-
-try:
-    from IPython.display import Image, display
-except ImportError:  # Optional dependency; visualization still works for files
-    Image, display = None, None
-
-from ppmat.datasets import DensityDataset
-from ppmat.datasets import SmallDensityDataset
-from ppmat.datasets.geometric_data_type.data import Data
-from ppmat.models import build_model
-from ppmat.utils import logger
-from ppmat.utils.misc import set_random_seed
-
-BOHR2ANG = 0.529177
-ANG2BOHR = 1.0 / BOHR2ANG
-
-
-def get_pretrained_model(cfg_path, model_path):
-    logger.info(f"from {cfg_path} loading config")
-    cfg = OmegaConf.load(cfg_path)
-    cfg = OmegaConf.to_container(cfg, resolve=True)
-
-    model = build_model(cfg["Model"])
-    logger.info(f"from {model_path}loading model")
-    # if a directory is given, pick best > latest > highest epoch > any pdparams
-    mpath = Path(model_path)
-    if mpath.is_dir():
-        candidates = list(mpath.glob("**/*.pdparams"))
-        chosen = None
-        for name in ["best.pdparams", "latest.pdparams"]:
-            hits = [c for c in candidates if c.name == name]
-            if hits:
-                chosen = hits[0]
-                break
-        if chosen is None:
-            epochs = []
-            for c in candidates:
-                stem = c.stem
-                if stem.startswith("epoch_"):
-                    try:
-                        ep = int(stem.split("_")[1])
-                        epochs.append((ep, c))
-                    except Exception:
-                        pass
-            if epochs:
-                epochs.sort(key=lambda x: -x[0])
-                chosen = epochs[0][1]
-        if chosen is None and candidates:
-            chosen = candidates[0]
-        if chosen is None:
-            raise FileNotFoundError(f"No .pdparams found under {model_path}")
-        model_path = str(chosen)
-        logger.info(f"Resolved checkpoint path: {model_path}")
-
-    state_dict = paddle.load(model_path)
-    if isinstance(state_dict, dict) and "model" in state_dict:
-        model.set_state_dict(state_dict["model"])
-    else:
-        model.set_state_dict(state_dict)
-    return model
-
-
-def inference_model(model, g, density, grid_coord, infos, grid_batch_size=8196):
-    with paddle.no_grad():
-        model.eval()
-        device = paddle.get_device()
-        prepared_infos = (
-            model._prepare_infos(infos, device) if hasattr(model, "_prepare_infos") else infos
-        )
-        if grid_batch_size is None:
-            if hasattr(model, "_forward_density"):
-                preds = model._forward_density(
-                    g.x, g.pos, grid_coord, g.batch, prepared_infos
-                ).squeeze(0)
-            else:
-                # Fallback for legacy models expecting raw tensors
-                preds = model(g.x, g.pos, grid_coord, g.batch, prepared_infos).squeeze(0)
-        else:
-            preds = []
-            total = grid_coord.shape[1]
-            step = grid_batch_size
-            num_iter = (total + step - 1) // step
-            for start in tqdm(range(0, total, step), total=num_iter):
-                end = min(start + step, total)
-                grid = grid_coord[:, start:end]
-                if hasattr(model, "_forward_density"):
-                    preds.append(
-                        model._forward_density(
-                            g.x, g.pos, grid, g.batch, prepared_infos
-                        ).squeeze(0)
-                    )
-                else:
-                    preds.append(
-                        model(g.x, g.pos, grid, g.batch, prepared_infos).squeeze(0)
-                    )
-            preds = paddle.concat(preds, axis=0)
-
-        if density is None:
-            return preds, None, None
-
-        mask = (density > 0).astype(dtype="float32")
-        preds = preds * mask
-        density = density * mask
-        diff = paddle.abs(preds - density)
-        loss = diff.pow(2).sum()
-        denom = paddle.clip(density.sum(), min=1e-12)
-        mae = diff.sum() / denom
-    return preds, loss, mae
-
-
-def draw_volume(
-    grid,
-    density,
-    atom_type,
-    atom_coord,
-    isomin=0.05,
-    isomax=None,
-    surface_count=5,
-    title=None,
-):
-    atom_colorscale = ["grey", "white", "red", "blue", "green"]
-    fig = go.Figure()
-    fig.add_trace(
-        go.Volume(
-            x=grid[..., 0],
-            y=grid[..., 1],
-            z=grid[..., 2],
-            value=density,
-            isomin=isomin,
-            isomax=isomax,
-            opacity=0.1,
-            surface_count=surface_count,
-            caps=dict(x_show=False, y_show=False, z_show=False),
-        )
-    )
-
-    axis_dict = dict(
-        showgrid=False,
-        showbackground=False,
-        zeroline=False,
-        visible=False,
-    )
-
-    fig.add_trace(
-        go.Scatter3d(
-            x=atom_coord[:, 0],
-            y=atom_coord[:, 1],
-            z=atom_coord[:, 2],
-            mode="markers",
-            marker=dict(
-                size=10,
-                color=atom_type,
-                cmin=0,
-                cmax=4,
-                colorscale=atom_colorscale,
-                opacity=0.6,
-            ),
-        )
-    )
-
-    if title is not None:
-        title = dict(
-            text=title,
-            x=0.5,
-            y=0.3,
-            xanchor="center",
-            yanchor="bottom",
-        )
-
-    fig.update_layout(
-        autosize=False,
-        width=800,
-        height=800,
-        showlegend=False,
-        scene=dict(xaxis=axis_dict, yaxis=axis_dict, zaxis=axis_dict),
-        title=title,
-        title_font_family="Times New Roman",
-    )
-
-    return fig
-
-
-def safe_write_image(fig, path, show_plot=False):
-    try:
-        fig.write_image(path)
-        logger.info(f"Image saved to: {path}")
-    except Exception as e:
-        logger.warning(f"Failed to save image {path}: {e}")
-        try:
-            html_path = path.with_suffix(".html")
-            fig.write_html(html_path)
-            logger.info(f"Saved interactive HTML instead: {html_path}")
-        except Exception as html_e:
-            logger.warning(f"Failed to save HTML fallback for {path}: {html_e}")
-
-    if show_plot:
-        try:
-            if Image is None or display is None:
-                raise ImportError("IPython not installed")
-            img_bytes = fig.to_image(format="png", scale=2)
-            display(Image(img_bytes))
-        except Exception as e:
-            logger.warning(f"Failed to display image: {e}")
-
-
-def maybe_downsample_volume(grid, values, shape, max_points=250_000):
-    """
-    Downsample a regular 3D grid for visualization to keep Plotly volume traces responsive.
-    grid: numpy array of shape (n_points, 3)
-    values: list of numpy arrays aligned with grid, each of shape (n_points,)
-    shape: original lattice shape [nx, ny, nz]
-    """
-    if shape is None or len(shape) != 3:
-        return grid, values, False, 1
-
-    try:
-        shape = [int(s) for s in shape]
-        total = shape[0] * shape[1] * shape[2]
-    except Exception:
-        return grid, values, False, 1
-
-    if total != grid.shape[0] or any(val.shape[0] != grid.shape[0] for val in values):
-        return grid, values, False, 1
-    if total <= max_points:
-        return grid, values, False, 1
-
-    stride = max(1, math.ceil((total / max_points) ** (1 / 3)))
-    try:
-        grid_view = grid.reshape(shape[0], shape[1], shape[2], 3)
-        grid_ds = grid_view[::stride, ::stride, ::stride, :].reshape(-1, 3)
-        values_ds = [
-            val.reshape(shape[0], shape[1], shape[2])[::stride, ::stride, ::stride].reshape(-1)
-            for val in values
-        ]
-    except Exception as e:
-        logger.warning(f"Failed to downsample grid for visualization: {e}")
-        return grid, values, False, 1
-
-    return grid_ds, values_ds, True, stride
-
-
-def write_cube_generic(fileobj, atom_type, atom_coord, density, info, idx2atom_num=None):
-    """
-    Minimal cube writer for datasets without a built-in write_cube method.
-    idx2atom_num maps dataset atom indices to atomic numbers (e.g., [6,1,8] for C/H/O).
-    """
-    fileobj.write("Cube file written on " + time.strftime("%c"))
-    fileobj.write("\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")
-    cell = info["cell"]
-    shape = info["shape"]
-    origin = info.get("origin", np.zeros(3, dtype=np.float32))
-    fileobj.write("{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(len(atom_type), *origin))
-    for s, c in zip(shape, cell):
-        d = c / s
-        fileobj.write("{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(s, *d))
-    for Z, (x, y, z) in zip(atom_type, atom_coord):
-        atomic_num = int(idx2atom_num[int(Z)]) if idx2atom_num is not None else int(Z)
-        fileobj.write(
-            "{0:5}{1:12.6f}{2:12.6f}{3:12.6f}{4:12.6f}\n".format(
-                atomic_num, float(atomic_num), x, y, z
-            )
-        )
-    density.tofile(fileobj, sep="\n", format="%e")
-
-
-def parse_grid_shape(shape_str):
-    parts = [p.strip() for p in str(shape_str).split(",") if p.strip()]
-    if len(parts) == 1:
-        n = int(parts[0])
-        if n <= 1:
-            raise ValueError(f"Invalid mol_grid_shape {shape_str}, each dimension must be > 1")
-        return [n, n, n]
-    if len(parts) == 3:
-        shape = [int(p) for p in parts]
-        if any(s <= 1 for s in shape):
-            raise ValueError(f"Invalid mol_grid_shape {shape_str}, each dimension must be > 1")
-        return shape
-    raise ValueError(f"Invalid mol_grid_shape {shape_str}, expected 'N' or 'Nx,Ny,Nz'")
-
-
-def normalize_element_symbol(symbol):
-    sym = str(symbol).strip()
-    if len(sym) == 0:
-        return sym
-    if len(sym) == 1:
-        return sym.upper()
-    return sym[0].upper() + sym[1:].lower()
-
-
-def load_atom_mapping(atom_file):
-    with Path(atom_file).open() as f:
-        atom_info = json.load(f)
-
-    atom_name2idx = {}
-    idx2atom_num = {}
-    for idx, item in enumerate(atom_info):
-        sym = normalize_element_symbol(item["name"])
-        atom_name2idx[sym] = idx
-        idx2atom_num[idx] = int(item["atom_num"])
-    return atom_name2idx, idx2atom_num
-
-
-def resolve_atom_file_for_mol(args_atom_file, dataset_atom_file):
-    candidates = []
-    if args_atom_file is not None:
-        candidates.append(Path(args_atom_file).expanduser())
-    if dataset_atom_file is not None:
-        candidates.append(Path(dataset_atom_file).expanduser())
-
-    for cand in candidates:
-        if cand.exists():
-            return cand
-
-    fallback = Path("/home/liuxuwei01/processed_output/omol25.json")
-    if fallback.exists():
-        logger.warning(
-            f"Configured atom_file not found ({candidates}); falling back to {fallback}"
-        )
-        return fallback
-
-    raise FileNotFoundError(
-        "Could not resolve atom_file for MOL inference. "
-        f"Checked: {[str(c) for c in candidates]} and fallback {fallback}"
-    )
-
-
-def collect_mol_files(mol_input, mol_pattern):
-    mol_path = Path(mol_input).expanduser()
-    if mol_path.is_file():
-        return [mol_path]
-    if not mol_path.is_dir():
-        raise FileNotFoundError(f"mol_input path not found: {mol_path}")
-
-    files = sorted([p for p in mol_path.glob(mol_pattern) if p.is_file()])
-    if not files:
-        files = sorted([p for p in mol_path.iterdir() if p.is_file() and p.suffix.lower() == ".mol"])
-    if not files:
-        raise FileNotFoundError(f"No .mol files found in directory: {mol_path}")
-    return files
-
-
-def open_text_maybe_compressed(path):
-    suffixes = "".join(path.suffixes).lower()
-    if suffixes.endswith(".lz4"):
-        import lz4.frame
-
-        return lz4.frame.open(path, mode="rt")
-    if suffixes.endswith(".xz"):
-        return lzma.open(path, mode="rt")
-    if suffixes.endswith(".gz"):
-        return gzip.open(path, mode="rt")
-    return path.open(mode="rt")
-
-
-def read_cube_density(path):
-    with open_text_maybe_compressed(path) as f:
-        f.readline()
-        f.readline()
-        line = f.readline().split()
-        if len(line) < 4:
-            raise ValueError(f"Invalid CUBE header (line 3) in {path}")
-        n_atom = int(line[0])
-        origin = np.array([float(x) for x in line[1:4]], dtype=np.float32)
-
-        shape = []
-        cell = np.zeros((3, 3), dtype=np.float32)
-        for i in range(3):
-            row = f.readline().split()
-            if len(row) < 4:
-                raise ValueError(f"Invalid CUBE axis line in {path}")
-            n, x, y, z = [float(s) for s in row[:4]]
-            shape.append(int(n))
-            cell[i] = np.array([x, y, z], dtype=np.float32)
-
-        x_coord = np.arange(shape[0], dtype=np.float32)[:, None] * cell[0][None, :]
-        y_coord = np.arange(shape[1], dtype=np.float32)[:, None] * cell[1][None, :]
-        z_coord = np.arange(shape[2], dtype=np.float32)[:, None] * cell[2][None, :]
-        grid_coord = (
-            x_coord.reshape(-1, 1, 1, 3)
-            + y_coord.reshape(1, -1, 1, 3)
-            + z_coord.reshape(1, 1, -1, 3)
-        ).reshape(-1, 3)
-        grid_coord = grid_coord + origin
-
-        atom_coord_ref = []
-        for _ in range(n_atom):
-            row = f.readline().split()
-            if len(row) < 5:
-                raise ValueError(f"Invalid CUBE atom line in {path}")
-            atom_coord_ref.append([float(row[2]), float(row[3]), float(row[4])])
-
-        n_grid = shape[0] * shape[1] * shape[2]
-        vals = []
-        for line in f:
-            parts = line.split()
-            if parts:
-                vals.extend(parts)
-        if len(vals) < n_grid:
-            raise ValueError(f"CUBE data too short in {path}: expect {n_grid}, got {len(vals)}")
-        density = np.array(vals[:n_grid], dtype=np.float32)
-
-    return (
-        paddle.to_tensor(density, dtype="float32"),
-        paddle.to_tensor(grid_coord, dtype="float32"),
-        {
-            "shape": shape,
-            "cell": paddle.to_tensor(cell, dtype="float32"),
-            "origin": paddle.to_tensor(origin, dtype="float32"),
-            "atom_coord_ref": np.asarray(atom_coord_ref, dtype=np.float32),
-        },
-    )
-
-
-def align_mol_atoms_to_cube(g, atom_coord_ref, sample_name, tol=0.05):
-    if atom_coord_ref is None:
-        return g
-    ref = np.asarray(atom_coord_ref, dtype=np.float32)
-    mol = g.pos.numpy().astype(np.float32)
-    if ref.ndim != 2 or ref.shape[1] != 3:
-        logger.warning(f"Invalid reference atom coordinates for {sample_name}, skip alignment")
-        return g
-    if mol.shape != ref.shape:
-        logger.warning(
-            f"Atom count mismatch for {sample_name} (mol={mol.shape[0]}, cube={ref.shape[0]}), "
-            "skip alignment"
-        )
-        return g
-
-    mol_center = mol.mean(axis=0)
-    ref_center = ref.mean(axis=0)
-    mol_c = mol - mol_center
-    ref_c = ref - ref_center
-    denom = float(np.sqrt((mol_c * mol_c).sum()))
-    numer = float(np.sqrt((ref_c * ref_c).sum()))
-    if denom < 1e-12 or numer < 1e-12:
-        return g
-
-    scale = numer / denom
-    aligned = mol_c * scale + ref_center
-    rms = float(np.sqrt(np.mean((aligned - ref) ** 2)))
-
-    # Typical unit mismatch is Angstrom->Bohr (about 1.8897).
-    # Apply alignment when scale obviously differs from 1.0 or residual is tiny after scaling.
-    if abs(scale - 1.0) > tol or rms < 1e-3:
-        g.pos = paddle.to_tensor(aligned, dtype="float32")
-        logger.info(
-            f"Aligned MOL coordinates to CUBE frame for {sample_name}: "
-            f"scale={scale:.6f} (A->Bohr~{ANG2BOHR:.6f}), rms={rms:.6e}"
-        )
-    else:
-        logger.info(
-            f"No coordinate rescale needed for {sample_name}: scale={scale:.6f}, rms={rms:.6e}"
-        )
-    return g
 
+from ppmat.predictor import FieldPredictor
+from ppmat.predictor.field_predictor import apply_predict_config
+from ppmat.utils.inference_cli import add_model_loading_arguments
+from ppmat.utils.inference_cli import validate_config_overrides
+from ppmat.utils.inference_cli import validate_model_loading_arguments
 
-def resolve_true_cube_for_mol(mol_path, true_cube_dir=None):
-    base = sanitize_base_name(mol_path.name)
-    base_density = f"{base[:-3]}Density" if base.endswith("Opt") else f"{base}Density"
-    roots = []
-    if true_cube_dir is not None:
-        roots.append(Path(true_cube_dir).expanduser())
-    roots.append(mol_path.parent)
+__all__ = ["FieldPredictor", "apply_predict_config", "build_parser", "main"]
 
-    stems = [base, f"{base}_true", base_density]
-    exts = [".cube", ".cub", ".cube.lz4", ".cube.gz", ".cube.xz", ".cub.lz4", ".cub.gz", ".cub.xz"]
-    name_candidates = []
-    for s in stems:
-        for ext in exts:
-            name_candidates.append(f"{s}{ext}")
 
-    seen = set()
-    uniq_candidates = []
-    for name in name_candidates:
-        if name not in seen:
-            uniq_candidates.append(name)
-            seen.add(name)
-
-    for root in roots:
-        if not root.exists():
-            continue
-        for name in uniq_candidates:
-            p = root / name
-            if p.is_file():
-                return p
-    return None
-
-
-def parse_mol_v2000(mol_path):
-    lines = mol_path.read_text(errors="replace").splitlines()
-    if len(lines) < 4:
-        raise ValueError(f"MOL file too short: {mol_path}")
-
-    counts = lines[3]
-    if "V3000" in counts.upper():
-        raise NotImplementedError(f"V3000 MOL is not supported yet: {mol_path}")
-
-    try:
-        n_atom = int(counts[:3])
-    except Exception:
-        parts = counts.split()
-        if len(parts) < 2:
-            raise ValueError(f"Failed to parse counts line in MOL file: {mol_path}")
-        n_atom = int(parts[0])
-
-    atom_start = 4
-    atom_end = atom_start + n_atom
-    if len(lines) < atom_end:
-        raise ValueError(f"Atom block incomplete in MOL file: {mol_path}")
-
-    coords = []
-    symbols = []
-    for line in lines[atom_start:atom_end]:
-        parts = line.split()
-        x = y = z = None
-        sym = None
-        if len(parts) >= 4:
-            try:
-                x, y, z = float(parts[0]), float(parts[1]), float(parts[2])
-                sym = parts[3]
-            except Exception:
-                x = y = z = None
-                sym = None
-        if x is None:
-            try:
-                x = float(line[0:10])
-                y = float(line[10:20])
-                z = float(line[20:30])
-                sym = line[31:34].strip()
-            except Exception as e:
-                raise ValueError(f"Failed to parse atom line in {mol_path}: {line}") from e
-
-        coords.append([x, y, z])
-        symbols.append(normalize_element_symbol(sym))
-
-    return np.asarray(coords, dtype=np.float32), symbols
-
-
-def build_mol_sample(mol_path, atom_name2idx, mol_grid_shape, mol_grid_padding):
-    atom_coord_np, atom_symbols = parse_mol_v2000(mol_path)
-
-    atom_type_idx = []
-    missing = set()
-    for sym in atom_symbols:
-        idx = atom_name2idx.get(sym)
-        if idx is None:
-            missing.add(sym)
-        else:
-            atom_type_idx.append(idx)
-    if missing:
-        raise ValueError(
-            f"Found atoms not covered by atom_file mapping in {mol_path}: {sorted(missing)}"
-        )
-
-    atom_type = paddle.to_tensor(atom_type_idx, dtype="int64")
-    atom_coord = paddle.to_tensor(atom_coord_np, dtype="float32")
-    g = Data(x=atom_type, pos=atom_coord)
-
-    shape = [int(s) for s in mol_grid_shape]
-    min_coord = atom_coord_np.min(axis=0)
-    max_coord = atom_coord_np.max(axis=0)
-    span = np.maximum(max_coord - min_coord, np.array([1e-3, 1e-3, 1e-3], dtype=np.float32))
-    axis_len = span + 2.0 * float(mol_grid_padding)
-    center = 0.5 * (min_coord + max_coord)
-    origin = center - 0.5 * axis_len
-
-    x = np.linspace(origin[0], origin[0] + axis_len[0], num=shape[0], endpoint=False, dtype=np.float32)
-    y = np.linspace(origin[1], origin[1] + axis_len[1], num=shape[1], endpoint=False, dtype=np.float32)
-    z = np.linspace(origin[2], origin[2] + axis_len[2], num=shape[2], endpoint=False, dtype=np.float32)
-    grid = np.stack(np.meshgrid(x, y, z, indexing="ij"), axis=-1).reshape(-1, 3).astype(np.float32)
-    grid_coord = paddle.to_tensor(grid, dtype="float32")
-
-    cell = np.diag(axis_len.astype(np.float32))
-    info = {
-        "shape": shape,
-        "cell": paddle.to_tensor(cell, dtype="float32"),
-        "origin": paddle.to_tensor(origin.astype(np.float32), dtype="float32"),
-        "file_name": mol_path.name,
-    }
-
-    return g, None, grid_coord, info
-
-
-def sanitize_base_name(sample_name):
-    base_name = Path(sample_name).name
-    for suf in [".lz4", ".zst", ".gz"]:
-        if base_name.endswith(suf):
-            base_name = base_name[: -len(suf)]
-    for suf in [".cube", ".CHGCAR", ".json", ".mol"]:
-        if base_name.endswith(suf):
-            base_name = base_name[: -len(suf)]
-    return base_name
-
-
-def prepare_info_cube(info, grid_coord):
-    info_cube = {}
-    shape = info.get("shape")
-    cell = info.get("cell")
-    origin = info.get("origin", None)
-    grid_np_full = grid_coord.detach().cpu().numpy()
-
-    if shape is not None and len(shape) == 3:
-        try:
-            shape_i = [int(s) for s in shape]
-            grid_view = grid_np_full.reshape(shape_i[0], shape_i[1], shape_i[2], 3)
-            origin_np = grid_view[0, 0, 0]
-            step_x = (
-                grid_view[1, 0, 0] - grid_view[0, 0, 0]
-                if shape_i[0] > 1
-                else np.zeros(3, dtype=np.float32)
-            )
-            step_y = (
-                grid_view[0, 1, 0] - grid_view[0, 0, 0]
-                if shape_i[1] > 1
-                else np.zeros(3, dtype=np.float32)
-            )
-            step_z = (
-                grid_view[0, 0, 1] - grid_view[0, 0, 0]
-                if shape_i[2] > 1
-                else np.zeros(3, dtype=np.float32)
-            )
-            cell_from_grid = np.stack(
-                [step_x * shape_i[0], step_y * shape_i[1], step_z * shape_i[2]], axis=0
-            )
-        except Exception:
-            origin_np = None
-            cell_from_grid = None
-    else:
-        origin_np = None
-        cell_from_grid = None
-
-    if shape is not None:
-        info_cube["shape"] = [int(s) for s in shape]
-    if cell is not None:
-        if hasattr(cell, "numpy"):
-            info_cube["cell"] = cell.numpy()
-        else:
-            info_cube["cell"] = np.array(cell, dtype=np.float32)
-    if cell_from_grid is not None:
-        info_cube["cell"] = cell_from_grid
-    if origin is not None:
-        if hasattr(origin, "numpy"):
-            info_cube["origin"] = origin.numpy()
-        else:
-            info_cube["origin"] = np.array(origin, dtype=np.float32)
-    if origin_np is not None:
-        info_cube["origin"] = origin_np
-    return info_cube
-
-
-def main():
-    parser = argparse.ArgumentParser(description="InfGCN electron density inference")
-    parser.add_argument(
-        "--config",
-        default="electronic_structure/configs/infgcn/infgcn_qm9.yaml",
-        help="Path to config yaml",
-    )
-    parser.add_argument(
-        "--checkpoint",
-        default="output/infgcn_qm9_best/infgcn_qm9.pdparams",
-        help="Checkpoint (.pdparams) to load",
-    )
+def build_parser():
+    parser = argparse.ArgumentParser(description="Electron density field inference")
+    add_model_loading_arguments(parser)
     parser.add_argument(
         "--split",
-        default="test",
+        default=None,
         choices=["train", "validation", "test"],
         help="Dataset split to sample from",
     )
     parser.add_argument(
         "--index",
-        default=0,
+        default=None,
         type=int,
         help="Index within the chosen split",
     )
@@ -718,33 +55,37 @@ def main():
     )
     parser.add_argument(
         "--output_dir",
-        default="./results",
+        default=None,
         help="Directory to store predictions/visualizations",
     )
     parser.add_argument(
         "--grid_batch_size",
-        default=4096,
+        default=None,
         type=int,
         help="Number of grid points per forward pass",
     )
     parser.add_argument(
         "--skip_vis",
         action="store_true",
+        default=None,
         help="Skip writing/visualizing density plots",
     )
     parser.add_argument(
         "--save_true_cube",
         action="store_true",
+        default=None,
         help="Save reference (DFT) electron density as a cube file",
     )
     parser.add_argument(
         "--save_pred_cube",
         action="store_true",
+        default=None,
         help="Save predicted electron density as a cube file",
     )
     parser.add_argument(
         "--save_html",
         action="store_true",
+        default=None,
         help="Save Plotly figures as interactive HTML (in addition to PNG)",
     )
     parser.add_argument(
@@ -755,26 +96,30 @@ def main():
     parser.add_argument(
         "--show_plot",
         action="store_true",
+        default=None,
         help="Display plotly figures inline (requires kaleido)",
     )
     parser.add_argument(
         "--mol_input",
         default=None,
-        help="Path to a .mol file or a directory of .mol files for direct structure inference",
+        help=(
+            "Path to a .mol file or a directory of .mol files for direct "
+            "structure inference"
+        ),
     )
     parser.add_argument(
         "--mol_pattern",
-        default="*.mol",
+        default=None,
         help="Glob pattern when --mol_input is a directory",
     )
     parser.add_argument(
         "--mol_grid_shape",
-        default="80,80,80",
+        default=None,
         help="Grid shape for MOL inference, e.g. '80' or '80,80,80'",
     )
     parser.add_argument(
         "--mol_grid_padding",
-        default=6.0,
+        default=None,
         type=float,
         help="Padding (Angstrom) around molecular coordinates for MOL grid generation",
     )
@@ -786,279 +131,22 @@ def main():
             "Expected names: <mol_basename>.cube or <mol_basename>_true.cube"
         ),
     )
-    args = parser.parse_args()
-
-    set_random_seed(42)
-
-    cfg = OmegaConf.load(args.config)
-    cfg = OmegaConf.to_container(cfg, resolve=True)
-
-    split_key = "val" if args.split == "validation" else args.split
-    ds_cfg_full = cfg["Dataset"][split_key]["dataset"]
-    dataset_cfg = ds_cfg_full.get("__init_params__", {})
-    dataset_params = copy.deepcopy(dataset_cfg)
-    dataset_params["split"] = args.split
-    if args.data_root is not None:
-        dataset_params["root"] = args.data_root
-    if args.split_file is not None:
-        dataset_params["split_file"] = args.split_file
-    if args.atom_file is not None:
-        dataset_params["atom_file"] = args.atom_file
-
-    use_mol_mode = args.mol_input is not None
-
-    dataset = None
-    cube_writer = None
-    idx2atom_num = None
-    atom_name2idx = None
-    mol_files = []
-    mol_grid_shape = None
-
-    if use_mol_mode:
-        atom_file_path = resolve_atom_file_for_mol(
-            args.atom_file,
-            dataset_params.get("atom_file"),
-        )
-        atom_name2idx, idx2atom_num = load_atom_mapping(atom_file_path)
-        mol_files = collect_mol_files(args.mol_input, args.mol_pattern)
-        mol_grid_shape = parse_grid_shape(args.mol_grid_shape)
-        cube_writer = lambda f, a, c, d, i: write_cube_generic(
-            f, a, c, d, i, idx2atom_num
-        )
-        logger.info(
-            f"MOL mode enabled: {len(mol_files)} file(s), atom_file={atom_file_path}, "
-            f"grid_shape={mol_grid_shape}, padding={args.mol_grid_padding}, "
-            f"true_cube_dir={args.mol_true_cube_dir}"
-        )
-    else:
-        dataset_cls_name = ds_cfg_full.get("__class_name__", "DensityDataset")
-        dataset_cls_map = {
-            "DensityDataset": DensityDataset,
-            "SmallDensityDataset": SmallDensityDataset,
-        }
-        if dataset_cls_name not in dataset_cls_map:
-            raise ValueError(f"Unsupported dataset class {dataset_cls_name}")
-        dataset = dataset_cls_map[dataset_cls_name](**dataset_params)
-        cube_writer = getattr(dataset, "write_cube", None)
-        idx2atom_num = getattr(dataset, "idx2atom_num", None)
-        if cube_writer is None:
-            if isinstance(dataset, SmallDensityDataset):
-                # Atom order in SmallDensityDataset: C=0, H=1, O=2
-                idx2atom_num = np.array([6, 1, 8], dtype=np.int64)
-                cube_writer = lambda f, a, c, d, i: write_cube_generic(
-                    f, a, c, d, i, idx2atom_num
-                )
-            else:
-                cube_writer = lambda *args, **kwargs: (_ for _ in ()).throw(
-                    AttributeError("Cube writer not available for this dataset")
-                )
-        if args.index >= len(dataset):
-            raise IndexError(
-                f"Index {args.index} exceeds dataset size {len(dataset)} for split {args.split}"
-            )
+    return parser
 
-    device = "gpu" if paddle.is_compiled_with_cuda() else "cpu"
-    paddle.set_device(device)
-    logger.info(f"Running inference on device: {device}")
 
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    cube_dir = Path(args.cube_dir) if args.cube_dir is not None else output_dir
-    cube_dir.mkdir(parents=True, exist_ok=True)
-
-    logger.info(f"Loading the pretrained model from {args.checkpoint}")
-    model = get_pretrained_model(args.config, args.checkpoint)
-    logger.info("Model loaded successfully.")
-
-    if use_mol_mode:
-        sample_iter = tqdm(mol_files, desc="MOL inference")
-    else:
-        sample_iter = [args.index]
-
-    for sample_item in sample_iter:
-        if use_mol_mode:
-            mol_path = sample_item
-            g, density, grid_coord, info = build_mol_sample(
-                mol_path,
-                atom_name2idx,
-                mol_grid_shape,
-                args.mol_grid_padding,
-            )
-            true_cube_path = resolve_true_cube_for_mol(mol_path, args.mol_true_cube_dir)
-            if true_cube_path is not None:
-                try:
-                    density, grid_coord, info_ref = read_cube_density(true_cube_path)
-                    g = align_mol_atoms_to_cube(g, info_ref.get("atom_coord_ref"), mol_path.name)
-                    info = dict(info_ref)
-                    info["file_name"] = mol_path.name
-                    info["true_cube_file"] = str(true_cube_path)
-                    logger.info(f"Using reference cube for {mol_path.name}: {true_cube_path}")
-                except Exception as e:
-                    logger.warning(
-                        f"Failed to read reference cube for {mol_path.name} at {true_cube_path}: {e}"
-                    )
-            sample_name = info.get("file_name", mol_path.name)
-        else:
-            sample_name = f"{args.split}_{args.index}"
-            g, density, grid_coord, info = dataset[args.index]
-            sample_name = info.get("file_name", sample_name)
-
-        g.batch = paddle.zeros_like(g.x)
-        g = g.to(device)
-        if density is not None:
-            density = density.to(device)
-        grid_coord = grid_coord.to(device)
-
-        logger.info(f"Starting prediction for sample: {sample_name}")
-        preds, loss, mae = inference_model(
-            model,
-            g,
-            density,
-            grid_coord[None],
-            [info],
-            grid_batch_size=args.grid_batch_size,
-        )
-        if loss is not None and mae is not None:
-            logger.info(
-                f"Prediction completed for {sample_name}, "
-                f"Loss: {float(loss):.6f}, MAE: {float(mae):.6f}"
-            )
-        else:
-            logger.info(f"Prediction completed for {sample_name} (no reference density)")
-
-        sample_tag = sanitize_base_name(sample_name)
-
-        if args.save_true_cube or args.save_pred_cube:
-            atom_type_np = g.x.detach().cpu().numpy()
-            atom_coord_np = g.pos.detach().cpu().numpy()
-            info_cube = prepare_info_cube(info, grid_coord)
-
-            if args.save_true_cube:
-                if density is None:
-                    logger.warning(
-                        f"Skipping true cube for {sample_name}: no reference density available"
-                    )
-                else:
-                    true_cube_path = cube_dir / f"{sample_tag}_true.cube"
-                    with true_cube_path.open("w") as f:
-                        cube_writer(
-                            f,
-                            atom_type_np,
-                            atom_coord_np,
-                            density.detach().cpu().numpy(),
-                            info_cube,
-                        )
-                    logger.info(f"Saved reference density cube to: {true_cube_path}")
-
-            if args.save_pred_cube:
-                pred_cube_path = cube_dir / f"{sample_tag}_pred.cube"
-                with pred_cube_path.open("w") as f:
-                    cube_writer(
-                        f,
-                        atom_type_np,
-                        atom_coord_np,
-                        preds.detach().cpu().numpy(),
-                        info_cube,
-                    )
-                logger.info(f"Saved predicted density cube to: {pred_cube_path}")
-
-        if not args.skip_vis:
-            grid_np = grid_coord.detach().cpu().numpy()
-            preds_np = preds.detach().cpu().numpy()
-            shape = info.get("shape")
-            atom_type = g.x.detach().cpu().numpy()
-            atom_coord = g.pos.detach().cpu().numpy()
-
-            if density is not None:
-                density_np = density.detach().cpu().numpy()
-                diff_np = density_np - preds_np
-                grid_vis, (density_vis, diff_vis, preds_vis), did_downsample, stride = (
-                    maybe_downsample_volume(
-                        grid_np,
-                        [density_np, diff_np, preds_np],
-                        shape if shape is None else [int(s) for s in shape],
-                    )
-                )
-                if did_downsample:
-                    logger.warning(
-                        f"Downsampled volume grid from {grid_np.shape[0]} to {grid_vis.shape[0]} "
-                        f"points for visualization (stride={stride}) to keep HTML output responsive."
-                    )
-
-                logger.info("Visualizing the DFT electron density")
-                fig = draw_volume(
-                    grid_vis,
-                    density_vis,
-                    atom_type,
-                    atom_coord,
-                    isomin=0.05,
-                    isomax=3.5,
-                    surface_count=5,
-                    title="DFT electron density",
-                )
-                true_density_path = output_dir / f"{sample_tag}_true_density.png"
-                safe_write_image(fig, true_density_path, show_plot=args.show_plot)
-                if args.save_html:
-                    fig.write_html(output_dir / f"{sample_tag}_true_density.html")
-
-                logger.info("Visualizing electron density difference")
-                fig = draw_volume(
-                    grid_vis,
-                    diff_vis,
-                    atom_type,
-                    atom_coord,
-                    isomin=-0.06,
-                    isomax=0.06,
-                    surface_count=4,
-                    title="Electron Density Difference",
-                )
-                diff_density_path = output_dir / f"{sample_tag}_diff_density.png"
-                safe_write_image(fig, diff_density_path, show_plot=args.show_plot)
-                if args.save_html:
-                    fig.write_html(output_dir / f"{sample_tag}_diff_density.html")
-
-                logger.info("Visualizing predicted electron density")
-                fig = draw_volume(
-                    grid_vis,
-                    preds_vis,
-                    atom_type,
-                    atom_coord,
-                    isomin=0.05,
-                    isomax=3.5,
-                    surface_count=5,
-                    title="Predicted Electron Density",
-                )
-                pred_density_path = output_dir / f"{sample_tag}_pred_density.png"
-                safe_write_image(fig, pred_density_path, show_plot=args.show_plot)
-                if args.save_html:
-                    fig.write_html(output_dir / f"{sample_tag}_pred_density.html")
-            else:
-                grid_vis, (preds_vis,), did_downsample, stride = maybe_downsample_volume(
-                    grid_np,
-                    [preds_np],
-                    shape if shape is None else [int(s) for s in shape],
-                )
-                if did_downsample:
-                    logger.warning(
-                        f"Downsampled volume grid from {grid_np.shape[0]} to {grid_vis.shape[0]} "
-                        f"points for visualization (stride={stride}) to keep HTML output responsive."
-                    )
-
-                logger.info("Visualizing predicted electron density")
-                fig = draw_volume(
-                    grid_vis,
-                    preds_vis,
-                    atom_type,
-                    atom_coord,
-                    isomin=0.05,
-                    isomax=3.5,
-                    surface_count=5,
-                    title="Predicted Electron Density",
-                )
-                pred_density_path = output_dir / f"{sample_tag}_pred_density.png"
-                safe_write_image(fig, pred_density_path, show_plot=args.show_plot)
-                if args.save_html:
-                    fig.write_html(output_dir / f"{sample_tag}_pred_density.html")
+def main():
+    parser = build_parser()
+    args, config_overrides = parser.parse_known_args()
+    validate_model_loading_arguments(parser, args)
+    validate_config_overrides(parser, config_overrides)
+    predictor = FieldPredictor(
+        model_name=args.model_name,
+        weights_name=args.weights_name,
+        config_path=args.config_path,
+        checkpoint_path=args.checkpoint_path,
+        config_overrides=config_overrides,
+    )
+    predictor.predict(args)
 
 
 if __name__ == "__main__":
diff --git a/electronic_structure/train.py b/electronic_structure/train.py
index 5ba2db62..e1cdc98e 100644
--- a/electronic_structure/train.py
+++ b/electronic_structure/train.py
@@ -13,8 +13,6 @@
 # limitations under the License.
 
 import argparse
-import datetime
-import math
 import os
 import os.path as osp
 
@@ -31,6 +29,7 @@
 from ppmat.utils import logger
 from ppmat.utils import misc
 from ppmat.utils.eager_comp_setting import setting_eager_mode
+from ppmat.utils.io import append_timestamp_to_output_dir
 
 
 def read_independent_dataloader_config(config):
@@ -93,10 +92,7 @@ def read_independent_dataloader_config(config):
     misc.set_random_seed(seed)
     logger.info(f"Set random seed to {seed}")
 
-    # add timestamp to output_dir
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    base_output_dir = config["Trainer"]["output_dir"]
-    config["Trainer"]["output_dir"] = f"{base_output_dir}_t_{timestamp}_s_{seed}"
+    append_timestamp_to_output_dir(config)
 
     # save config to output_dir, only rank 0 process will do this
     if dist.get_rank() == 0:
diff --git a/interatomic_potentials/train.py b/interatomic_potentials/train.py
index 09fd8880..04e607ed 100644
--- a/interatomic_potentials/train.py
+++ b/interatomic_potentials/train.py
@@ -29,6 +29,7 @@
 from ppmat.utils import logger
 from ppmat.utils import misc
 from ppmat.utils.eager_comp_setting import setting_eager_mode
+from ppmat.utils.io import append_timestamp_to_output_dir
 
 if dist.get_world_size() > 1:
     fleet.init(is_collective=True)
@@ -49,6 +50,9 @@
     cli_config = OmegaConf.from_dotlist(dynamic_args)
     config = OmegaConf.merge(config, cli_config)
 
+    seed = config["Trainer"].get("seed", 42)
+    append_timestamp_to_output_dir(config)
+
     # save config to output_dir, only rank 0 process will do this
     if dist.get_rank() == 0:
         os.makedirs(config["Trainer"]["output_dir"], exist_ok=True)
@@ -63,7 +67,6 @@
     logger.info(f"Logger saved to {logger_path}")
 
     # set random seed
-    seed = config["Trainer"].get("seed", 42)
     misc.set_random_seed(seed)
     logger.info(f"Set random seed to {seed}")
 
diff --git a/ppmat/datasets/msd_nmr_dataset.py b/ppmat/datasets/msd_nmr_dataset.py
index 2406bb5e..3bd7ae89 100644
--- a/ppmat/datasets/msd_nmr_dataset.py
+++ b/ppmat/datasets/msd_nmr_dataset.py
@@ -1286,9 +1286,12 @@ def __init__(self, dataloaders, cfg, recompute_statistics=False):
             self.edge_types = dataloaders.edge_counts()
             self.valency_distribution = dataloaders.valency_count(self.max_n_nodes)
 
-        self.train_smiles = get_train_smiles(
-            cfg, dataloaders.train_dataloader, self, evaluate_dataset=False
-        )
+        if cfg.get("load_train_smiles", True):
+            self.train_smiles = get_train_smiles(
+                cfg, dataloaders.train_dataloader, self, evaluate_dataset=False
+            )
+        else:
+            self.train_smiles = None
 
     def complete_infos(self, n_nodes, node_types):
         self.input_dims = None
@@ -1341,39 +1344,50 @@ def compute_input_output_dims(
         }
 
 
-def get_train_smiles(cfg, dataloader, dataset_infos, evaluate_dataset=False):
-    if evaluate_dataset:
-        assert (
-            dataset_infos is not None
-        ), "If wanting to evaluate dataset, need to pass dataset_infos"
-    if not osp.exists(cfg["datadir"]):
+def _get_msd_nmr_subdataset_name(data_flag: str):
+    if data_flag == "n<15":
+        return "msd_nmr_nless15"
+    if data_flag == "n<20":
+        return "msd_nmr_nless20"
+    if data_flag == "n<30":
+        return "msd_nmr_nless30"
+    if data_flag == "n<35":
+        return "msd_nmr_nless35"
+    raise ValueError(
+        f"Unknown data_flag: {data_flag}. Expected one of "
+        f"{'n<15', 'n<20', 'n<30', 'n<35'}."
+    )
+
+
+def _resolve_msd_nmr_subset_dir(datadir: str, data_flag: str):
+    subdataset_name = _get_msd_nmr_subdataset_name(data_flag)
+    if not osp.exists(datadir):
         logger.message(
             "The dataset directory is not found. Will save it to default path now."
         )
         root_path = download.get_datasets_path_from_url(
             MSDnmrDataset.url, MSDnmrDataset.md5
         )
-        path = osp.join(root_path, MSDnmrDataset.name, osp.basename(cfg["datadir"]))
-        if cfg["data_flag"] == "n<15":
-            subdataset_name = "msd_nmr_nless15"
-        elif cfg["data_flag"] == "n<20":
-            subdataset_name = "msd_nmr_nless20"
-        elif cfg["data_flag"] == "n<30":
-            subdataset_name = "msd_nmr_nless30"
-        elif cfg["data_flag"] == "n<35":
-            subdataset_name = "msd_nmr_nless35"
-        else:
-            raise ValueError(
-                f"Unknown data_flag: {cfg['data_flag']}. Expected one of "
-                f"{'n<15', 'n<20', 'n<30', 'n<35'}."
-            )
-        path = osp.join(root_path, MSDnmrDataset.name, subdataset_name)
+        return osp.join(root_path, MSDnmrDataset.name, subdataset_name)
+
+    if osp.exists(osp.join(datadir, "train.csv")):
+        return datadir
+
+    return osp.join(datadir, subdataset_name)
+
+
+def get_train_smiles(cfg, dataloader, dataset_infos, evaluate_dataset=False):
+    if evaluate_dataset:
+        assert (
+            dataset_infos is not None
+        ), "If wanting to evaluate dataset, need to pass dataset_infos"
 
     remove_h = cfg["build_graph_cfg"]["__init_params__"]["remove_h"]
     atom_decoder = dataset_infos.atom_decoder
 
     smiles_file_name = "train_smiles_no_h.npy" if remove_h else "train_smiles_h.npy"
-    smiles_path = os.path.join(path + "_cache", "train", smiles_file_name)
+    dataset_dir = _resolve_msd_nmr_subset_dir(cfg["datadir"], cfg["data_flag"])
+    smiles_path = os.path.join(dataset_dir + "_cache", "train", smiles_file_name)
     if os.path.exists(smiles_path):
         logger.message("Dataset smiles were found")
         train_smiles = np.load(smiles_path)
diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py
index 32f1c729..1a86d497 100644
--- a/ppmat/models/__init__.py
+++ b/ppmat/models/__init__.py
@@ -14,8 +14,6 @@
 
 import copy
 import inspect
-import os
-import os.path as osp
 from typing import Any
 from typing import Dict
 from typing import Optional
@@ -36,18 +34,20 @@
 from ppmat.models.diffnmr.diffnmr import MolecularGraphFormer
 from ppmat.models.diffnmr.diffnmr import NMRNetCLIP
 from ppmat.models.dimenetpp.dimenetpp import DimeNetPlusPlus
+from ppmat.models.infgcn.infgcn import InfGCN
+from ppmat.models.mateno.mateno import MatENO
 from ppmat.models.mattergen.mattergen import MatterGen
 from ppmat.models.mattergen.mattergen import MatterGenWithCondition
 from ppmat.models.mattersim.m3gnet import M3GNet
 from ppmat.models.mattersim.m3gnet_graph_converter import M3GNetGraphConvertor
 from ppmat.models.megnet.megnet import MEGNetPlus
-from ppmat.models.infgcn.infgcn import InfGCN
-from ppmat.models.mateno.mateno import MatENO
 from ppmat.models.sfin.sfin import SFIN
 from ppmat.models.spherenet.spherenet import SphereNet
 from ppmat.utils import download
 from ppmat.utils import logger
 from ppmat.utils import save_load
+from ppmat.utils.model_package import get_model_config_path
+from ppmat.utils.model_package import resolve_model_package_dir
 
 __all__ = [
     "iComformer",
@@ -78,21 +78,21 @@
 # Warning: The key of the dictionary must be consistent with the file name of the value
 MODEL_REGISTRY = {
     "comformer_mp2018_train_60k_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_mp2018_train_60k_e_form.zip",
-    "comformer_mp2018_train_60k_band_gap": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_mp2018_train_60k_band_gap.zip",
+    "comformer_mp2018_train_60k_band_gap": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/comformer/comformer_mp2018_train_60k_band_gap.zip",
     "comformer_mp2018_train_60k_G": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_mp2018_train_60k_G.zip",
     "comformer_mp2018_train_60k_K": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_mp2018_train_60k_K.zip",
     "comformer_mp2024_train_130k_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_mp2024_train_130k_e_form.zip",
-    "comformer_jarvis_dft_2d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_jarvis_dft_2d_e_form.zip",
-    "comformer_jarvis_dft_3d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_jarvis_dft_3d_e_form.zip",
-    "comformer_jarvis_alex_pbe_2d_all_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/comformer/comformer_jarvis_alex_pbe_2d_all_e_form.zip",
+    "comformer_jarvis_dft_2d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/comformer/comformer_jarvis_dft_2d_e_form.zip",
+    "comformer_jarvis_dft_3d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/comformer/comformer_jarvis_dft_3d_e_form.zip",
+    "comformer_jarvis_alex_pbe_2d_all_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/comformer/comformer_jarvis_alex_pbe_2d_all_e_form.zip",
     "megnet_mp2018_train_60k_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_mp2018_train_60k_e_form.zip",
     "megnet_mp2018_train_60k_band_gap": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_mp2018_train_60k_band_gap.zip",
     "megnet_mp2018_train_60k_G": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_mp2018_train_60k_G.zip",
     "megnet_mp2018_train_60k_K": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_mp2018_train_60k_K.zip",
     "megnet_mp2024_train_130k_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_mp2024_train_130k_e_form.zip",
-    "megnet_jarvis_dft_2d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_jarvis_dft_2d_e_form.zip",
-    "megnet_jarvis_dft_3d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_jarvis_dft_3d_e_form.zip",
-    "megnet_jarvis_alex_pbe_2d_all_e_form": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/property_prediction/megnet/megnet_jarvis_alex_pbe_2d_all_e_form.zip",
+    "megnet_jarvis_dft_2d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/megnet/megnet_jarvis_dft_2d_e_form.zip",
+    "megnet_jarvis_dft_3d_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/megnet/megnet_jarvis_dft_3d_e_form.zip",
+    "megnet_jarvis_alex_pbe_2d_all_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/megnet/megnet_jarvis_alex_pbe_2d_all_e_form.zip",
     "diffcsp_mp20": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/structure_generation/diffcsp/diffcsp_mp20.zip",
     "mattergen_mp20": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/structure_generation/mattergen/mattergen_mp20.zip",
     "mattergen_mp20_chemical_system": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/structure_generation/mattergen/mattergen_mp20_chemical_system.zip",
@@ -108,14 +108,14 @@
     "mattergen_alex_mp20_chemical_system_energy_above_hull": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/structure_generation/mattergen/mattergen_alex_mp20_chemical_system_energy_above_hull.zip",
     "mattergen_alex_mp20_dft_mag_density_hhi_score": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/structure_generation/mattergen/mattergen_alex_mp20_dft_mag_density_hhi_score.zip",
     "chgnet_mptrj": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/interatomic_potentials/chgnet/chgnet_mptrj.zip",
-    "dimenetpp_mp2018_train_60k_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet%2B%2B/dimenetpp_mp2018_train_60k_e_form.zip",
-    "dimenetpp_mp2018_train_60k_band_gap": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet%2B%2B/dimenetpp_mp2018_train_60k_band_gap.zip",
-    "dimenetpp_mp2018_train_60k_G": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet%2B%2B/dimenetpp_mp2018_train_60k_G.zip",
-    "dimenetpp_mp2018_train_60k_K": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet%2B%2B/dimenetpp_mp2018_train_60k_K.zip",
+    "dimenetpp_mp2018_train_60k_e_form": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet++/dimenetpp_mp2018_train_60k_e_form.zip",
+    "dimenetpp_mp2018_train_60k_band_gap": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet++/dimenetpp_mp2018_train_60k_band_gap.zip",
+    "dimenetpp_mp2018_train_60k_G": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet++/dimenetpp_mp2018_train_60k_G.zip",
+    "dimenetpp_mp2018_train_60k_K": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/dimenet++/dimenetpp_mp2018_train_60k_K.zip",
     "mattersim_1M": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/interatomic_potentials/mattersim/mattersim_1M.zip",
     "mattersim_5M": "https://paddle-org.bj.bcebos.com/paddlematerial/checkpoints/interatomic_potentials/mattersim/mattersim_5M.zip",
-    "mattergen_ml2ddb": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb.zip",
-    "mattergen_ml2ddb_chemical_system": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_chemical_system.zip",
+    "mattergen_ml2ddb": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/MatterGen/mattergen_ml2ddb.zip",
+    "mattergen_ml2ddb_chemical_system": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/MatterGen/mattergen_ml2ddb_chemical_system.zip",
     "mattergen_ml2ddb_space_group": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_space_group.zip",
     "sfin_haadf_enhance": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_enhance.zip",
     "sfin_haadf_detect": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_detect.zip",
@@ -141,6 +141,16 @@
     "spherenet_md17_salicylic": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_salicylic.zip",
     "spherenet_md17_toluene": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_toluene.zip",
     "spherenet_md17_uracil": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_uracil.zip",
+    "infgcn_md17_benzene": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_benzene.zip",
+    "infgcn_md17_ethane": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_ethane.zip",
+    "infgcn_md17_ethanol": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_ethanol.zip",
+    "infgcn_md17_malonaldehyde": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_malonaldehyde.zip",
+    "infgcn_md17_phenol": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_phenol.zip",
+    "infgcn_md17_resorcinol": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_md17_resorcinol.zip",
+    "infgcn_mp": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_mp.zip",
+    "infgcn_omol25_mc_5k_trimmed": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_omol25_mc_5k_trimmed.zip",
+    "infgcn_qm9": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/electronic_structure/infgcn/infgcn_qm9.zip",
+    "diffnmr_msdnmr_nless15": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/diffnmr_msdnmr_nless15.zip",
 }
 
 
@@ -227,26 +237,10 @@ def build_model(
 
 
 def build_model_from_name(model_name: str, weights_name: Optional[str] = None):
-    path = download.get_weights_path_from_url(MODEL_REGISTRY[model_name])
-    path = osp.join(path, model_name)
+    extracted_path = download.get_weights_path_from_url(MODEL_REGISTRY[model_name])
+    path = resolve_model_package_dir(model_name, extracted_path)
     logger.info(f"Save model and configuration files in path: {path}")
-    config_path = osp.join(path, f"{model_name}.yaml")
-    if not osp.exists(config_path):
-        logger.warning(
-            f"Config file not found: {config_path}, try find other yaml files."
-        )
-        file_list = os.listdir(path)
-        find_list = []
-        for file in file_list:
-            if file.endswith(".yaml") or file.endswith(".yml"):
-                find_list.append(osp.join(path, file))
-        if len(find_list) == 1:
-            config_path = find_list[0]
-        else:
-            raise ValueError(
-                f"Multiple yaml files found: {find_list}, must be only one"
-            )
-        logger.warning(f"Find config file: {config_path}, using this file.")
+    config_path = get_model_config_path(model_name, path)
 
     config = OmegaConf.load(config_path)
     config = OmegaConf.to_container(config, resolve=True)
diff --git a/ppmat/models/diffcsp/diffcsp.py b/ppmat/models/diffcsp/diffcsp.py
index 05721eb1..0cd5e696 100644
--- a/ppmat/models/diffcsp/diffcsp.py
+++ b/ppmat/models/diffcsp/diffcsp.py
@@ -340,6 +340,8 @@ def forward(
 
 
 class DiffCSP(paddle.nn.Layer):
+    supports_num_atoms_sampling = False
+
     """Crystal Structure Prediction by Joint Equivariant Diffusion
 
     https://arxiv.org/abs/2309.04475
diff --git a/ppmat/models/diffnmr/diffnmr.py b/ppmat/models/diffnmr/diffnmr.py
index c3ad924e..6baf19f7 100644
--- a/ppmat/models/diffnmr/diffnmr.py
+++ b/ppmat/models/diffnmr/diffnmr.py
@@ -451,11 +451,12 @@ def __init__(
         extra_features,
         domain_features,
         clip,
-        connector_cfg = None,
+        connector_cfg=None,
     ) -> None:
         super().__init__()
 
         # configure general variables settings
+        self.conditioning_mode = "spectrum"
         self.T = diffmodel_cfg["diffusion_steps"]
 
         # configure datasets inter-varibles
diff --git a/ppmat/models/infgcn/infgcn.py b/ppmat/models/infgcn/infgcn.py
index 047408aa..4adfb905 100644
--- a/ppmat/models/infgcn/infgcn.py
+++ b/ppmat/models/infgcn/infgcn.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import paddle
-from paddle_scatter import scatter
 
 from ppmat.models.common.e3nn import o3
 from ppmat.models.common.e3nn.math import soft_one_hot_linspace
@@ -26,6 +25,7 @@
 
 from ppmat.datasets.graph_utils.infgcn_graph_utils import radius
 from ppmat.datasets.graph_utils.infgcn_graph_utils import radius_graph
+from ppmat.utils.scatter import scatter
 
 
 class GCNLayer(paddle.nn.Layer):
diff --git a/ppmat/models/mateno/mateno.py b/ppmat/models/mateno/mateno.py
index ad67e11c..592e0a16 100644
--- a/ppmat/models/mateno/mateno.py
+++ b/ppmat/models/mateno/mateno.py
@@ -13,9 +13,9 @@
 # limitations under the License.
 
 import paddle
-from paddle_scatter import scatter
 
 from ppmat.datasets.graph_utils.infgcn_graph_utils import radius, radius_graph
+from ppmat.utils.scatter import scatter
 
 from ppmat.models.common.e3nn import o3
 from ppmat.models.common.e3nn.math import soft_one_hot_linspace
diff --git a/ppmat/predictor/__init__.py b/ppmat/predictor/__init__.py
index 34a0df0d..e43b5d5c 100644
--- a/ppmat/predictor/__init__.py
+++ b/ppmat/predictor/__init__.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 
 from ppmat.predictor.base import BasePredictor
+from ppmat.predictor.field_predictor import FieldPredictor
 
 __all__ = [
     "BasePredictor",
+    "FieldPredictor",
 ]
diff --git a/ppmat/predictor/field_io.py b/ppmat/predictor/field_io.py
new file mode 100644
index 00000000..cffbe6a5
--- /dev/null
+++ b/ppmat/predictor/field_io.py
@@ -0,0 +1,151 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gzip
+import lzma
+import time
+from pathlib import Path
+
+import numpy as np
+import paddle
+
+
+def write_cube(fileobj, atom_type, atom_coord, density, info, idx2atom_num=None):
+    """Write a Gaussian CUBE file for an electron-density prediction."""
+
+    fileobj.write("Cube file written on " + time.strftime("%c"))
+    fileobj.write("\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")
+    cell = info["cell"]
+    shape = info["shape"]
+    origin = info.get("origin", np.zeros(3, dtype=np.float32))
+    fileobj.write("{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(len(atom_type), *origin))
+    for size, vector in zip(shape, cell):
+        step = vector / size
+        fileobj.write("{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(size, *step))
+    for atom, (x_coord, y_coord, z_coord) in zip(atom_type, atom_coord):
+        atomic_num = (
+            int(idx2atom_num[int(atom)]) if idx2atom_num is not None else int(atom)
+        )
+        fileobj.write(
+            "{0:5}{1:12.6f}{2:12.6f}{3:12.6f}{4:12.6f}\n".format(
+                atomic_num,
+                float(atomic_num),
+                x_coord,
+                y_coord,
+                z_coord,
+            )
+        )
+    density.tofile(fileobj, sep="\n", format="%e")
+
+
+def unavailable_cube_writer(*args, **kwargs):
+    raise AttributeError("Cube writer not available for this dataset")
+
+
+def open_text_maybe_compressed(path):
+    path = Path(path)
+    suffixes = "".join(path.suffixes).lower()
+    if suffixes.endswith(".lz4"):
+        import lz4.frame
+
+        return lz4.frame.open(path, mode="rt")
+    if suffixes.endswith(".xz"):
+        return lzma.open(path, mode="rt")
+    if suffixes.endswith(".gz"):
+        return gzip.open(path, mode="rt")
+    return path.open(mode="rt")
+
+
+def read_cube_density(path):
+    """Read density values, grid coordinates, and metadata from a CUBE file."""
+
+    with open_text_maybe_compressed(path) as file_obj:
+        file_obj.readline()
+        file_obj.readline()
+        line = file_obj.readline().split()
+        if len(line) < 4:
+            raise ValueError(f"Invalid CUBE header (line 3) in {path}")
+        num_atoms = int(line[0])
+        origin = np.array([float(value) for value in line[1:4]], dtype=np.float32)
+
+        shape = []
+        cell = np.zeros((3, 3), dtype=np.float32)
+        for index in range(3):
+            row = file_obj.readline().split()
+            if len(row) < 4:
+                raise ValueError(f"Invalid CUBE axis line in {path}")
+            size, x_coord, y_coord, z_coord = [float(value) for value in row[:4]]
+            shape.append(int(size))
+            cell[index] = np.array([x_coord, y_coord, z_coord], dtype=np.float32)
+
+        x_coords = np.arange(shape[0], dtype=np.float32)[:, None] * cell[0][None, :]
+        y_coords = np.arange(shape[1], dtype=np.float32)[:, None] * cell[1][None, :]
+        z_coords = np.arange(shape[2], dtype=np.float32)[:, None] * cell[2][None, :]
+        grid_coord = (
+            x_coords.reshape(-1, 1, 1, 3)
+            + y_coords.reshape(1, -1, 1, 3)
+            + z_coords.reshape(1, 1, -1, 3)
+        ).reshape(-1, 3)
+        grid_coord += origin
+
+        atom_coord_ref = []
+        for _ in range(num_atoms):
+            row = file_obj.readline().split()
+            if len(row) < 5:
+                raise ValueError(f"Invalid CUBE atom line in {path}")
+            atom_coord_ref.append([float(row[2]), float(row[3]), float(row[4])])
+
+        num_grid_points = shape[0] * shape[1] * shape[2]
+        values = [value for line in file_obj for value in line.split()]
+        if len(values) < num_grid_points:
+            raise ValueError(
+                f"CUBE data too short in {path}: expected {num_grid_points}, "
+                f"got {len(values)}"
+            )
+        density = np.asarray(values[:num_grid_points], dtype=np.float32)
+
+    return (
+        paddle.to_tensor(density, dtype="float32"),
+        paddle.to_tensor(grid_coord, dtype="float32"),
+        {
+            "shape": shape,
+            "cell": paddle.to_tensor(cell, dtype="float32"),
+            "origin": paddle.to_tensor(origin, dtype="float32"),
+            "atom_coord_ref": np.asarray(atom_coord_ref, dtype=np.float32),
+        },
+    )
+
+
+def prepare_cube_info(info, grid_coord):
+    """Build CUBE metadata from dataset information and an explicit grid."""
+
+    shape = info.get("shape")
+    if shape is None or len(shape) != 3:
+        raise ValueError("CUBE output requires a three-dimensional grid shape.")
+
+    shape = [int(size) for size in shape]
+    grid = grid_coord.detach().cpu().numpy().reshape(*shape, 3)
+    origin = grid[0, 0, 0]
+    steps = np.stack(
+        [
+            grid[1, 0, 0] - origin if shape[0] > 1 else np.zeros(3, dtype=np.float32),
+            grid[0, 1, 0] - origin if shape[1] > 1 else np.zeros(3, dtype=np.float32),
+            grid[0, 0, 1] - origin if shape[2] > 1 else np.zeros(3, dtype=np.float32),
+        ]
+    )
+    return {
+        "shape": shape,
+        "cell": steps * np.asarray(shape, dtype=np.float32)[:, None],
+        "origin": origin,
+    }
diff --git a/ppmat/predictor/field_predictor.py b/ppmat/predictor/field_predictor.py
new file mode 100644
index 00000000..2b70f675
--- /dev/null
+++ b/ppmat/predictor/field_predictor.py
@@ -0,0 +1,841 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import copy
+import json
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import paddle
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from ppmat.datasets import DensityDataset
+from ppmat.datasets import SmallDensityDataset
+from ppmat.datasets.geometric_data_type.data import Data
+from ppmat.models import build_model
+from ppmat.models import build_model_from_name
+from ppmat.predictor.base import BasePredictor
+from ppmat.predictor.field_io import prepare_cube_info
+from ppmat.predictor.field_io import read_cube_density
+from ppmat.predictor.field_io import unavailable_cube_writer
+from ppmat.predictor.field_io import write_cube
+from ppmat.utils import logger
+from ppmat.utils import save_load
+from ppmat.utils.misc import set_random_seed
+from ppmat.utils.visualization import draw_volume
+from ppmat.utils.visualization import maybe_downsample_volume
+from ppmat.utils.visualization import safe_write_image
+
+BOHR2ANG = 0.529177
+ANG2BOHR = 1.0 / BOHR2ANG
+
+
+def apply_predict_config(args, cfg):
+    predict_cfg = cfg.get("Predict", {}) or {}
+    defaults = {
+        "split": "test",
+        "index": 0,
+        "data_root": None,
+        "split_file": None,
+        "atom_file": None,
+        "output_dir": "./results",
+        "grid_batch_size": 4096,
+        "skip_vis": False,
+        "save_true_cube": False,
+        "save_pred_cube": False,
+        "save_html": False,
+        "cube_dir": None,
+        "show_plot": False,
+        "mol_pattern": "*.mol",
+        "mol_grid_shape": "80,80,80",
+        "mol_grid_padding": 6.0,
+        "mol_true_cube_dir": None,
+    }
+
+    for name, default in defaults.items():
+        if getattr(args, name) is None:
+            setattr(args, name, predict_cfg.get(name, default))
+
+    return args
+
+
+def inference_model(model, g, density, grid_coord, infos, grid_batch_size=8196):
+    with paddle.no_grad():
+        model.eval()
+        device = paddle.get_device()
+        prepared_infos = (
+            model._prepare_infos(infos, device)
+            if hasattr(model, "_prepare_infos")
+            else infos
+        )
+        if grid_batch_size is None:
+            if hasattr(model, "_forward_density"):
+                preds = model._forward_density(
+                    g.x, g.pos, grid_coord, g.batch, prepared_infos
+                ).squeeze(0)
+            else:
+                # Fallback for legacy models expecting raw tensors
+                preds = model(g.x, g.pos, grid_coord, g.batch, prepared_infos).squeeze(
+                    0
+                )
+        else:
+            preds = []
+            total = grid_coord.shape[1]
+            step = grid_batch_size
+            num_iter = (total + step - 1) // step
+            for start in tqdm(range(0, total, step), total=num_iter):
+                end = min(start + step, total)
+                grid = grid_coord[:, start:end]
+                if hasattr(model, "_forward_density"):
+                    preds.append(
+                        model._forward_density(
+                            g.x, g.pos, grid, g.batch, prepared_infos
+                        ).squeeze(0)
+                    )
+                else:
+                    preds.append(
+                        model(g.x, g.pos, grid, g.batch, prepared_infos).squeeze(0)
+                    )
+            preds = paddle.concat(preds, axis=0)
+
+        if density is None:
+            return preds, None, None
+
+        mask = (density > 0).astype(dtype="float32")
+        preds = preds * mask
+        density = density * mask
+        diff = paddle.abs(preds - density)
+        loss = diff.pow(2).sum()
+        denom = paddle.clip(density.sum(), min=1e-12)
+        mae = diff.sum() / denom
+    return preds, loss, mae
+
+
+def parse_grid_shape(shape_str):
+    parts = [p.strip() for p in str(shape_str).split(",") if p.strip()]
+    if len(parts) == 1:
+        n = int(parts[0])
+        if n <= 1:
+            raise ValueError(
+                f"Invalid mol_grid_shape {shape_str}, each dimension must be > 1"
+            )
+        return [n, n, n]
+    if len(parts) == 3:
+        shape = [int(p) for p in parts]
+        if any(s <= 1 for s in shape):
+            raise ValueError(
+                f"Invalid mol_grid_shape {shape_str}, each dimension must be > 1"
+            )
+        return shape
+    raise ValueError(f"Invalid mol_grid_shape {shape_str}, expected 'N' or 'Nx,Ny,Nz'")
+
+
+def normalize_element_symbol(symbol):
+    sym = str(symbol).strip()
+    if len(sym) == 0:
+        return sym
+    if len(sym) == 1:
+        return sym.upper()
+    return sym[0].upper() + sym[1:].lower()
+
+
+def load_atom_mapping(atom_file):
+    with Path(atom_file).open() as f:
+        atom_info = json.load(f)
+
+    atom_name2idx = {}
+    idx2atom_num = {}
+    for idx, item in enumerate(atom_info):
+        sym = normalize_element_symbol(item["name"])
+        atom_name2idx[sym] = idx
+        idx2atom_num[idx] = int(item["atom_num"])
+    return atom_name2idx, idx2atom_num
+
+
+def resolve_atom_file_for_mol(args_atom_file, dataset_atom_file):
+    candidates = []
+    if args_atom_file is not None:
+        candidates.append(Path(args_atom_file).expanduser())
+    if dataset_atom_file is not None:
+        candidates.append(Path(dataset_atom_file).expanduser())
+
+    for cand in candidates:
+        if cand.exists():
+            return cand
+
+    raise FileNotFoundError(
+        "Could not resolve atom_file for MOL inference. Set --atom_file or "
+        "Predict.atom_file to an existing file. "
+        f"Checked: {[str(candidate) for candidate in candidates]}"
+    )
+
+
+def collect_mol_files(mol_input, mol_pattern):
+    mol_path = Path(mol_input).expanduser()
+    if mol_path.is_file():
+        return [mol_path]
+    if not mol_path.is_dir():
+        raise FileNotFoundError(f"mol_input path not found: {mol_path}")
+
+    files = sorted([p for p in mol_path.glob(mol_pattern) if p.is_file()])
+    if not files:
+        files = sorted(
+            [
+                p
+                for p in mol_path.iterdir()
+                if p.is_file() and p.suffix.lower() == ".mol"
+            ]
+        )
+    if not files:
+        raise FileNotFoundError(f"No .mol files found in directory: {mol_path}")
+    return files
+
+
+def align_mol_atoms_to_cube(g, atom_coord_ref, sample_name, tol=0.05):
+    if atom_coord_ref is None:
+        return g
+    ref = np.asarray(atom_coord_ref, dtype=np.float32)
+    mol = g.pos.numpy().astype(np.float32)
+    if ref.ndim != 2 or ref.shape[1] != 3:
+        logger.warning(
+            f"Invalid reference atom coordinates for {sample_name}, skip alignment"
+        )
+        return g
+    if mol.shape != ref.shape:
+        logger.warning(
+            f"Atom count mismatch for {sample_name} "
+            f"(mol={mol.shape[0]}, cube={ref.shape[0]}), "
+            "skip alignment"
+        )
+        return g
+
+    mol_center = mol.mean(axis=0)
+    ref_center = ref.mean(axis=0)
+    mol_c = mol - mol_center
+    ref_c = ref - ref_center
+    denom = float(np.sqrt((mol_c * mol_c).sum()))
+    numer = float(np.sqrt((ref_c * ref_c).sum()))
+    if denom < 1e-12 or numer < 1e-12:
+        return g
+
+    scale = numer / denom
+    aligned = mol_c * scale + ref_center
+    rms = float(np.sqrt(np.mean((aligned - ref) ** 2)))
+
+    # Typical unit mismatch is Angstrom->Bohr (about 1.8897).
+    # Apply alignment when scale differs from 1.0 or residual is tiny after scaling.
+    if abs(scale - 1.0) > tol or rms < 1e-3:
+        g.pos = paddle.to_tensor(aligned, dtype="float32")
+        logger.info(
+            f"Aligned MOL coordinates to CUBE frame for {sample_name}: "
+            f"scale={scale:.6f} (A->Bohr~{ANG2BOHR:.6f}), rms={rms:.6e}"
+        )
+    else:
+        logger.info(
+            f"No coordinate rescale needed for {sample_name}: "
+            f"scale={scale:.6f}, rms={rms:.6e}"
+        )
+    return g
+
+
+def resolve_true_cube_for_mol(mol_path, true_cube_dir=None):
+    base = sanitize_base_name(mol_path.name)
+    base_density = f"{base[:-3]}Density" if base.endswith("Opt") else f"{base}Density"
+    roots = []
+    if true_cube_dir is not None:
+        roots.append(Path(true_cube_dir).expanduser())
+    roots.append(mol_path.parent)
+
+    stems = [base, f"{base}_true", base_density]
+    exts = [
+        ".cube",
+        ".cub",
+        ".cube.lz4",
+        ".cube.gz",
+        ".cube.xz",
+        ".cub.lz4",
+        ".cub.gz",
+        ".cub.xz",
+    ]
+    name_candidates = []
+    for s in stems:
+        for ext in exts:
+            name_candidates.append(f"{s}{ext}")
+
+    seen = set()
+    uniq_candidates = []
+    for name in name_candidates:
+        if name not in seen:
+            uniq_candidates.append(name)
+            seen.add(name)
+
+    for root in roots:
+        if not root.exists():
+            continue
+        for name in uniq_candidates:
+            p = root / name
+            if p.is_file():
+                return p
+    return None
+
+
+def parse_mol_v2000(mol_path):
+    lines = mol_path.read_text(errors="replace").splitlines()
+    if len(lines) < 4:
+        raise ValueError(f"MOL file too short: {mol_path}")
+
+    counts = lines[3]
+    if "V3000" in counts.upper():
+        raise NotImplementedError(f"V3000 MOL is not supported yet: {mol_path}")
+
+    try:
+        n_atom = int(counts[:3])
+    except Exception:
+        parts = counts.split()
+        if len(parts) < 2:
+            raise ValueError(f"Failed to parse counts line in MOL file: {mol_path}")
+        n_atom = int(parts[0])
+
+    atom_start = 4
+    atom_end = atom_start + n_atom
+    if len(lines) < atom_end:
+        raise ValueError(f"Atom block incomplete in MOL file: {mol_path}")
+
+    coords = []
+    symbols = []
+    for line in lines[atom_start:atom_end]:
+        parts = line.split()
+        x = y = z = None
+        sym = None
+        if len(parts) >= 4:
+            try:
+                x, y, z = float(parts[0]), float(parts[1]), float(parts[2])
+                sym = parts[3]
+            except Exception:
+                x = y = z = None
+                sym = None
+        if x is None:
+            try:
+                x = float(line[0:10])
+                y = float(line[10:20])
+                z = float(line[20:30])
+                sym = line[31:34].strip()
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to parse atom line in {mol_path}: {line}"
+                ) from e
+
+        coords.append([x, y, z])
+        symbols.append(normalize_element_symbol(sym))
+
+    return np.asarray(coords, dtype=np.float32), symbols
+
+
+def build_mol_sample(mol_path, atom_name2idx, mol_grid_shape, mol_grid_padding):
+    atom_coord_np, atom_symbols = parse_mol_v2000(mol_path)
+
+    atom_type_idx = []
+    missing = set()
+    for sym in atom_symbols:
+        idx = atom_name2idx.get(sym)
+        if idx is None:
+            missing.add(sym)
+        else:
+            atom_type_idx.append(idx)
+    if missing:
+        raise ValueError(
+            "Found atoms not covered by atom_file mapping in "
+            f"{mol_path}: {sorted(missing)}"
+        )
+
+    atom_type = paddle.to_tensor(atom_type_idx, dtype="int64")
+    atom_coord = paddle.to_tensor(atom_coord_np, dtype="float32")
+    g = Data(x=atom_type, pos=atom_coord)
+
+    shape = [int(s) for s in mol_grid_shape]
+    min_coord = atom_coord_np.min(axis=0)
+    max_coord = atom_coord_np.max(axis=0)
+    span = np.maximum(
+        max_coord - min_coord, np.array([1e-3, 1e-3, 1e-3], dtype=np.float32)
+    )
+    axis_len = span + 2.0 * float(mol_grid_padding)
+    center = 0.5 * (min_coord + max_coord)
+    origin = center - 0.5 * axis_len
+
+    x = np.linspace(
+        origin[0],
+        origin[0] + axis_len[0],
+        num=shape[0],
+        endpoint=False,
+        dtype=np.float32,
+    )
+    y = np.linspace(
+        origin[1],
+        origin[1] + axis_len[1],
+        num=shape[1],
+        endpoint=False,
+        dtype=np.float32,
+    )
+    z = np.linspace(
+        origin[2],
+        origin[2] + axis_len[2],
+        num=shape[2],
+        endpoint=False,
+        dtype=np.float32,
+    )
+    grid = (
+        np.stack(np.meshgrid(x, y, z, indexing="ij"), axis=-1)
+        .reshape(-1, 3)
+        .astype(np.float32)
+    )
+    grid_coord = paddle.to_tensor(grid, dtype="float32")
+
+    cell = np.diag(axis_len.astype(np.float32))
+    info = {
+        "shape": shape,
+        "cell": paddle.to_tensor(cell, dtype="float32"),
+        "origin": paddle.to_tensor(origin.astype(np.float32), dtype="float32"),
+        "file_name": mol_path.name,
+    }
+
+    return g, None, grid_coord, info
+
+
+def sanitize_base_name(sample_name):
+    base_name = Path(sample_name).name
+    for suf in [".lz4", ".zst", ".gz"]:
+        if base_name.endswith(suf):
+            base_name = base_name[: -len(suf)]
+    for suf in [".cube", ".CHGCAR", ".json", ".mol"]:
+        if base_name.endswith(suf):
+            base_name = base_name[: -len(suf)]
+    return base_name
+
+
+class FieldPredictor(BasePredictor):
+    """Electron-density field predictor."""
+
+    def __init__(
+        self,
+        model_name=None,
+        weights_name=None,
+        config_path=None,
+        checkpoint_path=None,
+        config_overrides=None,
+        seed=42,
+    ):
+        super().__init__(
+            model_name=model_name,
+            weights_name=weights_name,
+            config_path=config_path,
+            checkpoint_path=checkpoint_path,
+            work_dir="",
+            device=None,
+        )
+        self.config_overrides = config_overrides
+        set_random_seed(seed)
+        self._load_model()
+        self.model.eval()
+        logger.info("Model loaded successfully.")
+
+    def _load_model(self):
+        if self.model_name is not None:
+            logger.info(f"Loading registered model: {self.model_name}")
+            self.model, self.config = build_model_from_name(
+                self.model_name, self.weights_name
+            )
+            if self.config_overrides:
+                cfg = OmegaConf.merge(
+                    OmegaConf.create(self.config),
+                    OmegaConf.from_dotlist(self.config_overrides),
+                )
+                self.config = OmegaConf.to_container(cfg, resolve=True)
+        else:
+            assert self.config_path is not None and self.checkpoint_path is not None, (
+                "config_path and checkpoint_path must be provided when model_name "
+                "is None."
+            )
+            logger.info(f"Loading the pretrained model from {self.checkpoint_path}")
+            cfg = OmegaConf.load(self.config_path)
+            if self.config_overrides:
+                cfg = OmegaConf.merge(
+                    cfg, OmegaConf.from_dotlist(self.config_overrides)
+                )
+            self.config = OmegaConf.to_container(cfg, resolve=True)
+            model_config = self.config.get("Model")
+            if model_config is None:
+                raise ValueError(f"Model config is missing from {self.config_path}.")
+            self.model = build_model(model_config)
+            save_load.load_pretrain(self.model, self.checkpoint_path)
+
+    def predict(self, args):
+        return run_prediction(args, self.model, self.config)
+
+    @staticmethod
+    def _save_cubes(
+        args,
+        cube_dir,
+        cube_writer,
+        sample_name,
+        sample_tag,
+        g,
+        density,
+        preds,
+        info,
+        grid_coord,
+    ):
+        if not (args.save_true_cube or args.save_pred_cube):
+            return
+
+        atom_type_np = g.x.detach().cpu().numpy()
+        atom_coord_np = g.pos.detach().cpu().numpy()
+        info_cube = prepare_cube_info(info, grid_coord)
+
+        if args.save_true_cube:
+            if density is None:
+                logger.warning(
+                    f"Skipping true cube for {sample_name}: "
+                    "no reference density available"
+                )
+            else:
+                true_cube_path = cube_dir / f"{sample_tag}_true.cube"
+                with true_cube_path.open("w") as f:
+                    cube_writer(
+                        f,
+                        atom_type_np,
+                        atom_coord_np,
+                        density.detach().cpu().numpy(),
+                        info_cube,
+                    )
+                logger.info(f"Saved reference density cube to: {true_cube_path}")
+
+        if args.save_pred_cube:
+            pred_cube_path = cube_dir / f"{sample_tag}_pred.cube"
+            with pred_cube_path.open("w") as f:
+                cube_writer(
+                    f,
+                    atom_type_np,
+                    atom_coord_np,
+                    preds.detach().cpu().numpy(),
+                    info_cube,
+                )
+            logger.info(f"Saved predicted density cube to: {pred_cube_path}")
+
+    @staticmethod
+    def _save_visualizations(
+        args,
+        output_dir,
+        sample_tag,
+        g,
+        density,
+        preds,
+        info,
+        grid_coord,
+    ):
+        if args.skip_vis:
+            return
+
+        grid_np = grid_coord.detach().cpu().numpy()
+        preds_np = preds.detach().cpu().numpy()
+        shape = info.get("shape")
+        atom_type = g.x.detach().cpu().numpy()
+        atom_coord = g.pos.detach().cpu().numpy()
+        shape = shape if shape is None else [int(s) for s in shape]
+
+        if density is not None:
+            density_np = density.detach().cpu().numpy()
+            diff_np = density_np - preds_np
+            (
+                grid_vis,
+                (density_vis, diff_vis, preds_vis),
+                did_downsample,
+                stride,
+            ) = maybe_downsample_volume(
+                grid_np,
+                [density_np, diff_np, preds_np],
+                shape,
+            )
+            FieldPredictor._log_downsample(grid_np, grid_vis, did_downsample, stride)
+
+            FieldPredictor._write_volume_plot(
+                args,
+                output_dir,
+                sample_tag,
+                "true_density",
+                "DFT electron density",
+                grid_vis,
+                density_vis,
+                atom_type,
+                atom_coord,
+                isomin=0.05,
+                isomax=3.5,
+                surface_count=5,
+            )
+            FieldPredictor._write_volume_plot(
+                args,
+                output_dir,
+                sample_tag,
+                "diff_density",
+                "Electron Density Difference",
+                grid_vis,
+                diff_vis,
+                atom_type,
+                atom_coord,
+                isomin=-0.06,
+                isomax=0.06,
+                surface_count=4,
+            )
+            FieldPredictor._write_volume_plot(
+                args,
+                output_dir,
+                sample_tag,
+                "pred_density",
+                "Predicted Electron Density",
+                grid_vis,
+                preds_vis,
+                atom_type,
+                atom_coord,
+                isomin=0.05,
+                isomax=3.5,
+                surface_count=5,
+            )
+            return
+
+        (
+            grid_vis,
+            (preds_vis,),
+            did_downsample,
+            stride,
+        ) = maybe_downsample_volume(grid_np, [preds_np], shape)
+        FieldPredictor._log_downsample(grid_np, grid_vis, did_downsample, stride)
+        FieldPredictor._write_volume_plot(
+            args,
+            output_dir,
+            sample_tag,
+            "pred_density",
+            "Predicted Electron Density",
+            grid_vis,
+            preds_vis,
+            atom_type,
+            atom_coord,
+            isomin=0.05,
+            isomax=3.5,
+            surface_count=5,
+        )
+
+    @staticmethod
+    def _log_downsample(grid_np, grid_vis, did_downsample, stride):
+        if did_downsample:
+            logger.warning(
+                f"Downsampled volume grid from {grid_np.shape[0]} "
+                f"to {grid_vis.shape[0]} points for visualization "
+                f"(stride={stride}) to keep HTML output responsive."
+            )
+
+    @staticmethod
+    def _write_volume_plot(
+        args,
+        output_dir,
+        sample_tag,
+        suffix,
+        title,
+        grid,
+        values,
+        atom_type,
+        atom_coord,
+        isomin,
+        isomax,
+        surface_count,
+    ):
+        logger.info(f"Visualizing {title}")
+        fig = draw_volume(
+            grid,
+            values,
+            atom_type,
+            atom_coord,
+            isomin=isomin,
+            isomax=isomax,
+            surface_count=surface_count,
+            title=title,
+        )
+        image_path = output_dir / f"{sample_tag}_{suffix}.png"
+        safe_write_image(fig, image_path, show_plot=args.show_plot)
+        if args.save_html:
+            fig.write_html(output_dir / f"{sample_tag}_{suffix}.html")
+
+
+def run_prediction(args, model, cfg):
+    apply_predict_config(args, cfg)
+
+    split_key = "val" if args.split == "validation" else args.split
+    ds_cfg_full = cfg["Dataset"][split_key]["dataset"]
+    dataset_cfg = ds_cfg_full.get("__init_params__", {})
+    dataset_params = copy.deepcopy(dataset_cfg)
+    dataset_params["split"] = args.split
+    if args.data_root is not None:
+        dataset_params["root"] = args.data_root
+    if args.split_file is not None:
+        dataset_params["split_file"] = args.split_file
+    if args.atom_file is not None:
+        dataset_params["atom_file"] = args.atom_file
+
+    use_mol_mode = args.mol_input is not None
+
+    dataset = None
+    cube_writer = None
+    idx2atom_num = None
+    atom_name2idx = None
+    mol_files = []
+    mol_grid_shape = None
+
+    if use_mol_mode:
+        atom_file_path = resolve_atom_file_for_mol(
+            args.atom_file,
+            dataset_params.get("atom_file"),
+        )
+        atom_name2idx, idx2atom_num = load_atom_mapping(atom_file_path)
+        mol_files = collect_mol_files(args.mol_input, args.mol_pattern)
+        mol_grid_shape = parse_grid_shape(args.mol_grid_shape)
+        cube_writer = partial(write_cube, idx2atom_num=idx2atom_num)
+        logger.info(
+            f"MOL mode enabled: {len(mol_files)} file(s), atom_file={atom_file_path}, "
+            f"grid_shape={mol_grid_shape}, padding={args.mol_grid_padding}, "
+            f"true_cube_dir={args.mol_true_cube_dir}"
+        )
+    else:
+        dataset_cls_name = ds_cfg_full.get("__class_name__", "DensityDataset")
+        dataset_cls_map = {
+            "DensityDataset": DensityDataset,
+            "SmallDensityDataset": SmallDensityDataset,
+        }
+        if dataset_cls_name not in dataset_cls_map:
+            raise ValueError(f"Unsupported dataset class {dataset_cls_name}")
+        dataset = dataset_cls_map[dataset_cls_name](**dataset_params)
+        cube_writer = getattr(dataset, "write_cube", None)
+        idx2atom_num = getattr(dataset, "idx2atom_num", None)
+        if cube_writer is None:
+            if isinstance(dataset, SmallDensityDataset):
+                # Atom order in SmallDensityDataset: C=0, H=1, O=2
+                idx2atom_num = np.array([6, 1, 8], dtype=np.int64)
+                cube_writer = partial(
+                    write_cube,
+                    idx2atom_num=idx2atom_num,
+                )
+            else:
+                cube_writer = unavailable_cube_writer
+        if args.index >= len(dataset):
+            raise IndexError(
+                f"Index {args.index} exceeds dataset size {len(dataset)} "
+                f"for split {args.split}"
+            )
+
+    device = "gpu" if paddle.is_compiled_with_cuda() else "cpu"
+    paddle.set_device(device)
+    logger.info(f"Running inference on device: {device}")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    cube_dir = Path(args.cube_dir) if args.cube_dir is not None else output_dir
+    cube_dir.mkdir(parents=True, exist_ok=True)
+
+    if use_mol_mode:
+        sample_iter = tqdm(mol_files, desc="MOL inference")
+    else:
+        sample_iter = [args.index]
+
+    for sample_item in sample_iter:
+        if use_mol_mode:
+            mol_path = sample_item
+            g, density, grid_coord, info = build_mol_sample(
+                mol_path,
+                atom_name2idx,
+                mol_grid_shape,
+                args.mol_grid_padding,
+            )
+            true_cube_path = resolve_true_cube_for_mol(mol_path, args.mol_true_cube_dir)
+            if true_cube_path is not None:
+                try:
+                    density, grid_coord, info_ref = read_cube_density(true_cube_path)
+                    g = align_mol_atoms_to_cube(
+                        g, info_ref.get("atom_coord_ref"), mol_path.name
+                    )
+                    info = dict(info_ref)
+                    info["file_name"] = mol_path.name
+                    info["true_cube_file"] = str(true_cube_path)
+                    logger.info(
+                        f"Using reference cube for {mol_path.name}: {true_cube_path}"
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to read reference cube for {mol_path.name} "
+                        f"at {true_cube_path}: {e}"
+                    )
+            sample_name = info.get("file_name", mol_path.name)
+        else:
+            sample_name = f"{args.split}_{args.index}"
+            g, density, grid_coord, info = dataset[args.index]
+            sample_name = info.get("file_name", sample_name)
+
+        g.batch = paddle.zeros_like(g.x)
+        g = g.to(device)
+        if density is not None:
+            density = density.to(device)
+        grid_coord = grid_coord.to(device)
+
+        logger.info(f"Starting prediction for sample: {sample_name}")
+        preds, loss, mae = inference_model(
+            model,
+            g,
+            density,
+            grid_coord[None],
+            [info],
+            grid_batch_size=args.grid_batch_size,
+        )
+        if loss is not None and mae is not None:
+            logger.info(
+                f"Prediction completed for {sample_name}, "
+                f"Loss: {float(loss):.6f}, MAE: {float(mae):.6f}"
+            )
+        else:
+            logger.info(
+                f"Prediction completed for {sample_name} (no reference density)"
+            )
+
+        sample_tag = sanitize_base_name(sample_name)
+
+        FieldPredictor._save_cubes(
+            args,
+            cube_dir,
+            cube_writer,
+            sample_name,
+            sample_tag,
+            g,
+            density,
+            preds,
+            info,
+            grid_coord,
+        )
+        FieldPredictor._save_visualizations(
+            args,
+            output_dir,
+            sample_tag,
+            g,
+            density,
+            preds,
+            info,
+            grid_coord,
+        )
diff --git a/ppmat/sampler/__init__.py b/ppmat/sampler/__init__.py
index 675cf5cb..d8616a83 100644
--- a/ppmat/sampler/__init__.py
+++ b/ppmat/sampler/__init__.py
@@ -11,3 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from ppmat.sampler.molecular_sampler import MolecularSampler
+from ppmat.sampler.structure_sampler import StructureSampler
+
+__all__ = ["MolecularSampler", "StructureSampler"]
diff --git a/ppmat/sampler/base_sampler.py b/ppmat/sampler/diffnmr.py
similarity index 71%
rename from ppmat/sampler/base_sampler.py
rename to ppmat/sampler/diffnmr.py
index 90fe2e8a..e6e17dcd 100644
--- a/ppmat/sampler/base_sampler.py
+++ b/ppmat/sampler/diffnmr.py
@@ -14,6 +14,7 @@
 
 import copy
 import os
+import os.path as osp
 import time
 from typing import Dict
 from typing import List
@@ -31,24 +32,25 @@
 from ppmat.datasets import build_dataloader
 from ppmat.datasets import build_dataset_infos
 from ppmat.datasets import set_signal_handlers
-from ppmat.datasets.msd_nmr_dataset import DataLoaderCollection
 from ppmat.datasets.transform import build_post_transforms
 from ppmat.metrics import DiffNMRStreamingAdapter
-from ppmat.metrics import build_metric
+from ppmat.models import MODEL_REGISTRY
 from ppmat.models import build_model
-from ppmat.models import build_model_from_name
 from ppmat.models.diffnmr.extra_features_graph import DummyExtraFeatures
 from ppmat.models.diffnmr.extra_features_graph import ExtraFeatures
 from ppmat.models.diffnmr.extra_features_molecular_graph import ExtraMolecularFeatures
 from ppmat.models.diffnmr.utils import diffgraphformer_utils
 from ppmat.schedulers import scheduling_diffnmr
+from ppmat.utils import download
 from ppmat.utils import logger
 from ppmat.utils import save_load
+from ppmat.utils.model_package import get_model_config_path
+from ppmat.utils.model_package import resolve_model_package_dir
 from ppmat.utils.visualization import MolecularVisualization
 
 
-class MolecularSampler:
-    """Molecular Sampler.
+class DiffNMRSampler:
+    """DiffNMR Sampler.
 
     This class provides an interface for sampling structures using pre-trained deep
     learning models. Supports two initialization modes:
@@ -87,113 +89,143 @@ def __init__(
         weights_name: Optional[str] = None,
         config_path: Optional[str] = None,
         checkpoint_path: Optional[str] = None,
+        config_overrides: Optional[List[str]] = None,
     ):
-        # if model_name is not None, then config_path and checkpoint_path must be
-        # provided
+        package_config_dir = None
         if model_name is None:
-            assert (
-                config_path is not None and checkpoint_path is not None
-            ), "config_path and checkpoint_path must be provided when model_name is "
-            "None."
+            assert config_path is not None and checkpoint_path is not None, (
+                "config_path and checkpoint_path must be provided when model_name is "
+                "None."
+            )
 
             logger.info(f"Loading model from {config_path} and {checkpoint_path}.")
 
+            config_base_dir = os.path.dirname(os.path.abspath(config_path))
+            checkpoint_dir = (
+                checkpoint_path
+                if checkpoint_path and os.path.isdir(checkpoint_path)
+                else None
+            )
             config = OmegaConf.load(config_path)
+            if config_overrides:
+                cli_config = OmegaConf.from_dotlist(config_overrides)
+                config = OmegaConf.merge(config, cli_config)
             config = OmegaConf.to_container(config, resolve=True)
+            self._resolve_package_paths(
+                config,
+                config_base_dir=config_base_dir,
+                checkpoint_dir=checkpoint_dir,
+            )
+        else:
+            logger.info(f"Loading registered model: {model_name}")
+            extracted_path = download.get_weights_path_from_url(
+                MODEL_REGISTRY[model_name]
+            )
+            checkpoint_path = resolve_model_package_dir(model_name, extracted_path)
+            config_path = get_model_config_path(model_name, checkpoint_path)
+            package_config_dir = checkpoint_path
+            config = OmegaConf.load(config_path)
+            if config_overrides:
+                cli_config = OmegaConf.from_dotlist(config_overrides)
+                config = OmegaConf.merge(config, cli_config)
+            config = OmegaConf.to_container(config, resolve=True)
+            self._resolve_package_paths(
+                config,
+                config_base_dir=package_config_dir,
+                checkpoint_dir=None,
+            )
 
-            model_config = config.get("Model", None)
-            assert model_config is not None, "Model config must be provided."
-
-            # TODO: optimize in the future
-            set_signal_handlers()
-            train_data_cfg = config["Dataset"].get("train")
-            train_loader = build_dataloader(train_data_cfg)
-
-            val_data_cfg = config["Dataset"].get("val")
-            val_loader = build_dataloader(val_data_cfg)
+        model_config = config.get("Model", None)
+        assert model_config is not None, "Model config must be provided."
 
-            test_data_cfg = config["Dataset"].get("test")
-            test_loader = build_dataloader(test_data_cfg)
+        set_signal_handlers()
+        sample_loader = build_dataloader(config["Sampler"]["data"])
 
-            # build datasetinfo
-            dataloaders = DataLoaderCollection(train_loader, val_loader, test_loader)
-            dataset_infos = build_dataset_infos(
-                dataloaders=dataloaders, cfg=config, recompute_statistics=False
-            )
-            train_smiles = dataset_infos.train_smiles
-
-            # extra features
-            if (
-                config["Model"]["__init_params__"]["diffmodel_cfg"]["extra_features"]
-                is not None
-            ):
-                extra_features = ExtraFeatures(
-                    config["Model"]["__init_params__"]["diffmodel_cfg"][
-                        "extra_features"
-                    ],
-                    dataset_infos=dataset_infos,
-                )
-                domain_features = ExtraMolecularFeatures(
-                    dataset_infos=dataset_infos,
-                )
-            else:
-                extra_features = DummyExtraFeatures()
-                domain_features = DummyExtraFeatures()
-            fallback_loader = train_loader or val_loader or test_loader
-            dataset_infos.compute_input_output_dims(
-                dataloader=fallback_loader,
-                extra_features=extra_features,
-                domain_features=domain_features,
-                conditionDim=config["Model"]["__init_params__"]["diffmodel_cfg"][
-                    "conditdim"
-                ],
-            )
+        # Build dataset infos without constructing full train/val/test dataloaders.
+        dataset_info_config = copy.deepcopy(config)
+        dataset_info_config["Dataset"]["train"]["dataset"]["__init_params__"][
+            "load_train_smiles"
+        ] = False
+        dataset_infos = build_dataset_infos(
+            dataloaders=None, cfg=dataset_info_config, recompute_statistics=False
+        )
+        train_smiles = dataset_infos.train_smiles
 
-            # CLIP for sample metric
-            model_cfg = config["CLIP"]
-            self.clip = build_model(
-                model_cfg,
-                extra_features=extra_features,
-                domain_features=domain_features,
+        # extra features
+        if (
+            config["Model"]["__init_params__"]["diffmodel_cfg"]["extra_features"]
+            is not None
+        ):
+            extra_features = ExtraFeatures(
+                config["Model"]["__init_params__"]["diffmodel_cfg"]["extra_features"],
                 dataset_infos=dataset_infos,
             )
-
-            # visualization tools
-            self.visualization_tools = MolecularVisualization(
+            domain_features = ExtraMolecularFeatures(
                 dataset_infos=dataset_infos,
-                output_dir=config["Trainer"]["output_dir"],
             )
+        else:
+            extra_features = DummyExtraFeatures()
+            domain_features = DummyExtraFeatures()
+        dataset_infos.compute_input_output_dims(
+            dataloader=sample_loader,
+            extra_features=extra_features,
+            domain_features=domain_features,
+            conditionDim=config["Model"]["__init_params__"]["diffmodel_cfg"][
+                "conditdim"
+            ],
+        )
 
-            model_cfg = config["Model"]
-            model = build_model(
-                model_cfg,
-                extra_features=extra_features,
-                domain_features=domain_features,
-                dataset_infos=dataset_infos,
-                visualization_tools=self.visualization_tools,
-                clip=self.clip,
-            )
+        # CLIP for sample metric
+        model_cfg = config["CLIP"]
+        self.clip = build_model(
+            model_cfg,
+            extra_features=extra_features,
+            domain_features=domain_features,
+            dataset_infos=dataset_infos,
+        )
 
-            self.pretrained_model_path = (
-                checkpoint_path
-                if checkpoint_path is not None
-                else config.get("pretrained_model_path", None)
-            )
-            self.pretrained_weight_name = (
-                weights_name
-                if weights_name is not None
-                else config.get("pretrained_weight_name", None)
+        # visualization tools
+        self.visualization_tools = MolecularVisualization(
+            dataset_infos=dataset_infos,
+            output_dir=config["Trainer"]["output_dir"],
+        )
+
+        model_cfg = config["Model"]
+        model = build_model(
+            model_cfg,
+            extra_features=extra_features,
+            domain_features=domain_features,
+            dataset_infos=dataset_infos,
+            visualization_tools=self.visualization_tools,
+            clip=self.clip,
+        )
+
+        self.pretrained_model_path = (
+            checkpoint_path
+            if checkpoint_path is not None
+            else config.get("pretrained_model_path", None)
+        )
+        self.pretrained_weight_name = weights_name
+        if self.pretrained_weight_name is None:
+            self.pretrained_weight_name = config.get("pretrained_weight_name", None)
+        if (
+            self.pretrained_weight_name is None
+            and self.pretrained_model_path is not None
+            and os.path.isdir(self.pretrained_model_path)
+        ):
+            sampler_pretrained_path = config.get("Sampler", {}).get(
+                "pretrained_model_path", None
             )
+            if sampler_pretrained_path is not None:
+                self.pretrained_weight_name = os.path.basename(sampler_pretrained_path)
+        if self.pretrained_model_path is not None:
             save_load.load_pretrain(
                 model, self.pretrained_model_path, self.pretrained_weight_name
             )
 
-        else:
-            logger.info("Since model_name is given, downloading it...")
-            model, config = build_model_from_name(model_name, weights_name)
-
         self.model = model
         self.config = config
+        self._sample_loader = sample_loader
 
         self.model.eval()
 
@@ -206,10 +238,12 @@ def __init__(
         self.number_chain_steps = sample_config["number_chain_steps"]
         self.sample_batch_iters = sample_config["sample_batch_iters"]
         self.metric_dict_sample = sample_config.get("out_dict", None)
-        self.flag_retrival_sampling = sample_config.get("flag_retrival_sampling", False)
+        self.flag_retrieval_sampling = sample_config.get(
+            "flag_retrieval_sampling", False
+        )
         self.flag_use_formula = sample_config.get("flag_use_formula", False)
-        self.flag_retrival_initilization = sample_config.get(
-            "flag_retrival_initilization", False
+        self.flag_retrieval_initialization = sample_config.get(
+            "flag_retrieval_initialization", False
         )
         self.num_candidates = sample_config.get("num_candidates", 1)
 
@@ -226,14 +260,14 @@ def __init__(
             else 0
         )
         self.output_dir = self.config.get("Sampler", {}).get("output_dir", "./outputs")
-        os.makedirs(self.output_dir, exist_ok=True)
-
-        if self.clip is not None:
-            setattr(self.model, "clip", self.clip)
+        self._set_output_dir(self.output_dir)
 
-        self.molecular_vectors, self.smiles_list = self._init_retrieval_bank(
-            self.sample_config,
-        )
+        if self.flag_retrieval_sampling or self.flag_retrieval_initialization:
+            self.molecular_vectors, self.smiles_list = self._init_retrieval_bank(
+                self.sample_config,
+            )
+        else:
+            self.molecular_vectors, self.smiles_list = None, None
 
         self.streaming = DiffNMRStreamingAdapter(
             t_scale=float(self.sample_config.get("t_scale", 1.0)),
@@ -246,29 +280,115 @@ def __init__(
             train_smiles=train_smiles,
             num_candidate=self.num_candidates,
         )
-        setattr(self.model, "streaming_adapter", self.streaming)
+
+    @staticmethod
+    def _resolve_pretrained_paths(
+        config: Dict,
+        config_base_dir: Optional[str],
+        checkpoint_dir: Optional[str] = None,
+    ):
+        DiffNMRSampler._resolve_package_paths(config, config_base_dir, checkpoint_dir)
+
+    @staticmethod
+    def _resolve_package_paths(
+        config: Dict,
+        config_base_dir: Optional[str],
+        checkpoint_dir: Optional[str] = None,
+    ):
+        def resolve_path(path):
+            if path is None or osp.isabs(path) or path.startswith("http"):
+                return path
+
+            if checkpoint_dir is not None:
+                candidate = osp.join(checkpoint_dir, osp.basename(path))
+                if osp.exists(candidate):
+                    return candidate
+
+            if config_base_dir is not None:
+                if path.startswith("./checkpoints/") or path.startswith("checkpoints/"):
+                    candidate = osp.join(
+                        config_base_dir, "checkpoints", osp.basename(path)
+                    )
+                    if osp.exists(candidate):
+                        return candidate
+                candidate = osp.normpath(osp.join(config_base_dir, path))
+                if osp.exists(candidate):
+                    return candidate
+
+            return path
+
+        def visit(obj):
+            if isinstance(obj, dict):
+                for key, value in obj.items():
+                    if key in {
+                        "pretrained_path",
+                        "pretrained_model_path",
+                        "vocab_peakwidth_path",
+                        "vocab_split_path",
+                        "retrieval_database_path",
+                        "path",
+                        "datadir",
+                    } and isinstance(value, str):
+                        resolved_path = resolve_path(value)
+                        if key in {
+                            "vocab_peakwidth_path",
+                            "vocab_split_path",
+                        } and not osp.isfile(resolved_path):
+                            raise FileNotFoundError(
+                                f"DiffNMR package is missing required {key}: "
+                                f"{resolved_path}"
+                            )
+                        obj[key] = resolved_path
+                    else:
+                        visit(value)
+            elif isinstance(obj, list):
+                for item in obj:
+                    visit(item)
+
+        visit(config)
+
+        sampler_config = config.get("Sampler", {})
+        retrieval_enabled = sampler_config.get(
+            "flag_retrieval_sampling", False
+        ) or sampler_config.get("flag_retrieval_initialization", False)
+        retrieval_path = sampler_config.get("retrieval_database_path")
+        if retrieval_enabled and (
+            retrieval_path is None or not osp.isfile(retrieval_path)
+        ):
+            raise FileNotFoundError(
+                "Retrieval sampling requires an existing "
+                "Sampler.retrieval_database_path."
+            )
+
+    def _set_output_dir(self, output_dir: str):
+        self.output_dir = output_dir
+        os.makedirs(self.output_dir, exist_ok=True)
+        if self.visualization_tools is not None:
+            self.visualization_tools.result_path = osp.join(self.output_dir, "graph")
 
     def compute_metric(
         self,
         save_path=None,
     ):
-        self.output_dir = save_path if save_path is not None else self.output_dir
-        metrics_cfg = self.sample_config.get("metrics")
-        assert metrics_cfg is not None, "metrics config must be provided."
-        metrics_fn = build_metric(metrics_cfg)
-
-        total_results = self.sample_by_dataloader(
+        if save_path is not None:
+            self._set_output_dir(save_path)
+        return self.sample_by_dataloader(
             self.output_dir,
         )
 
-        metric = metrics_fn(total_results)
-        return metric
-
     def post_process(self, data):
         if self.post_transforms is None:
             return data
         return self.post_transforms(data)
 
+    @staticmethod
+    def _clamp_keep_chain(keep_chain: int, n_nodes: Union[int, paddle.Tensor]):
+        if keep_chain <= 0:
+            return 0
+        if isinstance(n_nodes, int):
+            return min(keep_chain, n_nodes)
+        return min(keep_chain, int(n_nodes.shape[0]))
+
     def sample(self, data, sample_params=None):
         if sample_params is None:
             sample_params = {}
@@ -280,10 +400,14 @@ def sample(self, data, sample_params=None):
     def sample_by_dataloader(
         self,
         save_path=None,
+        data_loader=None,
     ):
-        self.output_dir = save_path if save_path is not None else self.output_dir
-        dataset_cfg = self.sample_config["data"]
-        data_loader = build_dataloader(dataset_cfg)
+        if save_path is not None:
+            self._set_output_dir(save_path)
+        if data_loader is None:
+            data_loader = getattr(self, "_sample_loader", None)
+        if data_loader is None:
+            data_loader = build_dataloader(self.sample_config["data"])
 
         # build_molecule_cfg = self.sample_config["build_molecule_cfg"]
         # molecule_converter = BuildMolecule(**build_molecule_cfg)
@@ -302,7 +426,7 @@ def sample_by_dataloader(
         metric_dict = self.sample_epoch(
             data_loader,
             epoch_id,
-            keep_onehot=self.flag_retrival_sampling,
+            keep_onehot=self.flag_retrieval_sampling,
             num_candidates=self.num_candidates,
         )
 
@@ -320,6 +444,7 @@ def sample_by_dataloader(
                         else f" | {k}(metric): {v:.5f}"
                     )
             logger.info(msg)
+        return metric_dict
 
     @paddle.no_grad()
     def sample_epoch(
@@ -344,9 +469,8 @@ def sample_epoch(
             Trainer / Runner object that holds the diffusion ``model``, runtime
             configs, logging utilities, etc.
         dataloader : paddle.io.DataLoader
-            Yields tuples ``(graph, aux_data)`` where `graph` is a *pgl* style
-            MiniBatchGraph and `aux_data` is a dict containing scalar labels,
-            condition vectors and atom counts. TODO: recheck details.
+            Yields batch dictionaries containing a PGL ``graph`` plus
+            ``property`` and ``spectrum`` dictionaries.
         epoch_id : int
             Current epoch index – propagated to the metric logger so that saved
             artefacts (csv / images) are grouped by epoch.
@@ -420,22 +544,25 @@ def sample_epoch(
             dense_data = dense_data.mask(node_mask)  # remove padding rows
 
             # basic batch tensors
-            batch_atomCount = paddle.to_tensor(
+            batch_atomCount = self._as_tensor(
                 batch_property["atom_count"]
             )  # [B] number of atoms
-            batch_y = paddle.to_tensor(batch_property["y"])  # labels (unused here)
+            batch_y = self._as_tensor(batch_property["y"])  # labels (unused here)
             batch_X, batch_E = dense_data.X, dense_data.E  # one‑hot Node / Edge
             bs = len(batch_y)  # batch size
 
             # 2.b build four‑branch NMR condition tensor list
-            if hasattr(self.model, "seq_len_H1"):
-                cond_H = paddle.to_tensor(batch_spectrum["H_nmr"])
-                cond_C = paddle.to_tensor(batch_spectrum["C_nmr"])
-                num_H_peak = paddle.to_tensor(batch_spectrum["num_H_peak"])
-                num_C_peak = paddle.to_tensor(batch_spectrum["num_C_peak"])
+            if getattr(self.model, "conditioning_mode", None) == "spectrum":
+                cond_H = self._as_tensor(batch_spectrum["H_nmr"])
+                cond_C = self._as_tensor(batch_spectrum["C_nmr"])
+                num_H_peak = self._as_tensor(batch_spectrum["num_H_peak"])
+                num_C_peak = self._as_tensor(batch_spectrum["num_C_peak"])
                 batch_nmr = [cond_H, num_H_peak, cond_C, num_C_peak]
             else:
-                batch_nmr = None  # TODO: re‑implement for single‑branch condition
+                raise NotImplementedError(
+                    "DiffNMRSampler currently supports only spectrum-conditioned "
+                    "models with four H1/C13 condition branches."
+                )
 
             # 2.c call `sample_batch` `num_candidates` times
             for c_idx in range(num_candidates):
@@ -457,9 +584,9 @@ def sample_epoch(
                     flag_useformula=self.flag_use_formula,
                     iter_idx=c_idx,
                 )
-                if self.flag_retrival_initilization:
+                if self.flag_retrieval_initialization:
                     kwargs.update(
-                        retrival_initilization=self.flag_retrival_initilization,
+                        retrieval_initialization=self.flag_retrieval_initialization,
                         clip=self.clip,
                         molecular_vectors=self.molecular_vectors,
                         smiles_list=self.smiles_list,
@@ -478,18 +605,9 @@ def sample_epoch(
                 if c_idx == 0:
                     samples["pred"].extend(mol_pred)
                     samples["true"].extend(mol_true)
-                # samples["n_all"] += len(batch_y) # TODO right?
 
             # 2‑d) meta‑info used by retrieval metrics
-            if batch_nmr is not None:
-                samples["batch_condition"] = [None for _ in range(4)]
-                for i, t in enumerate(batch_nmr):
-                    if samples["batch_condition"][i] is None:
-                        samples["batch_condition"][i] = paddle.to_tensor(t)
-                    else:
-                        samples["batch_condition"][i] = paddle.concat(
-                            [samples["batch_condition"][i], paddle.to_tensor(t)], axis=0
-                        )
+            self._append_batch_conditions(samples, batch_nmr)
             samples["node_mask_meta"].extend(batch_atomCount)
             samples["n_all"] += bs
 
@@ -514,6 +632,27 @@ def sample_epoch(
 
         return metric_dict
 
+    @staticmethod
+    def _append_batch_conditions(samples, batch_condition):
+        """Append one batch of NMR conditions without losing earlier batches."""
+
+        tensors = [DiffNMRSampler._as_tensor(value) for value in batch_condition]
+        if not samples["batch_condition"]:
+            samples["batch_condition"] = tensors
+            return
+        if len(samples["batch_condition"]) != len(tensors):
+            raise ValueError(
+                "Inconsistent number of NMR condition branches across batches."
+            )
+        samples["batch_condition"] = [
+            paddle.concat([previous, current], axis=0)
+            for previous, current in zip(samples["batch_condition"], tensors)
+        ]
+
+    @staticmethod
+    def _as_tensor(value):
+        return value if isinstance(value, paddle.Tensor) else paddle.to_tensor(value)
+
     @paddle.no_grad()
     def sample_batch(
         self,
@@ -531,7 +670,7 @@ def sample_batch(
         num_nodes: Union[int, paddle.Tensor] = None,
         flag_useformula: bool = False,
         return_onehot: bool = False,
-        retrival_initilization: bool = False,
+        retrieval_initialization: bool = False,
         clip: paddle.nn.Layer = None,
         molecular_vectors: paddle.Tensor = None,
         smiles_list: List = None,
@@ -561,7 +700,7 @@ def sample_batch(
         batch_y : paddle.Tensor
             Additional labels (if any) required by the model.
         iter_idx : int
-            Current iteration index for obtain candidates for retrival.
+            Current iteration index for obtain candidates for retrieval.
         num_nodes : int | paddle.Tensor | None
             Number of nodes per graph. When *None* the model samples from its own
             learned distribution.
@@ -571,7 +710,7 @@ def sample_batch(
         return_onehot : bool
             Whether to return the *padded* one‑hot tensors (`X_hot`, `E_hot`) in
             addition to discrete index lists – required by molVec retrieval.
-        retrival_initilization : bool, default False
+        retrieval_initialization : bool, default False
             Whether to enable **retrieval‑based initialization**.
             If True, the model will fetch the closest reference molecules
             (using `molecular_vectors`) and use them as the first step of
@@ -602,9 +741,12 @@ def sample_batch(
             n_nodes = model.node_dist.sample_n(batch_size)
         elif isinstance(num_nodes, int):
             n_nodes = paddle.full([batch_size], num_nodes, dtype="int64")
+        elif isinstance(num_nodes, paddle.Tensor):
+            n_nodes = num_nodes
         else:
-            n_nodes = paddle.to_tensor(num_nodes)  # assume Tensor
+            n_nodes = paddle.to_tensor(num_nodes)
 
+        keep_chain = self._clamp_keep_chain(keep_chain, n_nodes)
         n_max: int = int(paddle.max(n_nodes).item())  # ***largest graph size***
 
         # `node_mask[b, i] == True` if node *i* is real for graph *b*
@@ -618,13 +760,18 @@ def sample_batch(
         )
         X_t, E_t, y_t = z_T.X, z_T.E, z_T.y
 
-        chain_X = paddle.zeros([number_chain_steps, keep_chain, n_max], dtype="int64")
-        chain_E = paddle.zeros(
-            [number_chain_steps, keep_chain, n_max, n_max], dtype="int64"
-        )
+        if keep_chain > 0:
+            chain_X = paddle.zeros(
+                [number_chain_steps, keep_chain, n_max], dtype="int64"
+            )
+            chain_E = paddle.zeros(
+                [number_chain_steps, keep_chain, n_max, n_max], dtype="int64"
+            )
+        else:
+            chain_X = chain_E = None
 
         # 3. Retrieval Initialization(Optional)
-        if retrival_initilization and batch_condition is not None:
+        if retrieval_initialization and batch_condition is not None:
             logger.info("Sampling Initializing using Retrieval Method.")
             output = clip.spectrum_encoder(batch_condition)
 
@@ -702,9 +849,10 @@ def sample_batch(
                 X_t = batch_X
 
             # save intermediate frames for the first `keep_chain` graphs
-            write_index = (s_int * number_chain_steps) // model.T
-            chain_X[write_index] = discrete_sampled_s.X[:keep_chain]
-            chain_E[write_index] = discrete_sampled_s.E[:keep_chain]
+            if keep_chain > 0:
+                write_index = (s_int * number_chain_steps) // model.T
+                chain_X[write_index] = discrete_sampled_s.X[:keep_chain]
+                chain_E[write_index] = discrete_sampled_s.E[:keep_chain]
 
         # 5. Collapse padding → obtain discrete indices; optionally keep one‑hot
         # Make a *clone* of `sampled_s` so that collapsing will not overwrite the
@@ -779,15 +927,15 @@ def sample_batch(
                 )
                 assert chain_X.shape[0] == (number_chain_steps + 10)
 
-            # 7.b use visulize tools
-            num_mols = chain_X.shape[1]
-            # draw animation of diffusion process of generated molecules
-            for i in range(num_mols):
-                chain_X_np = chain_X[:, i, :].numpy()
-                chain_E_np = chain_E[:, i, :, :].numpy()
-                self.visualization_tools.visualize_chain(
-                    batch_id, i, chain_X_np, chain_E_np
-                )
+                # 7.b use visulize tools
+                num_mols = chain_X.shape[1]
+                # draw animation of diffusion process of generated molecules
+                for i in range(num_mols):
+                    chain_X_np = chain_X[:, i, :].numpy()
+                    chain_E_np = chain_E[:, i, :, :].numpy()
+                    self.visualization_tools.visualize_chain(
+                        batch_id, i, chain_X_np, chain_E_np
+                    )
             # draw picture of predicted and true molecules
             self.visualization_tools.visualizeNmr(
                 batch_id,
@@ -807,7 +955,7 @@ def _init_retrieval_bank(self, cfg):
         """
         if not cfg:
             return None, None
-        path = cfg.get("retrival_database_path", None)
+        path = cfg.get("retrieval_database_path", None)
         if path is None or not os.path.exists(path):
             logger.warning(f"[retrieval_bank] path missing or not found: {path}")
             return None, None
diff --git a/ppmat/sampler/molecular_sampler.py b/ppmat/sampler/molecular_sampler.py
new file mode 100644
index 00000000..ea0f03d8
--- /dev/null
+++ b/ppmat/sampler/molecular_sampler.py
@@ -0,0 +1,127 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+from typing import List
+from typing import Optional
+
+from omegaconf import OmegaConf
+
+SAMPLER_REGISTRY = {
+    "diffnmr": "ppmat.sampler.diffnmr:DiffNMRSampler",
+}
+
+MODEL_CLASS_TO_SAMPLER = {
+    "DiffNMR": "diffnmr",
+}
+
+MODEL_NAME_TO_SAMPLER = {
+    "diffnmr_msdnmr_nless15": "diffnmr",
+}
+
+
+def _get_sampler_name_from_overrides(config_overrides: Optional[List[str]]):
+    if not config_overrides:
+        return None
+    config = OmegaConf.from_dotlist(config_overrides)
+    sampler_config = config.get("Sampler", None)
+    if sampler_config is None:
+        return None
+    return sampler_config.get("name", None)
+
+
+def _load_config(config_path: Optional[str], config_overrides: Optional[List[str]]):
+    assert (
+        config_path is not None
+    ), "config_path must be provided when model_name is None."
+
+    config = OmegaConf.load(config_path)
+    if config_overrides:
+        cli_config = OmegaConf.from_dotlist(config_overrides)
+        config = OmegaConf.merge(config, cli_config)
+    return OmegaConf.to_container(config, resolve=True)
+
+
+def _infer_sampler_name(
+    model_name: Optional[str],
+    config_path: Optional[str],
+    config_overrides: Optional[List[str]],
+):
+    sampler_name = _get_sampler_name_from_overrides(config_overrides)
+    if sampler_name is not None:
+        return sampler_name
+
+    if model_name is not None:
+        if model_name in MODEL_NAME_TO_SAMPLER:
+            return MODEL_NAME_TO_SAMPLER[model_name]
+        raise ValueError(
+            f"Unable to infer sampler type from model_name '{model_name}'. "
+            "Please add it to MODEL_NAME_TO_SAMPLER or set `Sampler.name`."
+        )
+
+    config = _load_config(config_path=config_path, config_overrides=config_overrides)
+    sampler_config = config.get("Sampler", None)
+    if sampler_config is not None:
+        sampler_name = sampler_config.get("name", None)
+        if sampler_name is not None:
+            return sampler_name
+
+    model_config = config.get("Model", None)
+    if model_config is not None:
+        model_class_name = model_config.get("__class_name__", None)
+        if model_class_name in MODEL_CLASS_TO_SAMPLER:
+            return MODEL_CLASS_TO_SAMPLER[model_class_name]
+
+    raise ValueError(
+        "Unable to infer sampler type from config. Please set `Sampler.name`."
+    )
+
+
+def _build_sampler_from_name(sampler_name: str):
+    if sampler_name not in SAMPLER_REGISTRY:
+        supported_names = ", ".join(sorted(SAMPLER_REGISTRY))
+        raise KeyError(
+            f"No sampler registered for '{sampler_name}'. "
+            f"Supported samplers: {supported_names}."
+        )
+
+    module_name, class_name = SAMPLER_REGISTRY[sampler_name].split(":")
+    module = importlib.import_module(module_name)
+    return getattr(module, class_name)
+
+
+class MolecularSampler:
+    """Build a molecular sampler from a registered or explicit configuration."""
+
+    def __new__(
+        cls,
+        model_name: Optional[str] = None,
+        weights_name: Optional[str] = None,
+        config_path: Optional[str] = None,
+        checkpoint_path: Optional[str] = None,
+        config_overrides: Optional[List[str]] = None,
+    ):
+        sampler_name = _infer_sampler_name(
+            model_name=model_name,
+            config_path=config_path,
+            config_overrides=config_overrides,
+        )
+        sampler_cls = _build_sampler_from_name(sampler_name)
+        return sampler_cls(
+            model_name=model_name,
+            weights_name=weights_name,
+            config_path=config_path,
+            checkpoint_path=checkpoint_path,
+            config_overrides=config_overrides,
+        )
diff --git a/ppmat/sampler/structure_sampler.py b/ppmat/sampler/structure_sampler.py
new file mode 100644
index 00000000..d9fc66e2
--- /dev/null
+++ b/ppmat/sampler/structure_sampler.py
@@ -0,0 +1,277 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Optional
+
+import numpy as np
+import paddle
+from omegaconf import OmegaConf
+from pymatgen.core import Composition
+from pymatgen.io.cif import CifWriter
+
+from ppmat.datasets import build_dataloader
+from ppmat.datasets.build_structure import BuildStructure
+from ppmat.datasets.transform import build_post_transforms
+from ppmat.metrics import build_metric
+from ppmat.models import build_model
+from ppmat.models import build_model_from_name
+from ppmat.utils import logger
+from ppmat.utils import save_load
+
+
+class StructureSampler:
+    """Structure Sampler.
+
+    This class provides an interface for sampling structures using pre-trained deep
+    learning models. Supports two initialization modes:
+
+    1. **Automatic Model Loading**
+       Specify `model_name` and `weights_name` to automatically download
+       and load pre-trained weights from the `MODEL_REGISTRY`.
+
+    2. **Custom Model Loading**
+       Provide explicit `config_path` and `checkpoint_path` to load
+       custom-trained models from local files.
+
+    Args:
+        model_name (Optional[str], optional): Name of the pre-defined model architecture
+            from the `MODEL_REGISTRY` registry. When specified, associated weights
+            will be automatically downloaded. Defaults to None.
+
+        weights_name (Optional[str], optional): Specific pre-trained weight identifier.
+            Used only when `model_name` is provided. Valid options include:
+            - 'best.pdparams' (highest validation performance)
+            - 'latest.pdparams' (most recent training checkpoint)
+            - Custom weight files ending with '.pdparams'
+            Defaults to None.
+
+        config_path (Optional[str], optional): Path to model configuration file (YAML)
+            for custom models. Required when not using predefined `model_name`.
+            Defaults to None.
+        checkpoint_path (Optional[str], optional): Path to model checkpoint file
+            (.pdparams) for custom models. Required when not using predefined
+            `model_name`. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        weights_name: Optional[str] = None,
+        config_path: Optional[str] = None,
+        checkpoint_path: Optional[str] = None,
+        config_overrides=None,
+    ):
+        # if model_name is not None, then config_path and checkpoint_path must be
+        # provided
+        if model_name is None:
+            assert (
+                config_path is not None and checkpoint_path is not None
+            ), "config_path and checkpoint_path must be provided when model_name is "
+            "None."
+
+            logger.info(f"Loading model from {config_path} and {checkpoint_path}.")
+
+            config = OmegaConf.load(config_path)
+            if config_overrides:
+                config = OmegaConf.merge(
+                    config, OmegaConf.from_dotlist(config_overrides)
+                )
+            config = OmegaConf.to_container(config, resolve=True)
+
+            model_config = config.get("Model", None)
+            assert model_config is not None, "Model config must be provided."
+            model = build_model(model_config)
+            save_load.load_pretrain(model, checkpoint_path)
+
+        else:
+            logger.info("Since model_name is given, downloading it...")
+            model, config = build_model_from_name(model_name, weights_name)
+            if config_overrides:
+                config = OmegaConf.merge(
+                    OmegaConf.create(config),
+                    OmegaConf.from_dotlist(config_overrides),
+                )
+                config = OmegaConf.to_container(config, resolve=True)
+
+        self.model = model
+        self.config = config
+
+        self.model.eval()
+
+        # sample config
+        sample_config = config.get("Sample", None)
+        self.sample_config = sample_config
+
+        self.post_transforms_cfg = self.sample_config.get("post_transforms", None)
+        if self.post_transforms_cfg is not None:
+            self.post_transforms = build_post_transforms(self.post_transforms_cfg)
+        else:
+            self.post_transforms = None
+
+    def compute_metric(
+        self,
+        save_path=None,
+    ):
+        metrics_cfg = self.sample_config.get("metrics")
+        assert metrics_cfg is not None, "metrics config must be provided."
+        metrics_fn = build_metric(metrics_cfg)
+
+        total_results = self.sample_by_dataloader(save_path)
+
+        metric = metrics_fn(total_results)
+        return metric
+
+    def post_process(self, data):
+        if self.post_transforms is None:
+            return data
+        return self.post_transforms(data)
+
+    def sample(self, data, sample_params=None):
+        if sample_params is None:
+            sample_params = self.sample_config.get("model_sample_params", {}) or {}
+        assert isinstance(sample_params, dict), "sample_params must be a dict or None."
+        pred_data = self.model.sample(data, **sample_params)
+        pred_data = self.post_process(pred_data)
+        return pred_data
+
+    def sample_by_dataloader(
+        self,
+        save_path=None,
+    ):
+        dataset_cfg = self.sample_config["data"]
+        data_loader = build_dataloader(dataset_cfg)
+
+        build_structure_cfg = self.sample_config["build_structure_cfg"]
+        structure_converter = BuildStructure(**build_structure_cfg)
+
+        logger.info(f"Total iterations: {len(data_loader)}")
+        logger.info("Start sampling process...\n")
+
+        total_results = []
+        for iter_id, batch_data in enumerate(data_loader):
+            pred_data = self.sample(batch_data)
+            structures = structure_converter(pred_data["result"])
+            if save_path is not None:
+                os.makedirs(save_path, exist_ok=True)
+                for i, structure in enumerate(structures):
+                    formula = structure.formula.replace(" ", "-")
+                    tar_file = os.path.join(
+                        save_path, f"{formula}_{iter_id + 1}_{i + 1}.cif"
+                    )
+                    if structure is not None:
+                        writer = CifWriter(structure)
+                        writer.write_file(tar_file)
+                    else:
+                        logger.info(
+                            f"No structure generated for iteration {iter_id}, index {i}"
+                        )
+            total_results.extend(pred_data["result"])
+        return total_results
+
+    def sample_by_num_atoms(self, num_atoms, save_path=None, sample_params=None):
+        assert isinstance(num_atoms, int), "num_atoms must be an integer."
+        if not getattr(self.model, "supports_num_atoms_sampling", True):
+            raise NotImplementedError(
+                f"{type(self.model).__name__} requires atom types. "
+                "Use sample_by_chemical_formula instead."
+            )
+        data = {
+            "structure_array": {
+                "num_atoms": paddle.to_tensor(np.array([num_atoms]).astype("int64")),
+            }
+        }
+
+        result = self.sample(data, sample_params=sample_params)
+        self._save_result(result, save_path)
+        return result
+
+    def sample_by_chemical_formula(
+        self, chemical_formula, save_path=None, sample_params=None
+    ):
+        assert isinstance(chemical_formula, str), "chemical_formula must be a string."
+        composition = Composition(chemical_formula)
+        atom_types = []
+        for elem, num in composition.items():
+            atom_types.extend([elem.Z] * int(num))
+        atom_types = np.array(atom_types).astype("int64")
+
+        data = {
+            "structure_array": {
+                "atom_types": paddle.to_tensor(atom_types),
+                "num_atoms": paddle.to_tensor(
+                    np.array([atom_types.shape[0]]).astype("int64")
+                ),
+            }
+        }
+        result = self.sample(data, sample_params=sample_params)
+        self._save_result(result, save_path)
+        return result
+
+    def sample_by_condition(
+        self,
+        num_atoms,
+        conditions,
+        save_path=None,
+        sample_params=None,
+    ):
+        """Sample structures from a conditional generation model."""
+        condition_names = getattr(self.model, "condition_names", None)
+        if not condition_names:
+            raise NotImplementedError(
+                f"{type(self.model).__name__} is not a conditional model."
+            )
+        missing = sorted(set(condition_names) - set(conditions))
+        extra = sorted(set(conditions) - set(condition_names))
+        if missing or extra:
+            raise ValueError(
+                f"Expected conditions {sorted(condition_names)}; "
+                f"missing={missing}, extra={extra}."
+            )
+
+        data = {
+            "structure_array": {
+                "num_atoms": paddle.to_tensor([num_atoms], dtype="int64"),
+            }
+        }
+        for name in condition_names:
+            value = conditions[name]
+            if isinstance(value, paddle.Tensor):
+                condition = value
+            elif isinstance(value, str):
+                condition = [value]
+            else:
+                condition = paddle.to_tensor([value], dtype="float32")
+            data[name] = condition
+
+        result = self.sample(data, sample_params=sample_params)
+        self._save_result(result, save_path)
+        return result
+
+    def _save_result(self, result, save_path):
+        if save_path is None:
+            return
+        os.makedirs(save_path, exist_ok=True)
+        logger.info(f"Save results to {save_path}")
+        structure_converter = BuildStructure(
+            **self.sample_config["build_structure_cfg"]
+        )
+        structures = structure_converter(result["result"])
+        for i, structure in enumerate(structures):
+            if structure is None:
+                logger.info(f"No structure generated for index {i}")
+                continue
+            formula = structure.formula.replace(" ", "-")
+            writer = CifWriter(structure)
+            writer.write_file(os.path.join(save_path, f"{formula}_{i + 1}.cif"))
diff --git a/ppmat/schedulers/scheduling_diffnmr.py b/ppmat/schedulers/scheduling_diffnmr.py
index d7b51dfe..1d23067f 100644
--- a/ppmat/schedulers/scheduling_diffnmr.py
+++ b/ppmat/schedulers/scheduling_diffnmr.py
@@ -521,6 +521,17 @@ def sample_discrete_feature_noise(limit_dist, node_mask):
     return ph.mask(node_mask)
 
 
+def _encode_spectrum_condition(model, condition):
+    """Encode the four-branch NMR condition through the declared interface."""
+
+    if model.flag_onlyH:
+        embedding, _ = model.encoder(condition)
+        return embedding, None
+
+    embedding, (token_encoding, _) = model.encoder(condition)
+    return embedding, token_encoding
+
+
 @paddle.no_grad()
 def step(
     model, s, t, X_t, E_t, y_t, node_mask, conditionVec, batch_X, batch_E, batch_y
@@ -564,16 +575,18 @@ def step(
         [noisy_data["y_t"].astype("float32"), extra_data.y.astype(dtype="float32")]
     )
 
-    from ppmat.models.diffnmr.diffnmr import DiffNMR
-
-    if isinstance(model, DiffNMR):
-        if model.flag_onlyH is True:
-            global_H, _ = model.encoder(conditionVec)
-            embeddings_spectrum = global_H
-        else:
-            embeddings_spectrum, (spectrum_encoding, _) = model.encoder(conditionVec)
+    if getattr(model, "conditioning_mode", None) == "spectrum":
+        embeddings_spectrum, spectrum_encoding = _encode_spectrum_condition(
+            model, conditionVec
+        )
         if model.connector_flag is True:
-            embeddings_spectrum = model.connector.sample(embeddings_spectrum, spectrum_encoding)
+            if spectrum_encoding is None:
+                raise NotImplementedError(
+                    "Connector sampling requires the joint H1/C13 spectrum encoder."
+                )
+            embeddings_spectrum = model.connector.sample(
+                embeddings_spectrum, spectrum_encoding
+            )
         input_y = paddle.concat([input_y, embeddings_spectrum], axis=1).astype(
             "float32"
         )
@@ -1006,14 +1019,8 @@ def reconstruction_logp(model, t, X, E, node_mask, condition_Spectrum):
     )
 
     ###########################################################
-    from ppmat.models.diffnmr.diffnmr import DiffNMR
-
-    if model.__class__ is DiffNMR:
-        if model.flag_onlyH is True:
-            global_H, _ = model.encoder(condition_Spectrum)
-            embeddings_spectrum = global_H
-        else:
-            embeddings_spectrum = model.encoder(condition_Spectrum)
+    if getattr(model, "conditioning_mode", None) == "spectrum":
+        embeddings_spectrum, _ = _encode_spectrum_condition(model, condition_Spectrum)
         input_y = paddle.concat([input_y, embeddings_spectrum], axis=1).astype(
             "float32"
         )
diff --git a/ppmat/utils/inference_cli.py b/ppmat/utils/inference_cli.py
new file mode 100644
index 00000000..3fc5d9f0
--- /dev/null
+++ b/ppmat/utils/inference_cli.py
@@ -0,0 +1,70 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from typing import Sequence
+
+
+def add_model_loading_arguments(parser: argparse.ArgumentParser) -> None:
+    """Add the common registered/local model loading arguments."""
+    parser.add_argument(
+        "--model_name",
+        default=None,
+        help="Registered model name.",
+    )
+    parser.add_argument(
+        "--weights_name",
+        default=None,
+        help="Optional weight filename in a model package or checkpoint directory.",
+    )
+    parser.add_argument(
+        "--config_path",
+        default=None,
+        help="Path to a local configuration file.",
+    )
+    parser.add_argument(
+        "--checkpoint_path",
+        default=None,
+        help="Path to a local checkpoint file or directory.",
+    )
+
+
+def validate_model_loading_arguments(
+    parser: argparse.ArgumentParser, args: argparse.Namespace
+) -> None:
+    """Require exactly one complete registered or local loading mode."""
+    uses_registered_model = args.model_name is not None
+    uses_local_model = args.config_path is not None or args.checkpoint_path is not None
+
+    if uses_registered_model and uses_local_model:
+        parser.error(
+            "--model_name cannot be combined with --config_path or --checkpoint_path"
+        )
+    if not uses_registered_model and not uses_local_model:
+        parser.error(
+            "provide --model_name, or both --config_path and --checkpoint_path"
+        )
+    if uses_local_model and (args.config_path is None or args.checkpoint_path is None):
+        parser.error("--config_path and --checkpoint_path must be provided together")
+
+
+def validate_config_overrides(
+    parser: argparse.ArgumentParser, config_overrides: Sequence[str]
+) -> None:
+    """Reject unknown options while allowing OmegaConf ``key=value`` overrides."""
+    invalid_overrides = [
+        value for value in config_overrides if value.startswith("-") or "=" not in value
+    ]
+    if invalid_overrides:
+        parser.error("unrecognized arguments: " + " ".join(invalid_overrides))
diff --git a/ppmat/utils/io.py b/ppmat/utils/io.py
index 28a326e8..a5377232 100644
--- a/ppmat/utils/io.py
+++ b/ppmat/utils/io.py
@@ -14,11 +14,13 @@
 
 import argparse
 import ast
+import datetime
 import hashlib
 import json
 import os
 import os.path as osp
 from typing import List
+from typing import Optional
 
 import numpy as np
 
@@ -127,6 +129,17 @@ def calc_md5(fullname):
     return calc_md5sum
 
 
+def append_timestamp_to_output_dir(
+    config,
+    now: Optional[datetime.datetime] = None,
+):
+    seed = config["Trainer"].get("seed", 42)
+    timestamp = (now or datetime.datetime.now()).strftime("%Y%m%d_%H%M%S")
+    base_output_dir = config["Trainer"]["output_dir"]
+    config["Trainer"]["output_dir"] = f"{base_output_dir}_t_{timestamp}_s_{seed}"
+    return config
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Calculate MD5 hash of a file")
     parser.add_argument("filename", help="Path to the file to hash")
diff --git a/ppmat/utils/model_package.py b/ppmat/utils/model_package.py
new file mode 100644
index 00000000..b0882527
--- /dev/null
+++ b/ppmat/utils/model_package.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+
+def resolve_model_package_dir(model_name: str, extracted_path: str) -> str:
+    """Resolve a downloaded model package without recursively guessing its layout.
+
+    A package must use one of the two layouts produced by the shared downloader:
+
+    .. code-block:: text
+
+        <extracted_path>/<model_name>.yaml
+
+    or:
+
+    .. code-block:: text
+
+        <extracted_path>/<model_name>/<model_name>.yaml
+    """
+
+    root = Path(extracted_path)
+    candidates = (root, root / model_name)
+    valid = [
+        candidate
+        for candidate in candidates
+        if candidate.is_dir()
+        and any(
+            (candidate / f"{model_name}{suffix}").is_file()
+            for suffix in (".yaml", ".yml")
+        )
+    ]
+    if len(valid) == 1:
+        return str(valid[0])
+    if len(valid) > 1:
+        raise ValueError(
+            f"Ambiguous package layout for '{model_name}' under {extracted_path}: "
+            f"{[str(path) for path in valid]}"
+        )
+    raise FileNotFoundError(
+        f"Invalid package for '{model_name}' under {extracted_path}. Expected "
+        f"'{model_name}.yaml' directly in the extracted directory or in its "
+        f"'{model_name}' child directory."
+    )
+
+
+def get_model_config_path(model_name: str, package_dir: str) -> str:
+    """Return the required model-specific YAML path from a validated package."""
+
+    package_path = Path(package_dir)
+    matches = [
+        package_path / f"{model_name}{suffix}"
+        for suffix in (".yaml", ".yml")
+        if (package_path / f"{model_name}{suffix}").is_file()
+    ]
+    if len(matches) == 1:
+        return str(matches[0])
+    if len(matches) > 1:
+        raise ValueError(
+            f"Package {package_dir} contains both YAML variants for '{model_name}'."
+        )
+    raise FileNotFoundError(
+        f"Package {package_dir} does not contain '{model_name}.yaml' or "
+        f"'{model_name}.yml'."
+    )
diff --git a/ppmat/utils/scatter.py b/ppmat/utils/scatter.py
index db055487..ef1fc329 100644
--- a/ppmat/utils/scatter.py
+++ b/ppmat/utils/scatter.py
@@ -76,11 +76,14 @@ def _scatter_sum(
     # FIXME: Paddle's put_along_axis backward (PutAlongAxisGradNode) crashes
     # for dim=0; use one-hot + matmul as drop-in replacement.
     if dim == 0:
-        # _broadcast expanded index to src.shape; collapse back to 1D via first column
-        idx_1d = index.reshape([-1, src.shape[1]])[:, 0] if index.ndim > 1 else index
+        # ``index`` is constant across all non-scatter dimensions after
+        # broadcasting. Collapse it back to one group id per source row.
+        idx_1d = index.reshape([index.shape[0], -1])[:, 0]
         one_hot = paddle.nn.functional.one_hot(idx_1d, out.shape[0]).cast(src.dtype)
-        # one_hot: [N, out_dim] -> [out_dim, N] @ [N, C] = [out_dim, C]
-        return paddle.mm(one_hot.t(), src)
+        # Flatten arbitrary trailing dimensions for matmul, then restore them.
+        flat_src = src.reshape([src.shape[0], -1])
+        flat_out = paddle.mm(one_hot.t(), flat_src)
+        return flat_out.reshape(out.shape)
     else:
         return paddle.put_along_axis(
             arr=out, indices=index, values=src, axis=dim, reduce="add"
diff --git a/ppmat/utils/visualization.py b/ppmat/utils/visualization.py
index 45799438..ce086b6a 100644
--- a/ppmat/utils/visualization.py
+++ b/ppmat/utils/visualization.py
@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 import os
 
 import imageio
 import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
+import plotly.graph_objects as go
 import rdkit
 from rdkit import Chem
 from rdkit import RDLogger
@@ -27,6 +29,13 @@
 
 from ppmat.utils import logger
 
+try:
+    from IPython.display import Image
+    from IPython.display import display
+except ImportError:
+    Image = None
+    display = None
+
 
 class MolecularVisualization:
     def __init__(self, dataset_infos, output_dir):
@@ -237,3 +246,135 @@ def visualize_chain(self, path, nodes_list, adjacency_matrix):
         )
         imgs.extend([imgs[-1]] * 10)
         imageio.mimsave(gif_path, imgs, subrectangles=True, duration=20)
+
+
+def draw_volume(
+    grid,
+    density,
+    atom_type,
+    atom_coord,
+    isomin=0.05,
+    isomax=None,
+    surface_count=5,
+    title=None,
+):
+    atom_colorscale = ["grey", "white", "red", "blue", "green"]
+
+    fig = go.Figure()
+    fig.add_trace(
+        go.Volume(
+            x=grid[..., 0],
+            y=grid[..., 1],
+            z=grid[..., 2],
+            value=density,
+            isomin=isomin,
+            isomax=isomax,
+            opacity=0.1,
+            surface_count=surface_count,
+            caps=dict(x_show=False, y_show=False, z_show=False),
+        )
+    )
+
+    axis_dict = dict(
+        showgrid=False,
+        showbackground=False,
+        zeroline=False,
+        visible=False,
+    )
+
+    fig.add_trace(
+        go.Scatter3d(
+            x=atom_coord[:, 0],
+            y=atom_coord[:, 1],
+            z=atom_coord[:, 2],
+            mode="markers",
+            marker=dict(
+                size=10,
+                color=atom_type,
+                cmin=0,
+                cmax=4,
+                colorscale=atom_colorscale,
+                opacity=0.6,
+            ),
+        )
+    )
+
+    if title is not None:
+        title = dict(
+            text=title,
+            x=0.5,
+            y=0.3,
+            xanchor="center",
+            yanchor="bottom",
+        )
+
+    fig.update_layout(
+        autosize=False,
+        width=800,
+        height=800,
+        showlegend=False,
+        scene=dict(xaxis=axis_dict, yaxis=axis_dict, zaxis=axis_dict),
+        title=title,
+        title_font_family="Times New Roman",
+    )
+
+    return fig
+
+
+def safe_write_image(fig, path, show_plot=False):
+    try:
+        fig.write_image(path)
+        logger.info(f"Image saved to: {path}")
+    except Exception as e:
+        logger.warning(f"Failed to save image {path}: {e}")
+        try:
+            html_path = path.with_suffix(".html")
+            fig.write_html(html_path)
+            logger.info(f"Saved interactive HTML instead: {html_path}")
+        except Exception as html_e:
+            logger.warning(f"Failed to save HTML fallback for {path}: {html_e}")
+
+    if show_plot:
+        try:
+            if Image is None or display is None:
+                raise ImportError("IPython is required to display image.")
+            img_bytes = fig.to_image(format="png", scale=2)
+            display(Image(img_bytes))
+        except Exception as e:
+            logger.warning(f"Failed to display image: {e}")
+
+
+def maybe_downsample_volume(grid, values, shape, max_points=250_000):
+    """
+    Downsample a regular 3D grid for visualization to keep Plotly volume
+    traces responsive.
+    """
+    if shape is None or len(shape) != 3:
+        return grid, values, False, 1
+
+    try:
+        shape = [int(s) for s in shape]
+        total = shape[0] * shape[1] * shape[2]
+    except Exception:
+        return grid, values, False, 1
+
+    if total != grid.shape[0] or any(val.shape[0] != grid.shape[0] for val in values):
+        return grid, values, False, 1
+    if total <= max_points:
+        return grid, values, False, 1
+
+    stride = max(1, math.ceil((total / max_points) ** (1 / 3)))
+    try:
+        grid_view = grid.reshape(shape[0], shape[1], shape[2], 3)
+        grid_ds = grid_view[::stride, ::stride, ::stride, :].reshape(-1, 3)
+        values_ds = [
+            val.reshape(shape[0], shape[1], shape[2])[
+                ::stride, ::stride, ::stride
+            ].reshape(-1)
+            for val in values
+        ]
+    except Exception as e:
+        logger.warning(f"Failed to downsample grid for visualization: {e}")
+        return grid, values, False, 1
+
+    return grid_ds, values_ds, True, stride
diff --git a/property_prediction/train.py b/property_prediction/train.py
index cde90a56..d7116388 100644
--- a/property_prediction/train.py
+++ b/property_prediction/train.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
-import datetime
 import os
 import os.path as osp
 
@@ -29,6 +28,7 @@
 from ppmat.trainer.base_trainer import BaseTrainer
 from ppmat.utils import logger
 from ppmat.utils import misc
+from ppmat.utils.io import append_timestamp_to_output_dir
 
 
 def read_independent_dataloader_config(config):
@@ -91,10 +91,7 @@ def read_independent_dataloader_config(config):
     misc.set_random_seed(seed)
     logger.info(f"Set random seed to {seed}")
 
-    # add timestamp to output_dir
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    base_output_dir = config["Trainer"]["output_dir"]
-    config["Trainer"]["output_dir"] = f"{base_output_dir}_t_{timestamp}_s_{seed}"
+    append_timestamp_to_output_dir(config)
 
     # save config to output_dir, only rank 0 process will do this
     if dist.get_rank() == 0:
diff --git a/requirements.txt b/requirements.txt
index 0c8da8a9..f60bf4e4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,6 +3,7 @@ colorlog==6.8.2
 Cython==3.0.12
 hydra-core==1.3.2
 einops==0.8.0
+imageio==2.37.3
 importlib_metadata==8.6.1
 jarvis_tools==2025.5.30
 lmdb==1.7.5
diff --git a/research/ML2DDB/README.md b/research/ML2DDB/README.md
index f4d4e910..a287baff 100644
--- a/research/ML2DDB/README.md
+++ b/research/ML2DDB/README.md
@@ -65,7 +65,7 @@ Refer to the [install doc](../../Install.md) to install PaddleMaterials.
 
 #### 1. Generate new 2D materials:
    ```bash
-   python structure_generation/sample.py --model_name='mattergen_ml2ddb' --mode='by_dataloader' --save_path='results_mattergen_ml2ddb'
+   python structure_generation/sample.py --model_name='mattergen_ml2ddb' --mode='by_dataloader' --output_dir='results_mattergen_ml2ddb'
    ```
    You can download the pre-trained model from [here](https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb.zip) and modify the `total_num` parameter in the configuration file to generate more structures.
 
@@ -80,13 +80,13 @@ Refer to the [install doc](../../Install.md) to install PaddleMaterials.
    ```
    Then you can generate more structures:
    ```bash
-   python structure_generation/sample.py --config_path='your config path after modify' --checkpoint_path='your downloaded checkpoint path(*.pdparams)'  --mode='by_dataloader' --save_path='results_mattergen_ml2ddb'
+   python structure_generation/sample.py --config_path='your config path after modify' --checkpoint_path='your downloaded checkpoint path(*.pdparams)'  --mode='by_dataloader' --output_dir='results_mattergen_ml2ddb'
    ```
 
 
 #### 2. Generate new 2D materials with specific chemical system:
    ```bash
-   python structure_generation/sample.py --model_name='mattergen_ml2ddb_chemical_system' --mode='by_dataloader' --save_path='results_mattergen_ml2ddb_chemical_system'
+   python structure_generation/sample.py --model_name='mattergen_ml2ddb_chemical_system' --mode='by_dataloader' --output_dir='results_mattergen_ml2ddb_chemical_system'
    ```
    The above command will generate structures with specific chemical system `Si` and `Mo`, if you want to generate structures with other chemical systems, you can download the pre-trained model from [here](https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_chemical_system.zip) and modify the `prop_values` parameter in the configuration file.
 
@@ -101,13 +101,13 @@ Refer to the [install doc](../../Install.md) to install PaddleMaterials.
    ```
    Then you can generate structures with specific chemical system:
    ```bash
-   python structure_generation/sample.py --config_path='your config path after modify' --checkpoint_path='your downloaded checkpoint path(*.pdparams)'  --mode='by_dataloader' --save_path='results_mattergen_ml2ddb_chemical_system'
+   python structure_generation/sample.py --config_path='your config path after modify' --checkpoint_path='your downloaded checkpoint path(*.pdparams)'  --mode='by_dataloader' --output_dir='results_mattergen_ml2ddb_chemical_system'
    ```
 
 
 #### 3. Generate new 2D materials with specific space group:
    ```bash
-   python structure_generation/sample.py --model_name='mattergen_ml2ddb_space_group' --mode='by_dataloader' --save_path='results_mattergen_ml2ddb_space_group'
+   python structure_generation/sample.py --model_name='mattergen_ml2ddb_space_group' --mode='by_dataloader' --output_dir='results_mattergen_ml2ddb_space_group'
    ```
    The above command will generate structures with specific space group `11`, if you want to generate structures with other space groups, you can download the pre-trained model from [here](https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_space_group.zip) and modify the `prop_values` parameter in the configuration file.
 
@@ -122,7 +122,7 @@ Refer to the [install doc](../../Install.md) to install PaddleMaterials.
    ```
    Then you can generate structures with specific space group:
    ```bash
-   python structure_generation/sample.py --config_path='your config path after modify' --checkpoint_path='your downloaded checkpoint path(*.pdparams)'  --mode='by_dataloader' --save_path='results_mattergen_ml2ddb_space_group'
+   python structure_generation/sample.py --config_path='your config path after modify' --checkpoint_path='your downloaded checkpoint path(*.pdparams)'  --mode='by_dataloader' --output_dir='results_mattergen_ml2ddb_space_group'
    ```
 
 
diff --git a/spectrum_elucidation/configs/diffnmr/DiffNMR.yaml b/spectrum_elucidation/configs/diffnmr/DiffNMR.yaml
index 00fcc8d8..0b4ff834 100644
--- a/spectrum_elucidation/configs/diffnmr/DiffNMR.yaml
+++ b/spectrum_elucidation/configs/diffnmr/DiffNMR.yaml
@@ -95,27 +95,30 @@ Trainer:
       sample: ["Accuracy"]
 
 Sampler:
+  name: diffnmr
   sample_every_val: 500
-  visual_num: 10 # number of visualize molecules in the sample
-  chains_to_save: 5 # less sample batch size, select representative sample in sample batch for visualize
+  visual_num: 1 # number of visualize molecules in the sample
+  chains_to_save: 0 # less sample batch size, select representative sample in sample batch for visualize
   number_chain_steps: 10 # Number of frames in each gif
-  sample_batch_iters: 100 # Number of sample batches
-  flag_retrival_sampling: False
+  sample_batch_iters: 1 # Number of sample batches
+  flag_retrieval_sampling: False
   flag_use_formula: False
-  flag_retrival_initilization: False
+  flag_retrieval_initialization: False
   num_candidates: 1
-  retrival_database_path: ./spectrum_elucidation/retrival_database/mol_rep_15.csv
-  pretrained_model_path: ./pretrained/DiffNMR_nless15_best.pdparams
+  retrieval_database_path: null
+  pretrained_model_path: ./checkpoints/best.pdparams
   output_dir: ./output/DiffNMR/DiffNMR/sample
   out_dict: {"Accuracy"}
   data:
     dataset:
       __class_name__: MSDnmrDataset
       __init_params__:
-        path: "./data/MSD_nmr/test.csv"
+        path: "./example/sample.csv"
         vocab_peakwidth_path: "./spectrum_elucidation/vocab/nless15/H1_statistic/delta_distribution.csv"
         vocab_split_path: "./spectrum_elucidation/vocab/nless15/H1_statistic/split_type_distribution.csv"
         cache: True
+        cache_path: "./output/diffnmr_example_cache"
+        overwrite: True
         data_flag: "n<15"
         max_atoms: 15
         build_molecule_cfg: ${Global.molecule_converter}
@@ -126,7 +129,7 @@ Sampler:
       __init_params__:
         shuffle: True
         drop_last: False
-        batch_size: 256
+        batch_size: 1
 
 
 Model:
@@ -134,7 +137,7 @@ Model:
   __init_params__:
     encoder_cfg:
       __name__: NMRNetCLIP
-      pretrained_path: "./pretrained/DiffNMR_NMRNet_nless15_best.pdparams"
+      pretrained_path: "./checkpoints/DiffNMR_NMRNet_nless15_best.pdparams"
       dim_enc_H: 1024
       dimff_enc_H: 2048
       dim_enc_C: 256
@@ -150,7 +153,7 @@ Model:
       onlyH: False # True if set NMRNet only encode H Spectrum info
     decoder_cfg:
       __name__: DiffGraphFormer
-      pretrained_path: "./pretrained/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+      pretrained_path: "./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
       hidden_mlp_dims: {
         'X': 256,
         'E': 128,
@@ -182,7 +185,7 @@ CLIP:
     __class_name__: NMRNetCLIP
     __init_params__:
       spectrum_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_NMRNet_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_NMRNet_nless15_best.pdparams"
         dim_enc_H: 1024
         dimff_enc_H: 2048
         dim_enc_C: 256
@@ -196,7 +199,7 @@ CLIP:
         peakwidthemb_num: 70
         integralemb_num: 26
       graph_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
         n_layers_GT: 5
         hidden_mlp_dims: {
           "X": 256,
@@ -296,4 +299,4 @@ Dataset:
 
 DataInfo:
     extra_features: 'all' # 'all', 'cycles', 'eigenvalues' or null
-    conditdim: 512
\ No newline at end of file
+    conditdim: 512
diff --git a/spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml b/spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml
index 8d9a486d..71c21f27 100644
--- a/spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml
+++ b/spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml
@@ -51,7 +51,7 @@ Trainer:
   # Evaluation frequency [epoch], for example, eval_freq=1 means evaluate every 1 epoch
   eval_freq: 1 # set 0 to disable evaluation during training
   # Pretrained model path, if null, no pretrained model will be loaded
-  pretrained_model_path: ./pretrained/DiffNMR_DiffGraphFormer_nless15_init.pdparams # please set your pretrained model path here when run trainer.test
+  pretrained_model_path: ./checkpoints/DiffNMR_DiffGraphFormer_nless15_init.pdparams # please set your pretrained model path here when run trainer.test
   # Resume from checkpoint path, useful for resuming training
   resume_from_checkpoint: null
   # whether use automatic mixed precision
@@ -99,12 +99,12 @@ Sampler:
   chains_to_save: 5 # less sample batch size, select representative sample in sample batch for visualize
   number_chain_steps: 10 # Number of frames in each gif
   sample_batch_iters: 1 # Number of sample batches
-  flag_retrival_sampling: False
+  flag_retrieval_sampling: False
   flag_use_formula: False
-  flag_retrival_initilization: False
+  flag_retrieval_initialization: False
   num_candidates: 1
-  retrival_database_path: ./spectrum_elucidation/retrival_database/mol_rep_15.csv
-  pretrained_model_path: ./pretrained/DiffNMR_DiffGraphFormer_nless15_best.pdparams
+  retrieval_database_path: ./spectrum_elucidation/retrieval_database/mol_rep_15.csv
+  pretrained_model_path: ./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams
   output_dir: ./output/DiffNMR/DiffGraphFormer/sample
   out_dict: {"Accuracy"}
   data:
@@ -178,7 +178,7 @@ CLIP:
     __class_name__: NMRNetCLIP
     __init_params__:
       spectrum_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_NMRNet_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_NMRNet_nless15_best.pdparams"
         dim_enc_H: 1024
         dimff_enc_H: 2048
         dim_enc_C: 256
@@ -192,7 +192,7 @@ CLIP:
         peakwidthemb_num: 70
         integralemb_num: 26
       graph_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
         n_layers_GT: 5
         hidden_mlp_dims: {
           "X": 256,
@@ -292,4 +292,4 @@ Dataset:
 
 DataInfo:
     extra_features: 'all' # 'all', 'cycles', 'eigenvalues' or null
-    conditdim: 512
\ No newline at end of file
+    conditdim: 512
diff --git a/spectrum_elucidation/configs/diffnmr/PP-DiffNMR.yaml b/spectrum_elucidation/configs/diffnmr/PP-DiffNMR.yaml
index 82fda86b..36d2202e 100644
--- a/spectrum_elucidation/configs/diffnmr/PP-DiffNMR.yaml
+++ b/spectrum_elucidation/configs/diffnmr/PP-DiffNMR.yaml
@@ -100,12 +100,12 @@ Sampler:
   chains_to_save: 5 # less sample batch size, select representative sample in sample batch for visualize
   number_chain_steps: 10 # Number of frames in each gif
   sample_batch_iters: 100 # Number of sample batches
-  flag_retrival_sampling: False
+  flag_retrieval_sampling: False
   flag_use_formula: False
-  flag_retrival_initilization: False
+  flag_retrieval_initialization: False
   num_candidates: 1
-  retrival_database_path: ./spectrum_elucidation/retrival_database/mol_rep_15.csv
-  pretrained_model_path: ./pretrained/DiffNMR_nless15_best.pdparams
+  retrieval_database_path: ./spectrum_elucidation/retrieval_database/mol_rep_15.csv
+  pretrained_model_path: ./checkpoints/DiffNMR_nless15_best.pdparams
   output_dir: ./output/DiffNMR/DiffNMR/sample
   out_dict: {"Accuracy"}
   data:
@@ -134,7 +134,7 @@ Model:
   __init_params__:
     encoder_cfg:
       __name__: NMRNetCLIP
-      pretrained_path: ./output/DiffNMR/NMRNet/checkpoints/best.pdparams # "./pretrained/step2_best.pdparams" #"/home/liuxuwei01/PaddleMaterial/output/step2_init_weight.pdparams" #"./output/step2_best.pdparams"
+      pretrained_path: ./output/DiffNMR/NMRNet/checkpoints/best.pdparams
       dim_enc_H: 1024
       dimff_enc_H: 2048
       dim_enc_C: 256
@@ -149,7 +149,7 @@ Model:
       integralemb_num: 26
     decoder_cfg:
       __name__: DiffGraphFormer
-      pretrained_path: ./output/DiffNMR/DiffGraphFormer/checkpoints/latest.pdparams # "./pretrained/step1_best.pdparams"
+      pretrained_path: ./output/DiffNMR/DiffGraphFormer/checkpoints/latest.pdparams
       hidden_mlp_dims: {
         'X': 256,
         'E': 128,
@@ -212,7 +212,7 @@ CLIP:
     __class_name__: NMRNetCLIP
     __init_params__:
       spectrum_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_NMRNet_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_NMRNet_nless15_best.pdparams"
         dim_enc_H: 1024
         dimff_enc_H: 2048
         dim_enc_C: 256
@@ -226,7 +226,7 @@ CLIP:
         peakwidthemb_num: 70
         integralemb_num: 26
       graph_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
         n_layers_GT: 5
         hidden_mlp_dims: {
           "X": 256,
@@ -323,4 +323,4 @@ Dataset:
 
 DataInfo:
     extra_features: 'all' # 'all', 'cycles', 'eigenvalues' or null
-    conditdim: 512
\ No newline at end of file
+    conditdim: 512
diff --git a/spectrum_elucidation/configs/diffnmr/PP-DiffNMR_DiffPrior.yaml b/spectrum_elucidation/configs/diffnmr/PP-DiffNMR_DiffPrior.yaml
index 308184af..4ae53f9a 100644
--- a/spectrum_elucidation/configs/diffnmr/PP-DiffNMR_DiffPrior.yaml
+++ b/spectrum_elucidation/configs/diffnmr/PP-DiffNMR_DiffPrior.yaml
@@ -100,12 +100,12 @@ Sampler:
   chains_to_save: 5 # less sample batch size, select representative sample in sample batch for visualize
   number_chain_steps: 10 # Number of frames in each gif
   sample_batch_iters: 100 # Number of sample batches
-  flag_retrival_sampling: False
+  flag_retrieval_sampling: False
   flag_use_formula: False
-  flag_retrival_initilization: False
+  flag_retrieval_initialization: False
   num_candidates: 1
-  retrival_database_path: ./spectrum_elucidation/retrival_database/mol_rep_15.csv
-  pretrained_model_path: ./pretrained/DiffNMR_nless15_best.pdparams
+  retrieval_database_path: ./spectrum_elucidation/retrieval_database/mol_rep_15.csv
+  pretrained_model_path: ./checkpoints/DiffNMR_nless15_best.pdparams
   output_dir: ./output/PP-DiffNMR/DiffNMR/sample
   out_dict: {"Accuracy"}
   data:
@@ -166,7 +166,7 @@ CLIP:
     __class_name__: NMRNetCLIP
     __init_params__:
       spectrum_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_NMRNet_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_NMRNet_nless15_best.pdparams"
         dim_enc_H: 1024
         dimff_enc_H: 2048
         dim_enc_C: 256
@@ -180,7 +180,7 @@ CLIP:
         peakwidthemb_num: 70
         integralemb_num: 26
       graph_encoder:
-        pretrained_model_path: "./pretrained/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+        pretrained_model_path: "./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
         n_layers_GT: 5
         hidden_mlp_dims: {
           "X": 256,
@@ -278,4 +278,4 @@ Dataset:
 
 DataInfo:
     extra_features: 'all' # 'all', 'cycles', 'eigenvalues' or null
-    conditdim: 512
\ No newline at end of file
+    conditdim: 512
diff --git a/spectrum_elucidation/configs/diffnmr/README.md b/spectrum_elucidation/configs/diffnmr/README.md
index 0dbe3df1..9fe58912 100644
--- a/spectrum_elucidation/configs/diffnmr/README.md
+++ b/spectrum_elucidation/configs/diffnmr/README.md
@@ -4,162 +4,160 @@
 
 ## Abstract
 
-Nuclear Magnetic Resonance (NMR) spectroscopy is a central characterization method for molecular structure elucidation, yet interpreting NMR spectra to deduce molecular structures remains challenging due to the complexity of spectral data and the vastness of the chemical space. In this work, we introduce DiffNMR, a novel end-to-end framework that leverages a conditional discrete diffusion model for de novo molecular structure elucidation from NMR spectra. DiffNMR refines molecular graphs iteratively through a diffusion-based generative process, ensuring global consistency and mitigating error accumulation inherent in autoregressive methods. The framework integrates a two-stage pretraining strategy that aligns spectral and molecular representations via diffusion autoencoder (Diff-AE) and contrastive learning, the incorporation of retrieval initialization and similarity filtering during inference, and a specialized NMR encoder with radial basis function (RBF) encoding for chemical shifts, preserving continuity and chemical correlation. Experimental results demonstrate that DiffNMR achieves competitive performance for NMR-based structure elucidation, offering an efficient and robust solution for automated molecular analysis.
+Nuclear Magnetic Resonance (NMR) spectroscopy is a central characterization method
+for molecular structure elucidation. DiffNMR formulates NMR-based structure
+elucidation as conditional molecular graph generation and uses a discrete diffusion
+model to iteratively refine molecular graphs from spectrum conditions.
 
 ![DiffNMR Overview](../../docs/diffnmr_overview.png)
 
-## Datasets:
+---
 
-- MSD-NMR:
+## Model Description
 
-    MSD-NMR Multimodal-Spectroscopic-Dataset (MSD-NMR) is a comprehensive dataset for molecular structure elucidation from NMR spectra. It contains 121,509 spectra, each corresponding to a molecular structure with up to 15 heavy atoms. Up to 574,799 spectra with up to 35 heavy atoms. The dataset is divided into training, validation, and test sets.
+### Overview
 
-    | Dataset | train | val | test | total |
-    |:--------|------:|----:|-----:|------:|
-    | [MSD-NMR](https://paddle-org.bj.bcebos.com/paddlematerial/datasets/msd/msd_nmr.zip) |  |  |  |  |
-    | n<15    | 109,358 | 6,076  | 6,075  | 121,509 |
-    | n<20    | 235,512 | 13,085 | 13,084 | 261,681 |
-    | n<25    | 351,273 | 19,516 | 19,515 | 390,304 |
-    | n<35    | 517,319 | 28,741 | 28,739 | 574,799 |
+DiffNMR is an end-to-end molecular generation model for NMR spectrum
+elucidation. Given tokenized `1H` and `13C` NMR signals, the model generates
+candidate molecular graphs and evaluates generated molecules against the input
+spectrum condition.
 
-## Data Preparation
+The framework contains:
 
-To set up the DiffNMR environment, please follow these steps:
+- `NMRNetCLIP`: spectrum encoder for NMR representation learning
+- `MolecularGraphFormer`: molecular graph encoder for graph representation learning
+- `DiffNMR`: conditional discrete diffusion model for molecular graph generation
 
-1. Download the required files:
-   - Vocabulary list: [vocab.tar.gz](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/vocab.tar.gz)
-   - Retrieval database: [retrival_database.zip](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/retrival_database.zip)
+### Method
 
-2. Place the downloaded files in the `spectrum_elucidation` directory
+DiffNMR uses a two-stage pretraining and fine-tuning workflow:
 
-3. Decompress the files using the following commands:
-   ```bash
-   tar -xvzf vocab.tar.gz
-   unzip retrival_database.zip
-   ```
+1. Pretrain the molecular encoder and decoder with a diffusion autoencoder.
+2. Pretrain the NMR spectrum encoder with contrastive learning.
+3. Fine-tune the conditional diffusion model for spectrum-conditioned molecular
+   graph generation.
+
+During sampling, the model starts from noisy discrete graph features and denoises
+them step by step under the NMR spectrum condition.
+
+---
+
+## Dataset Description
+
+### MSD-NMR
+
+MSD-NMR is a multimodal spectroscopic dataset for molecular structure
+elucidation. PaddleMaterials uses the preprocessed CSV format with the following
+columns:
+
+- `smiles`: molecular SMILES
+- `tokenized_input`: JSON string containing `1HNMR` and `13CNMR`
+- `atom_count`: number of atoms for filtering and batching
+
+| Dataset | Train | Val | Test | Total |
+| --- | ---: | ---: | ---: | ---: |
+| MSD-NMR n<15 | 109,358 | 6,076 | 6,075 | 121,509 |
+| MSD-NMR n<20 | 235,512 | 13,085 | 13,084 | 261,681 |
+| MSD-NMR n<25 | 351,273 | 19,516 | 19,515 | 390,304 |
+| MSD-NMR n<35 | 517,319 | 28,741 | 28,739 | 574,799 |
+
+### Data Preparation
+
+Download the dataset and support files:
+
+- [MSD-NMR dataset](https://paddle-org.bj.bcebos.com/paddlematerial/datasets/msd/msd_nmr.zip)
+- [Vocabulary files](https://paddle-org.bj.bcebos.com/paddlematerials/assets/vocabs/msd_nmr_vocab.zip)
+- [MSD-NMR n<15 retrieval database](https://paddle-org.bj.bcebos.com/paddlematerials/assets/databases/msd_nmr_nless15_retrieval_molecular_representations.zip)
+- [MSD-NMR n<20 retrieval database](https://paddle-org.bj.bcebos.com/paddlematerials/assets/databases/msd_nmr_nless20_retrieval_molecular_representations.zip)
+
+Place the files under the repository root and extract them:
+
+```bash
+mkdir -p data spectrum_elucidation
+unzip msd_nmr.zip -d data
+unzip msd_nmr_vocab.zip -d spectrum_elucidation
+unzip msd_nmr_nless15_retrieval_molecular_representations.zip -d spectrum_elucidation
+```
+
+The default configs expect:
+
+```text
+data/MSD_nmr/train.csv
+data/MSD_nmr/val.csv
+data/MSD_nmr/test.csv
+spectrum_elucidation/vocab/nless15/
+spectrum_elucidation/retrieval_database/
+```
+
+For a quick sampling smoke test, a bundled one-row sample from the MSD-NMR n<15
+validation split is provided at `example/sample.csv`.
+
+---
 
 ## Results
 
-<table>
-    <head>
-        <tr>
-            <th  nowrap="nowrap">Model</th>
-            <th  nowrap="nowrap">Dataset</th>
-            <th  nowrap="nowrap">Loss</th>
-            <th  nowrap="nowrap">Negative log likelihood</th>
-            <th  nowrap="nowrap">GPUs</th>
-            <th  nowrap="nowrap">Training time</th>
-            <th  nowrap="nowrap">Config</th>
-            <th  nowrap="nowrap">Checkpoint | Log</th>
-        </tr>
-    </head>
-    <body>
-        <tr>
-            <td  nowrap="nowrap">diffnmr_diffgraphfromer_msdnmr_nless15</td>
-            <td  nowrap="nowrap">msdnmr_nless15</td>
-            <td  nowrap="nowrap">1.946618</td>
-            <td  nowrap="nowrap">66.028621</td>
-            <td  nowrap="nowrap">4</td>
-            <td  nowrap="nowrap">~34.15 hours</td>
-            <td  nowrap="nowrap"><a href="DiffNMR_DiffGraphFormer.yaml">DiffNMR_DiffGraphFormer</a></td>
-            <td  nowrap="nowrap"><a href="">checkpoint | log</a></td>
-        </tr>  
-    </body>
-    <body>
-        <tr>
-            <td  nowrap="nowrap">diffnmr_nmrnet_msdnmr_nless15</td>
-            <td  nowrap="nowrap">msdnmr_nless15</td>
-            <td  nowrap="nowrap">3.217951</td>
-            <td  nowrap="nowrap">-</td>
-            <td  nowrap="nowrap">4</td>
-            <td  nowrap="nowrap">~6.5 hours</td>
-            <td  nowrap="nowrap"><a href="DiffNMR_DiffGraphFormer.yaml">DiffNMR_NMRNet</a></td>
-            <td  nowrap="nowrap"><a href="">checkpoint | log</a></td>
-        </tr>  
-    </body>
-    <body>
-        <tr>
-            <td  nowrap="nowrap">diffnmr_msdnmr_nless15</td>
-            <td  nowrap="nowrap">msdnmr_nless15</td>
-            <td  nowrap="nowrap">1.946618</td>
-            <td  nowrap="nowrap">66.028621</td>
-            <td  nowrap="nowrap">4</td>
-            <td  nowrap="nowrap">~30.24 hours</td>
-            <td  nowrap="nowrap"><a href="DiffNMR_DiffGraphFormer.yaml">DiffNMR</a></td>
-            <td  nowrap="nowrap"><a href="">checkpoint | log</a></td>
-        </tr>  
-    </body>
-</table>
-
-Note: please refer to the following pretrained weights:
-- [DiffNMR_DiffGraphFormer_nless15_best.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_DiffGraphFormer_nless15_best.pdparams)
-- [DiffNMR_DiffGraphFormer_nless15_init.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_DiffGraphFormer_nless15_init.pdparams)
-- [DiffNMR_NMRNet_nless15_best.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_NMRNet_nless15_best.pdparams)
-- [DiffNMR_NMRNet_nless15_init.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_NMRNet_nless15_init.pdparams)
-- [DiffNMR_NMRNet_nless15_init_v2.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_NMRNet_nless15_init_v2.pdparams)
-- [DiffNMR_nless15_best.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_nless15_best.pdparams)
-- [DiffNMR_nless15_onlyH_best.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/DiffNMR_nless15_onlyH_best.pdparams) 
+| Model Name | Dataset | Loss | Negative Log Likelihood | GPUs | Training Time | Config | Checkpoint / Log |
+| --- | --- | ---: | ---: | ---: | --- | --- | --- |
+| diffnmr_diffgraphformer_msdnmr_nless15 | MSD-NMR n<15 | 1.946618 | 66.028621 | 4 | ~34.15 hours | [DiffNMR_DiffGraphFormer.yaml](DiffNMR_DiffGraphFormer.yaml) | - |
+| diffnmr_nmrnet_msdnmr_nless15 | MSD-NMR n<15 | 3.217951 | - | 4 | ~6.5 hours | [DiffNMR_NMRNet.yaml](DiffNMR_NMRNet.yaml) | - |
+| diffnmr_msdnmr_nless15 | MSD-NMR n<15 | 1.946618 | 66.028621 | 4 | ~30.24 hours | [DiffNMR.yaml](DiffNMR.yaml) | [checkpoint](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_elucidation/diffnmr/diffnmr_msdnmr_nless15.zip) |
+
+---
+
+## Command
 
 ### Training
+
 ```bash
-## 2 stage pretraining
-### stage 1: pretrain Diff-AE of Molecular Encoder and Molecular Decoder
-# multi-gpu training, we use 4 gpus here
+# stage 1: pretrain molecular encoder and decoder
 python -m paddle.distributed.launch --gpus="0,1,2,3" spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml
-# single-gpu training
 python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml
-### stage 2: pretrain NMR Spectrum Encoder NMRNet by CLIP
+
+# stage 2: pretrain NMR spectrum encoder
 python -m paddle.distributed.launch --gpus="0,1,2,3" spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_NMRNet.yaml
-# single-gpu training
 python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_NMRNet.yaml
-## fine-tuning
-# multi-gpu training, we use 4 gpus here
+
+# fine-tune DiffNMR
 python -m paddle.distributed.launch --gpus="0,1,2,3" spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR.yaml
-# single-gpu training
 python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR.yaml
 ```
 
 ### Validation
+
 ```bash
-# Adjust program behavior on-the-fly using command-line parameters – this provides a convenient way to customize settings without modifying the configuration file directly.
-# such as: --Global.do_eval=True
-## 2 stage pretraining
-### stage 1: pretrain Diff-AE of Molecular Encoder and Molecular Decoder
-python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='your model path(*.pdparams)'
-### stage 2: pretrain NMR Spectrum Encoder NMRNet by CLIP
-python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_NMRNet.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='your model path(*.pdparams)'
-## fine-tuning
-python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='your model path(*.pdparams)'
+python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='path/to/model.pdparams'
 ```
 
 ### Testing
+
 ```bash
-# This command is used to evaluate the model's performance on the test dataset.
-## 2 stage pretraining
-### stage 1: pretrain Diff-AE of Molecular Encoder and Molecular Decoder
-python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_DiffGraphFormer.yaml Global.do_eval=False Global.do_train=False Global.do_test=True Trainer.pretrained_model_path='your model path(*.pdparams)'
-### stage 2: pretrain NMR Spectrum Encoder NMRNet by CLIP
-python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR_NMRNet.yaml Global.do_eval=False Global.do_train=False Global.do_test=True Trainer.pretrained_model_path='your model path(*.pdparams)'
-## fine-tuning
-python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR.yaml Global.do_eval=False Global.do_train=False Global.do_test=True Trainer.pretrained_model_path='your model path(*.pdparams)'
+python spectrum_elucidation/train.py -c spectrum_elucidation/configs/diffnmr/DiffNMR.yaml Global.do_eval=False Global.do_train=False Global.do_test=True Trainer.pretrained_model_path='path/to/model.pdparams'
 ```
 
 ### Sample
+
 ```bash
-# This command is used to predict the  crystal structure using a trained model.
-# Note: The model_name and weights_name parameters are used to specify the pre-trained model and its corresponding weights. 
-# The prediction results will be saved in the folder specified by the `save_path` parameter, with the default set to `result`.
+# This command is used to sample molecular structures conditioned on NMR spectra.
+# Mode 1: Use a self-contained pre-trained model package.
+# Mode 2: Use a custom configuration file with local checkpoints and vocabulary.
+# Results are saved to the folder specified by --output_dir (default: results).
 
-# Mode 1: Use a custom configuration file and checkpoint for crystal structure prediction. This approach allows for more flexibility and customization.
-python spectrum_elucidation/sample.py --config_path='spectrum_elucidation/configs/diffnmr/DiffNMR.yaml' --weights_name='DiffNMR_nless15_best.pdparams' --save_path='result_diffnmr_nless15/' --checkpoint_path="pretrained"
+python spectrum_elucidation/sample.py --model_name='diffnmr_msdnmr_nless15' --weights_name='best.pdparams' --output_dir='result_diffnmr_sample/'
+python spectrum_elucidation/sample.py --config_path='spectrum_elucidation/configs/diffnmr/DiffNMR.yaml' --checkpoint_path='./checkpoints' --weights_name='best.pdparams' --output_dir='result_diffnmr_sample/'
 
+# By default, sampling uses the bundled one-row validation example.
+# To sample the full test split, override:
+# Sampler.data.dataset.__init_params__.path='./data/MSD_nmr/test.csv'
 ```
 
+---
+
 ## Citation
-```
+
+```bibtex
 @article{yang2025diffnmr,
   title={DiffNMR: Diffusion Models for Nuclear Magnetic Resonance Spectra Elucidation},
-  author= {Yang, Qingsong and Wu, Binglan and Liu, Xuwei and Chen, Bo and Li, Wei and Long, Gen and Chen, Xin and Xiao, Mingjun},
+  author={Yang, Qingsong and Wu, Binglan and Liu, Xuwei and Chen, Bo and Li, Wei and Long, Gen and Chen, Xin and Xiao, Mingjun},
   journal={arXiv preprint arXiv:2507.08854},
   year={2025}
 }
diff --git a/spectrum_elucidation/configs/diffnmr/example/sample.csv b/spectrum_elucidation/configs/diffnmr/example/sample.csv
new file mode 100644
index 00000000..78ebb459
--- /dev/null
+++ b/spectrum_elucidation/configs/diffnmr/example/sample.csv
@@ -0,0 +1,2 @@
+smiles,tokenized_input,atom_count
+CSc1ccc(C(C)C(=O)O)cc1F,"{""1HNMR"": [[7.27, 0.11, ""m"", ""2H"", []], [7.04, 0.06, ""ddd"", ""1H"", [0.74, 2.04, 12.23]], [3.85, 0.08, ""dt"", ""1H"", [7.18, 7.89]], [2.46, 0.03, ""s"", ""3H"", []], [1.5, 0.05, ""d"", ""4H"", [7.14]]], ""13CNMR"": [180.3, 162.4, 160.4, 140.3, 140.2, 128.6, 128.6, 125.8, 125.7, 124.4, 124.4, 116.4, 116.2, 44.9, 44.9, 18.1, 16.4, 16.4]}",14
diff --git a/spectrum_elucidation/configs/diffnmr/spectrum_elucidation/vocab/nless15/H1_statistic/delta_distribution.csv b/spectrum_elucidation/configs/diffnmr/spectrum_elucidation/vocab/nless15/H1_statistic/delta_distribution.csv
new file mode 100644
index 00000000..7acf3b22
--- /dev/null
+++ b/spectrum_elucidation/configs/diffnmr/spectrum_elucidation/vocab/nless15/H1_statistic/delta_distribution.csv
@@ -0,0 +1,68 @@
+Value,Count
+0.0,383
+0.01,7164
+0.02,70351
+0.03,116454
+0.04,139756
+0.05,123225
+0.06,124624
+0.07,92986
+0.08,72088
+0.09,48022
+0.1,36717
+0.11,23754
+0.12,14381
+0.13,8728
+0.14,5245
+0.15,3636
+0.16,2796
+0.17,1972
+0.18,1466
+0.19,1089
+0.2,868
+0.21,658
+0.22,531
+0.23,494
+0.24,549
+0.25,595
+0.26,396
+0.27,275
+0.28,244
+0.29,198
+0.3,190
+0.31,144
+0.32,123
+0.33,132
+0.34,102
+0.35,130
+0.36,215
+0.37,151
+0.38,72
+0.39,90
+0.4,68
+0.41,59
+0.42,56
+0.43,52
+0.44,36
+0.45,49
+0.46,38
+0.47,35
+0.48,37
+0.49,20
+0.5,14
+0.51,22
+0.52,16
+0.53,13
+0.54,10
+0.55,11
+0.56,8
+0.57,6
+0.58,5
+0.59,6
+0.6,1
+0.61,5
+0.62,4
+0.63,4
+0.69,2
+0.7,1
+0.81,1
diff --git a/spectrum_elucidation/configs/diffnmr/spectrum_elucidation/vocab/nless15/H1_statistic/split_type_distribution.csv b/spectrum_elucidation/configs/diffnmr/spectrum_elucidation/vocab/nless15/H1_statistic/split_type_distribution.csv
new file mode 100644
index 00000000..84f4de14
--- /dev/null
+++ b/spectrum_elucidation/configs/diffnmr/spectrum_elucidation/vocab/nless15/H1_statistic/split_type_distribution.csv
@@ -0,0 +1,114 @@
+Type,Count
+t,84634
+dt,33917
+dd,91890
+ddt,19659
+ddd,39143
+ddddd,1153
+dqd,2709
+td,18728
+m,198374
+s,148561
+q,21232
+dtdd,967
+d,134963
+qt,3063
+tq,1807
+tdd,1564
+qd,3769
+dtq,408
+dhept,249
+dp,4632
+p,12094
+dtt,3043
+dtd,8248
+tt,14468
+tdt,834
+dddt,2558
+tdq,299
+qq,399
+dddd,13661
+dq,14363
+h,2959
+hept,2363
+dqt,308
+dpd,159
+ttd,730
+pt,710
+tdp,78
+ddq,3308
+ddddq,55
+ddtd,1224
+tddd,340
+dddtd,57
+pd,1775
+dh,730
+qdt,149
+dddq,360
+ddtq,79
+dtddd,34
+hd,272
+tp,744
+ddqt,51
+ddtdd,57
+pq,177
+qtd,80
+qp,41
+ttdd,42
+dtdq,49
+ddtt,286
+th,158
+dtdt,277
+ttq,130
+qdd,406
+ddp,425
+dddddd,10
+ht,86
+dddp,76
+dttd,52
+ddddt,185
+ddqd,63
+ttt,138
+dtp,83
+dqq,46
+tdtd,18
+ddh,48
+qdq,39
+dttt,25
+heptd,120
+ddtdt,18
+tqd,54
+dddqd,2
+tdddd,19
+dddtt,28
+dqdt,9
+tpd,9
+tddt,79
+dqdd,79
+qddt,14
+dtqd,5
+tddq,16
+qddd,33
+pdd,59
+dpdd,10
+tddtd,2
+dtddt,10
+pdt,9
+tqt,28
+ttdt,23
+qtt,7
+tdtt,6
+tttd,2
+hdd,24
+dhd,2
+dddddt,13
+ddttd,1
+dpt,7
+dqddd,1
+ptd,4
+tdddt,2
+qqd,1
+ddqdd,1
+ddddtd,2
+dtdtd,2
+tqdd,1
diff --git a/spectrum_elucidation/sample.py b/spectrum_elucidation/sample.py
index 88cb6f93..11898b1b 100644
--- a/spectrum_elucidation/sample.py
+++ b/spectrum_elucidation/sample.py
@@ -14,60 +14,49 @@
 
 import argparse
 
-from ppmat.sampler.base_sampler import MolecularSampler
+from ppmat.sampler import MolecularSampler
 from ppmat.utils import logger
-
-if __name__ == "__main__":
-
-    argparse = argparse.ArgumentParser()
-
-    argparse.add_argument("--model_name", type=str, default=None)
-    argparse.add_argument(
-        "--weights_name",
-        type=str,
-        default=None,
-        help="Weights name, e.g., best.pdparams, latest.pdparams.",
-    )
-    argparse.add_argument(
-        "--config_path",
-        type=str,
-        default=None,
-        help="Path to the configuration file.",
+from ppmat.utils.inference_cli import add_model_loading_arguments
+from ppmat.utils.inference_cli import validate_config_overrides
+from ppmat.utils.inference_cli import validate_model_loading_arguments
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(description="Molecular structure sampling")
+    add_model_loading_arguments(parser)
+    parser.add_argument(
+        "--output_dir",
+        default="results",
+        help="Directory in which sampling results are saved.",
     )
-    argparse.add_argument(
-        "--checkpoint_path",
-        type=str,
-        default=None,
-        help="Path to the checkpoint file.",
-    )
-    argparse.add_argument("--save_path", type=str, default="results")
-    argparse.add_argument(
+    parser.add_argument(
         "--mode",
-        type=str,
-        choices=[
-            "by_dataloader",
-            "compute_metric",
-        ],
+        choices=["by_dataloader", "compute_metric"],
         default="by_dataloader",
     )
+    return parser
+
 
-    args = argparse.parse_args()
+def main():
+    parser = build_parser()
+    args, config_overrides = parser.parse_known_args()
+    validate_model_loading_arguments(parser, args)
+    validate_config_overrides(parser, config_overrides)
 
     sampler = MolecularSampler(
         model_name=args.model_name,
         weights_name=args.weights_name,
         config_path=args.config_path,
         checkpoint_path=args.checkpoint_path,
+        config_overrides=config_overrides,
     )
     if args.mode == "compute_metric":
-        metric_result = sampler.compute_metric(
-            save_path=args.save_path,
-        )
+        metric_result = sampler.compute_metric(save_path=args.output_dir)
         for metric_name, metric_value in metric_result.items():
             logger.info(f"{metric_name}: {metric_value}")
-    elif args.mode == "by_dataloader":
-        result = sampler.sample_by_dataloader(
-            save_path=args.save_path,
-        )
     else:
-        raise ValueError(f"Unknown mode: {args.mode}")
+        sampler.sample_by_dataloader(save_path=args.output_dir)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/spectrum_elucidation/train.py b/spectrum_elucidation/train.py
index 8926e118..5b934228 100644
--- a/spectrum_elucidation/train.py
+++ b/spectrum_elucidation/train.py
@@ -32,6 +32,7 @@
 from ppmat.trainer.base_trainer import BaseTrainer
 from ppmat.utils import logger
 from ppmat.utils import misc
+from ppmat.utils.io import append_timestamp_to_output_dir
 from ppmat.utils.visualization import MolecularVisualization
 
 if dist.get_world_size() > 1:
@@ -55,6 +56,9 @@
     cli_config = OmegaConf.from_dotlist(dynamic_args)
     config = OmegaConf.merge(config, cli_config)
 
+    seed = config["Trainer"].get("seed", 42)
+    append_timestamp_to_output_dir(config)
+
     # save config to output_dir, only rank 0 process will do this
     if dist.get_rank() == 0:
         os.makedirs(config["Trainer"]["output_dir"], exist_ok=True)
@@ -69,22 +73,32 @@
     logger.info(f"Logger saved to {logger_path}")
 
     # set random seed
-    seed = config["Trainer"].get("seed", 42)
     misc.set_random_seed(seed)
     logger.info(f"Set random seed to {seed}")
 
     # load dataloader from config
     set_signal_handlers()
-    if config["Global"].get("do_train", True):
+    do_train = config["Global"].get("do_train", True)
+    do_eval = config["Global"].get("do_eval", False)
+    do_test = config["Global"].get("do_test", False)
+    if not any((do_train, do_eval, do_test)):
+        raise ValueError(
+            "At least one of Global.do_train, Global.do_eval, or Global.do_test "
+            "must be True."
+        )
+
+    # DiffNMR dataset statistics and novelty metrics use the training SMILES even
+    # during evaluation-only and test-only runs.
+    if do_train or do_eval or do_test:
         train_data_cfg = config["Dataset"].get("train")
         assert (
             train_data_cfg is not None
-        ), "train_data_cfg must be defined, when do_train is true"
+        ), "train_data_cfg must be defined for DiffNMR train/eval/test."
         train_loader = build_dataloader(train_data_cfg)
     else:
         train_loader = None
 
-    if config["Global"].get("do_eval", False) or config["Global"].get("do_train", True):
+    if do_eval or do_train:
         val_data_cfg = config["Dataset"].get("val")
         if val_data_cfg is not None:
             val_loader = build_dataloader(val_data_cfg)
@@ -94,7 +108,7 @@
     else:
         val_loader = None
 
-    if config["Global"].get("do_test", False):
+    if do_test:
         test_data_cfg = config["Dataset"].get("test")
         assert (
             test_data_cfg is not None
diff --git a/spectrum_enhancement/train.py b/spectrum_enhancement/train.py
index 8c5436c5..794aac09 100644
--- a/spectrum_enhancement/train.py
+++ b/spectrum_enhancement/train.py
@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import argparse
-import datetime
 import os
 import os.path as osp
 from typing import Any
@@ -33,14 +32,15 @@
 from ppmat.utils import logger
 from ppmat.utils import misc
 from ppmat.utils.eager_comp_setting import setting_eager_mode
+from ppmat.utils.io import append_timestamp_to_output_dir
 
 
 def read_independent_dataloader_config(config: Dict[str, Any]):
     if config["Global"].get("do_train", True):
         train_data_cfg = config["Dataset"].get("train")
-        assert train_data_cfg is not None, (
-            "train_data_cfg must be defined when Global.do_train is True"
-        )
+        assert (
+            train_data_cfg is not None
+        ), "train_data_cfg must be defined when Global.do_train is True"
         train_loader = build_dataloader(train_data_cfg)
     else:
         train_loader = None
@@ -57,9 +57,9 @@ def read_independent_dataloader_config(config: Dict[str, Any]):
 
     if config["Global"].get("do_test", False):
         test_data_cfg = config["Dataset"].get("test")
-        assert test_data_cfg is not None, (
-            "test_data_cfg must be defined when Global.do_test is True"
-        )
+        assert (
+            test_data_cfg is not None
+        ), "test_data_cfg must be defined when Global.do_test is True"
         test_loader = build_dataloader(test_data_cfg)
     else:
         test_loader = None
@@ -81,11 +81,6 @@ def parse_args():
         default="./spectrum_enhancement/configs/sfin/sfin_haadf_enhance.yaml",
         help="Path to config file.",
     )
-    parser.add_argument(
-        "--append_timestamp",
-        action="store_true",
-        help="Append timestamp to Trainer.output_dir.",
-    )
     return parser.parse_known_args()
 
 
@@ -99,11 +94,7 @@ def main():
     cli_cfg = OmegaConf.from_dotlist(dynamic_args)
     cfg = OmegaConf.merge(cfg, cli_cfg)
 
-    if args.append_timestamp or cfg["Trainer"].get("append_timestamp", False):
-        seed = cfg["Trainer"].get("seed", 42)
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        base_output_dir = cfg["Trainer"]["output_dir"]
-        cfg["Trainer"]["output_dir"] = f"{base_output_dir}_t_{timestamp}_s_{seed}"
+    append_timestamp_to_output_dir(cfg)
 
     if dist.get_rank() == 0:
         os.makedirs(cfg["Trainer"]["output_dir"], exist_ok=True)
@@ -140,12 +131,12 @@ def main():
     model = build_model(model_cfg)
 
     if config.get("Optimizer") is not None and config["Global"].get("do_train", True):
-        assert train_loader is not None, (
-            "train_loader must be defined when Optimizer is provided."
-        )
-        assert config["Trainer"].get("max_epochs") is not None, (
-            "Trainer.max_epochs must be defined when Optimizer is provided."
-        )
+        assert (
+            train_loader is not None
+        ), "train_loader must be defined when Optimizer is provided."
+        assert (
+            config["Trainer"].get("max_epochs") is not None
+        ), "Trainer.max_epochs must be defined when Optimizer is provided."
         optimizer, lr_scheduler = build_optimizer(
             config["Optimizer"],
             model,
diff --git a/structure_generation/configs/diffcsp/README.md b/structure_generation/configs/diffcsp/README.md
index 790f0519..108d359f 100644
--- a/structure_generation/configs/diffcsp/README.md
+++ b/structure_generation/configs/diffcsp/README.md
@@ -198,13 +198,13 @@ python structure_generation/train.py -c structure_generation/configs/diffcsp/dif
 # Predict crystal structures using a trained model.
 # Mode 1: Use a pre-trained model (downloads automatically).
 # Mode 2: Use a custom configuration file and checkpoint.
-# Results are saved to the folder specified by --save_path (default: result).
+# Results are saved to the folder specified by --output_dir (default: results).
 
 # Mode 1: pre-trained model
-python structure_generation/sample.py --model_name='diffcsp_mp20' --weights_name='latest.pdparams' --save_path='result_diffcsp_mp20/' --chemical_formula='LiMnO2'
+python structure_generation/sample.py --model_name='diffcsp_mp20' --weights_name='latest.pdparams' --output_dir='result_diffcsp_mp20/' --chemical_formula='LiMnO2'
 
 # Mode 2: custom config + checkpoint
-python structure_generation/sample.py --config_path='structure_generation/configs/diffcsp/diffcsp_mp20.yaml' --checkpoint_path='./output/diffcsp_mp20/checkpoints/latest.pdparams' --save_path='result_diffcsp_mp20/' --chemical_formula='LiMnO2'
+python structure_generation/sample.py --config_path='structure_generation/configs/diffcsp/diffcsp_mp20.yaml' --checkpoint_path='./output/diffcsp_mp20/checkpoints/latest.pdparams' --output_dir='result_diffcsp_mp20/' --chemical_formula='LiMnO2'
 ```
 
 ---
diff --git a/structure_generation/configs/mattergen/README.md b/structure_generation/configs/mattergen/README.md
index d779119b..565a7fc3 100644
--- a/structure_generation/configs/mattergen/README.md
+++ b/structure_generation/configs/mattergen/README.md
@@ -291,61 +291,61 @@ python structure_generation/train.py -c structure_generation/configs/mattergen/m
 # This command is used to predict the crystal structure using a trained model.
 # Mode 1: Use a pre-trained model (downloads automatically).
 # Mode 2: Use a custom configuration file and checkpoint.
-# Results are saved to the folder specified by --save_path (default: result).
+# Results are saved to the folder specified by --output_dir (default: results).
 
 # mp20 dataset, without conditional constraints
-python structure_generation/sample.py --model_name='mattergen_mp20' --weights_name='latest.pdparams' --save_path='result_mattergen_mp20/' --mode='by_num_atoms' --num_atoms=4
-python structure_generation/sample.py --model_name='mattergen_mp20' --weights_name='latest.pdparams' --save_path='result_mattergen_mp20/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20.yaml' --checkpoint_path='./output/mattergen_mp20/checkpoints/latest.pdparams' --save_path='result_mattergen_mp20/' --mode='by_num_atoms' --num_atoms=4
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20.yaml' --checkpoint_path='./output/mattergen_mp20/checkpoints/latest.pdparams' --save_path='result_mattergen_mp20/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_mp20' --weights_name='latest.pdparams' --output_dir='result_mattergen_mp20/' --mode='by_num_atoms' --num_atoms=4
+python structure_generation/sample.py --model_name='mattergen_mp20' --weights_name='latest.pdparams' --output_dir='result_mattergen_mp20/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20.yaml' --checkpoint_path='./output/mattergen_mp20/checkpoints/latest.pdparams' --output_dir='result_mattergen_mp20/' --mode='by_num_atoms' --num_atoms=4
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20.yaml' --checkpoint_path='./output/mattergen_mp20/checkpoints/latest.pdparams' --output_dir='result_mattergen_mp20/' --mode='by_dataloader'
 
 # mp20 dataset, with chemical system constraints
-python structure_generation/sample.py --model_name='mattergen_mp20_chemical_system' --weights_name='latest.pdparams' --save_path='result_mattergen_mp20_chemical_system/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_chemical_system.yaml' --checkpoint_path='./output/mattergen_mp20_chemical_system/checkpoints/latest.pdparams' --save_path='result_mattergen_mp20_chemical_system/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_mp20_chemical_system' --weights_name='latest.pdparams' --output_dir='result_mattergen_mp20_chemical_system/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_chemical_system.yaml' --checkpoint_path='./output/mattergen_mp20_chemical_system/checkpoints/latest.pdparams' --output_dir='result_mattergen_mp20_chemical_system/' --mode='by_dataloader'
 
 # mp20 dataset, with dft_band_gap constraints
-python structure_generation/sample.py --model_name='mattergen_mp20_dft_band_gap' --weights_name='latest.pdparams' --save_path='result_mattergen_mp20_dft_band_gap/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_dft_band_gap.yaml' --checkpoint_path='./output/mattergen_mp20_dft_band_gap/checkpoints/latest.pdparams' --save_path='result_mattergen_mp20_dft_band_gap/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_mp20_dft_band_gap' --weights_name='latest.pdparams' --output_dir='result_mattergen_mp20_dft_band_gap/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_dft_band_gap.yaml' --checkpoint_path='./output/mattergen_mp20_dft_band_gap/checkpoints/latest.pdparams' --output_dir='result_mattergen_mp20_dft_band_gap/' --mode='by_dataloader'
 
 # mp20 dataset, with dft_bulk_modulus constraints
-python structure_generation/sample.py --model_name='mattergen_mp20_dft_bulk_modulus' --weights_name='latest.pdparams' --save_path='result_mattergen_mp20_dft_bulk_modulus/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_dft_bulk_modulus.yaml' --checkpoint_path='./output/mattergen_mp20_dft_bulk_modulus/checkpoints/latest.pdparams' --save_path='result_mattergen_mp20_dft_bulk_modulus/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_mp20_dft_bulk_modulus' --weights_name='latest.pdparams' --output_dir='result_mattergen_mp20_dft_bulk_modulus/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_dft_bulk_modulus.yaml' --checkpoint_path='./output/mattergen_mp20_dft_bulk_modulus/checkpoints/latest.pdparams' --output_dir='result_mattergen_mp20_dft_bulk_modulus/' --mode='by_dataloader'
 
 # mp20 dataset, with dft_mag_density constraints
-python structure_generation/sample.py --model_name='mattergen_mp20_dft_mag_density' --weights_name='latest.pdparams' --save_path='result_mattergen_mp20_dft_mag_density/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_dft_mag_density.yaml' --checkpoint_path='./output/mattergen_mp20_dft_mag_density/checkpoints/latest.pdparams' --save_path='result_mattergen_mp20_dft_mag_density/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_mp20_dft_mag_density' --weights_name='latest.pdparams' --output_dir='result_mattergen_mp20_dft_mag_density/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_mp20_dft_mag_density.yaml' --checkpoint_path='./output/mattergen_mp20_dft_mag_density/checkpoints/latest.pdparams' --output_dir='result_mattergen_mp20_dft_mag_density/' --mode='by_dataloader'
 
 # alex_mp20 dataset, without conditional constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20.yaml' --checkpoint_path='./output/mattergen_alex_mp20/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20.yaml' --checkpoint_path='./output/mattergen_alex_mp20/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with dft_band_gap constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_dft_band_gap' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_dft_band_gap/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_dft_band_gap.yaml' --checkpoint_path='./output/mattergen_alex_mp20_dft_band_gap/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_dft_band_gap/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_dft_band_gap' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_dft_band_gap/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_dft_band_gap.yaml' --checkpoint_path='./output/mattergen_alex_mp20_dft_band_gap/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_dft_band_gap/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with chemical system constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_chemical_system' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_chemical_system/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_chemical_system.yaml' --checkpoint_path='./output/mattergen_alex_mp20_chemical_system/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_chemical_system/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_chemical_system' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_chemical_system/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_chemical_system.yaml' --checkpoint_path='./output/mattergen_alex_mp20_chemical_system/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_chemical_system/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with dft_mag_density constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_dft_mag_density' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_dft_mag_density/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_dft_mag_density.yaml' --checkpoint_path='./output/mattergen_alex_mp20_dft_mag_density/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_dft_mag_density/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_dft_mag_density' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_dft_mag_density/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_dft_mag_density.yaml' --checkpoint_path='./output/mattergen_alex_mp20_dft_mag_density/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_dft_mag_density/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with ml_bulk_modulus constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_ml_bulk_modulus' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_ml_bulk_modulus/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_ml_bulk_modulus.yaml' --checkpoint_path='./output/mattergen_alex_mp20_ml_bulk_modulus/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_ml_bulk_modulus/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_ml_bulk_modulus' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_ml_bulk_modulus/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_ml_bulk_modulus.yaml' --checkpoint_path='./output/mattergen_alex_mp20_ml_bulk_modulus/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_ml_bulk_modulus/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with space_group constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_space_group' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_space_group/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_space_group.yaml' --checkpoint_path='./output/mattergen_alex_mp20_space_group/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_space_group/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_space_group' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_space_group/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_space_group.yaml' --checkpoint_path='./output/mattergen_alex_mp20_space_group/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_space_group/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with chemical system and energy above hull constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_chemical_system_energy_above_hull' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_chemical_system_energy_above_hull/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_chemical_system_energy_above_hull.yaml' --checkpoint_path='./output/mattergen_alex_mp20_chemical_system_energy_above_hull/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_chemical_system_energy_above_hull/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_chemical_system_energy_above_hull' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_chemical_system_energy_above_hull/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_chemical_system_energy_above_hull.yaml' --checkpoint_path='./output/mattergen_alex_mp20_chemical_system_energy_above_hull/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_chemical_system_energy_above_hull/' --mode='by_dataloader'
 
 # alex_mp20 dataset, with dft_mag_density and hhi_score constraints
-python structure_generation/sample.py --model_name='mattergen_alex_mp20_dft_mag_density_hhi_score' --weights_name='latest.pdparams' --save_path='result_mattergen_alex_mp20_dft_mag_density_hhi_score/' --mode='by_dataloader'
-python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_dft_mag_density_hhi_score.yaml' --checkpoint_path='./output/mattergen_alex_mp20_dft_mag_density_hhi_score/checkpoints/latest.pdparams' --save_path='result_mattergen_alex_mp20_dft_mag_density_hhi_score/' --mode='by_dataloader'
+python structure_generation/sample.py --model_name='mattergen_alex_mp20_dft_mag_density_hhi_score' --weights_name='latest.pdparams' --output_dir='result_mattergen_alex_mp20_dft_mag_density_hhi_score/' --mode='by_dataloader'
+python structure_generation/sample.py --config_path='structure_generation/configs/mattergen/mattergen_alex_mp20_dft_mag_density_hhi_score.yaml' --checkpoint_path='./output/mattergen_alex_mp20_dft_mag_density_hhi_score/checkpoints/latest.pdparams' --output_dir='result_mattergen_alex_mp20_dft_mag_density_hhi_score/' --mode='by_dataloader'
 ```
 
 ---
diff --git a/structure_generation/sample.py b/structure_generation/sample.py
index 51441fc2..e5e87b59 100644
--- a/structure_generation/sample.py
+++ b/structure_generation/sample.py
@@ -1,11 +1,11 @@
 # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-
+#
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
-
+#
 #     http://www.apache.org/licenses/LICENSE-2.0
-
+#
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -13,295 +13,88 @@
 # limitations under the License.
 
 import argparse
-import os
-from typing import Optional
 
-import numpy as np
-import paddle
 from omegaconf import OmegaConf
-from pymatgen.core import Composition
-from pymatgen.io.cif import CifWriter
 
-from ppmat.datasets import build_dataloader
-from ppmat.datasets.build_structure import BuildStructure
-from ppmat.datasets.transform import build_post_transforms
-from ppmat.metrics import build_metric
-from ppmat.models import build_model
-from ppmat.models import build_model_from_name
+from ppmat.sampler import StructureSampler
 from ppmat.utils import logger
-from ppmat.utils import save_load
-
-
-class StructureSampler:
-    """Structure Sampler.
-
-    This class provides an interface for sampling structures using pre-trained deep
-    learning models. Supports two initialization modes:
-
-    1. **Automatic Model Loading**
-       Specify `model_name` and `weights_name` to automatically download
-       and load pre-trained weights from the `MODEL_REGISTRY`.
-
-    2. **Custom Model Loading**
-       Provide explicit `config_path` and `checkpoint_path` to load
-       custom-trained models from local files.
-
-    Args:
-        model_name (Optional[str], optional): Name of the pre-defined model architecture
-            from the `MODEL_REGISTRY` registry. When specified, associated weights
-            will be automatically downloaded. Defaults to None.
-
-        weights_name (Optional[str], optional): Specific pre-trained weight identifier.
-            Used only when `model_name` is provided. Valid options include:
-            - 'best.pdparams' (highest validation performance)
-            - 'latest.pdparams' (most recent training checkpoint)
-            - Custom weight files ending with '.pdparams'
-            Defaults to None.
-
-        config_path (Optional[str], optional): Path to model configuration file (YAML)
-            for custom models. Required when not using predefined `model_name`.
-            Defaults to None.
-        checkpoint_path (Optional[str], optional): Path to model checkpoint file
-            (.pdparams) for custom models. Required when not using predefined
-            `model_name`. Defaults to None.
-    """
-
-    def __init__(
-        self,
-        model_name: Optional[str] = None,
-        weights_name: Optional[str] = None,
-        config_path: Optional[str] = None,
-        checkpoint_path: Optional[str] = None,
-    ):
-        # if model_name is not None, then config_path and checkpoint_path must be
-        # provided
-        if model_name is None:
-            assert (
-                config_path is not None and checkpoint_path is not None
-            ), "config_path and checkpoint_path must be provided when model_name is "
-            "None."
-
-            logger.info(f"Loading model from {config_path} and {checkpoint_path}.")
-
-            config = OmegaConf.load(config_path)
-            config = OmegaConf.to_container(config, resolve=True)
-
-            model_config = config.get("Model", None)
-            assert model_config is not None, "Model config must be provided."
-            model = build_model(model_config)
-            save_load.load_pretrain(model, checkpoint_path)
-
-        else:
-            logger.info("Since model_name is given, downloading it...")
-            model, config = build_model_from_name(model_name, weights_name)
-
-        self.model = model
-        self.config = config
-
-        self.model.eval()
-
-        # sample config
-        sample_config = config.get("Sample", None)
-        self.sample_config = sample_config
-
-        self.post_transforms_cfg = self.sample_config.get("post_transforms", None)
-        if self.post_transforms_cfg is not None:
-            self.post_transforms = build_post_transforms(self.post_transforms_cfg)
-        else:
-            self.post_transforms = None
-
-    def compute_metric(
-        self,
-        save_path=None,
-    ):
-        metrics_cfg = self.sample_config.get("metrics")
-        assert metrics_cfg is not None, "metrics config must be provided."
-        metrics_fn = build_metric(metrics_cfg)
-
-        total_results = self.sample_by_dataloader(save_path)
-
-        metric = metrics_fn(total_results)
-        return metric
-
-    def post_process(self, data):
-        if self.post_transforms is None:
-            return data
-        return self.post_transforms(data)
+from ppmat.utils.inference_cli import add_model_loading_arguments
+from ppmat.utils.inference_cli import validate_config_overrides
+from ppmat.utils.inference_cli import validate_model_loading_arguments
 
-    def sample(self, data, sample_params=None):
-        if sample_params is None:
-            sample_params = {}
-        assert isinstance(sample_params, dict), "sample_params must be a dict or None."
-        pred_data = self.model.sample(data, **sample_params)
-        pred_data = self.post_process(pred_data)
-        return pred_data
+__all__ = ["StructureSampler", "build_parser", "main"]
 
-    def sample_by_dataloader(
-        self,
-        save_path=None,
-    ):
-        dataset_cfg = self.sample_config["data"]
-        data_loader = build_dataloader(dataset_cfg)
 
-        build_structure_cfg = self.sample_config["build_structure_cfg"]
-        structure_converter = BuildStructure(**build_structure_cfg)
-
-        logger.info(f"Total iterations: {len(data_loader)}")
-        logger.info("Start sampling process...\n")
-
-        total_results = []
-        for iter_id, batch_data in enumerate(data_loader):
-            pred_data = self.model.sample(batch_data)
-            structures = structure_converter(pred_data["result"])
-            if save_path is not None:
-                os.makedirs(save_path, exist_ok=True)
-                for i, structure in enumerate(structures):
-                    formula = structure.formula.replace(" ", "-")
-                    tar_file = os.path.join(
-                        save_path, f"{formula}_{iter_id + 1}_{i + 1}.cif"
-                    )
-                    if structure is not None:
-                        writer = CifWriter(structure)
-                        writer.write_file(tar_file)
-                    else:
-                        logger.info(
-                            f"No structure generated for iteration {iter_id}, index {i}"
-                        )
-            total_results.extend(pred_data["result"])
-        return total_results
-
-    def sample_by_num_atoms(self, num_atoms, save_path=None, sample_params=None):
-        assert isinstance(num_atoms, int), "num_atoms must be an integer."
-        data = {
-            "structure_array": {
-                "num_atoms": paddle.to_tensor(np.array([num_atoms]).astype("int64")),
-            }
-        }
-
-        result = self.sample(data, sample_params=sample_params)
-
-        if save_path is not None:
-            os.makedirs(save_path, exist_ok=True)
-            logger.info(f"Save results to {save_path}")
-            build_structure_cfg = self.sample_config["build_structure_cfg"]
-            structure_converter = BuildStructure(**build_structure_cfg)
-            structures = structure_converter(result["result"])
-            for i, structure in enumerate(structures):
-                formula = structure.formula.replace(" ", "-")
-                tar_file = os.path.join(save_path, f"{formula}_{i + 1}.cif")
-                if structure is not None:
-                    writer = CifWriter(structure)
-                    writer.write_file(tar_file)
-                else:
-                    logger.info(f"No structure generated for index {i}")
-
-        return result
-
-    def sample_by_chemical_formula(
-        self, chemical_formula, save_path=None, sample_params=None
-    ):
-        assert isinstance(chemical_formula, str), "chemical_formula must be a string."
-        composition = Composition(chemical_formula)
-        atom_types = []
-        for elem, num in composition.items():
-            atom_types.extend([elem.Z] * int(num))
-        atom_types = np.array(atom_types).astype("int64")
-
-        data = {
-            "structure_array": {
-                "atom_types": paddle.to_tensor(atom_types),
-                "num_atoms": paddle.to_tensor(
-                    np.array([atom_types.shape[0]]).astype("int64")
-                ),
-            }
-        }
-        result = self.sample(data, sample_params=sample_params)
-
-        if save_path is not None:
-            os.makedirs(save_path, exist_ok=True)
-            logger.info(f"Save results to {save_path}")
-            build_structure_cfg = self.sample_config["build_structure_cfg"]
-            structure_converter = BuildStructure(**build_structure_cfg)
-            structures = structure_converter(result["result"])
-            for i, structure in enumerate(structures):
-                formula = structure.formula.replace(" ", "-")
-                tar_file = os.path.join(save_path, f"{formula}_{i + 1}.cif")
-                if structure is not None:
-                    writer = CifWriter(structure)
-                    writer.write_file(tar_file)
-                else:
-                    logger.info(f"No structure generated for index {i}")
-
-        return result
-
-    def sample_by_condition(self, composition, save_path=None, sample_params=None):
-        # todo: implement this function
-        pass
-
-
-if __name__ == "__main__":
-
-    argparse = argparse.ArgumentParser()
-
-    argparse.add_argument("--model_name", type=str, default=None)
-    argparse.add_argument(
-        "--weights_name",
-        type=str,
-        default=None,
-        help="Weights name, e.g., best.pdparams, latest.pdparams.",
-    )
-    argparse.add_argument(
-        "--config_path",
-        type=str,
-        default=None,
-        help="Path to the configuration file.",
+def build_parser():
+    parser = argparse.ArgumentParser(description="Crystal structure sampling")
+    add_model_loading_arguments(parser)
+    parser.add_argument(
+        "--output_dir",
+        default="results",
+        help="Directory in which generated structures are saved.",
     )
-    argparse.add_argument(
-        "--checkpoint_path",
-        type=str,
-        default=None,
-        help="Path to the checkpoint file.",
+    parser.add_argument("--chemical_formula", default="LiMnO2")
+    parser.add_argument("--num_atoms", type=int, default=4)
+    parser.add_argument(
+        "--condition",
+        action="append",
+        default=[],
+        metavar="NAME=VALUE",
+        help="Conditional property value; repeat for multi-property models.",
     )
-    argparse.add_argument("--save_path", type=str, default="results")
-    argparse.add_argument("--chemical_formula", type=str, default="LiMnO2")
-    argparse.add_argument("--num_atoms", type=int, default=4)
-    argparse.add_argument(
+    parser.add_argument(
         "--mode",
-        type=str,
         choices=[
             "by_chemical_formula",
             "by_num_atoms",
+            "by_condition",
             "by_dataloader",
             "compute_metric",
         ],
         default="by_chemical_formula",
     )
+    return parser
+
 
-    args = argparse.parse_args()
+def main():
+    parser = build_parser()
+    args, config_overrides = parser.parse_known_args()
+    validate_model_loading_arguments(parser, args)
+    validate_config_overrides(parser, config_overrides)
 
     sampler = StructureSampler(
         model_name=args.model_name,
         weights_name=args.weights_name,
         config_path=args.config_path,
         checkpoint_path=args.checkpoint_path,
+        config_overrides=config_overrides,
     )
     if args.mode == "compute_metric":
-        metric_result = sampler.compute_metric(save_path=args.save_path)
+        metric_result = sampler.compute_metric(save_path=args.output_dir)
         for metric_name, metric_value in metric_result.items():
             logger.info(f"{metric_name}: {metric_value}")
     elif args.mode == "by_chemical_formula":
-        result = sampler.sample_by_chemical_formula(
+        sampler.sample_by_chemical_formula(
             chemical_formula=args.chemical_formula,
-            save_path=args.save_path,
+            save_path=args.output_dir,
         )
     elif args.mode == "by_num_atoms":
-        result = sampler.sample_by_num_atoms(
+        sampler.sample_by_num_atoms(
             num_atoms=args.num_atoms,
-            save_path=args.save_path,
+            save_path=args.output_dir,
         )
-    elif args.mode == "by_dataloader":
-        result = sampler.sample_by_dataloader(
-            save_path=args.save_path,
+    elif args.mode == "by_condition":
+        conditions = OmegaConf.to_container(
+            OmegaConf.from_dotlist(args.condition), resolve=True
+        )
+        sampler.sample_by_condition(
+            num_atoms=args.num_atoms,
+            conditions=conditions,
+            save_path=args.output_dir,
         )
     else:
-        raise ValueError(f"Unknown mode: {args.mode}")
+        sampler.sample_by_dataloader(save_path=args.output_dir)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/structure_generation/train.py b/structure_generation/train.py
index 09a7f94a..1f564644 100644
--- a/structure_generation/train.py
+++ b/structure_generation/train.py
@@ -28,6 +28,7 @@
 from ppmat.trainer.base_trainer import BaseTrainer
 from ppmat.utils import logger
 from ppmat.utils import misc
+from ppmat.utils.io import append_timestamp_to_output_dir
 
 if dist.get_world_size() > 1:
     fleet.init(is_collective=True)
@@ -49,6 +50,9 @@
     cli_config = OmegaConf.from_dotlist(dynamic_args)
     config = OmegaConf.merge(config, cli_config)
 
+    seed = config["Trainer"].get("seed", 42)
+    append_timestamp_to_output_dir(config)
+
     # save config to output_dir, only rank 0 process will do this
     if dist.get_rank() == 0:
         os.makedirs(config["Trainer"]["output_dir"], exist_ok=True)
@@ -63,7 +67,6 @@
     logger.info(f"Logger saved to {logger_path}")
 
     # set random seed
-    seed = config["Trainer"].get("seed", 42)
     misc.set_random_seed(seed)
     logger.info(f"Set random seed to {seed}")
 
diff --git a/test/test_all_registered_models_integration.py b/test/test_all_registered_models_integration.py
new file mode 100644
index 00000000..4fbb8046
--- /dev/null
+++ b/test/test_all_registered_models_integration.py
@@ -0,0 +1,75 @@
+"""Opt-in integration tests for every package in ``MODEL_REGISTRY``.
+
+Run with:
+
+    RUN_ALL_MODEL_PACKAGE_INTEGRATION=1 \
+      pytest test/test_all_registered_models_integration.py -q
+"""
+
+import gc
+import os
+import re
+from pathlib import Path
+
+import paddle
+import pytest
+
+from ppmat.models import MODEL_REGISTRY
+from ppmat.utils import download
+from ppmat.utils.model_package import resolve_model_package_dir
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("RUN_ALL_MODEL_PACKAGE_INTEGRATION") != "1",
+    reason="set RUN_ALL_MODEL_PACKAGE_INTEGRATION=1 to test every model package",
+)
+
+
+@pytest.mark.parametrize("model_name", sorted(MODEL_REGISTRY))
+def test_all_registered_model_packages_load(model_name):
+    """Download the published package and load its configured checkpoint."""
+    if model_name == "diffnmr_msdnmr_nless15":
+        from ppmat.sampler import MolecularSampler
+
+        model_or_sampler = MolecularSampler(
+            model_name=model_name,
+            weights_name="best.pdparams",
+        )
+        model = model_or_sampler.model
+    else:
+        from ppmat.models import build_model_from_name
+
+        model_or_sampler, config = build_model_from_name(model_name)
+        assert config.get("Model") is not None
+        model = model_or_sampler
+
+    extracted_path = download.get_weights_path_from_url(MODEL_REGISTRY[model_name])
+    package_dir = Path(resolve_model_package_dir(model_name, extracted_path))
+    checkpoint_path = _select_checkpoint(package_dir)
+    checkpoint = paddle.load(str(checkpoint_path))
+    if "state_dict" in checkpoint:
+        checkpoint = checkpoint["state_dict"]
+    assert set(checkpoint) == set(model.state_dict()), (
+        f"{model_name} checkpoint keys differ: "
+        f"missing={sorted(set(model.state_dict()) - set(checkpoint))}, "
+        f"unexpected={sorted(set(checkpoint) - set(model.state_dict()))}"
+    )
+
+    del model_or_sampler
+    gc.collect()
+
+
+def _select_checkpoint(package_dir: Path) -> Path:
+    checkpoints = list(package_dir.rglob("*.pdparams"))
+    assert checkpoints, f"No checkpoint found under {package_dir}"
+    for preferred_name in ("best.pdparams", "latest.pdparams"):
+        preferred = [path for path in checkpoints if path.name == preferred_name]
+        if preferred:
+            return sorted(preferred)[0]
+    epochs = [
+        (int(match.group(1)), path)
+        for path in checkpoints
+        if (match := re.fullmatch(r"epoch_(\d+)\.pdparams", path.name))
+    ]
+    if epochs:
+        return max(epochs)[1]
+    return sorted(checkpoints)[0]
diff --git a/test/test_inference_cli.py b/test/test_inference_cli.py
new file mode 100644
index 00000000..572fbf90
--- /dev/null
+++ b/test/test_inference_cli.py
@@ -0,0 +1,167 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from pathlib import Path
+
+import pytest
+
+from ppmat.utils.inference_cli import add_model_loading_arguments
+from ppmat.utils.inference_cli import validate_config_overrides
+from ppmat.utils.inference_cli import validate_model_loading_arguments
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _parse(arguments):
+    parser = argparse.ArgumentParser()
+    add_model_loading_arguments(parser)
+    args, overrides = parser.parse_known_args(arguments)
+    validate_model_loading_arguments(parser, args)
+    validate_config_overrides(parser, overrides)
+    return args, overrides
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        ["--model_name", "registered_model"],
+        ["--config_path", "model.yaml", "--checkpoint_path", "model.pdparams"],
+    ],
+)
+def test_model_loading_modes(arguments):
+    args, overrides = _parse(arguments)
+    assert args.model_name or (args.config_path and args.checkpoint_path)
+    assert overrides == []
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        [],
+        ["--config_path", "model.yaml"],
+        ["--checkpoint_path", "model.pdparams"],
+        [
+            "--model_name",
+            "registered_model",
+            "--config_path",
+            "model.yaml",
+            "--checkpoint_path",
+            "model.pdparams",
+        ],
+        ["--model_name", "registered_model", "--config", "model.yaml"],
+    ],
+)
+def test_invalid_model_loading_modes(arguments):
+    with pytest.raises(SystemExit):
+        _parse(arguments)
+
+
+def test_config_overrides():
+    _, overrides = _parse(
+        ["--model_name", "registered_model", "Predict.grid_batch_size=128"]
+    )
+    assert overrides == ["Predict.grid_batch_size=128"]
+
+
+def test_structure_sampler_uses_common_cli_contract():
+    from structure_generation.sample import build_parser
+
+    parser = build_parser()
+    args, overrides = parser.parse_known_args(
+        [
+            "--model_name",
+            "mattergen_mp20",
+            "--output_dir",
+            "generated",
+            "Sample.foo=bar",
+        ]
+    )
+    validate_model_loading_arguments(parser, args)
+    validate_config_overrides(parser, overrides)
+
+    assert args.output_dir == "generated"
+    assert overrides == ["Sample.foo=bar"]
+
+
+def test_structure_sampler_is_exported_from_ppmat():
+    from ppmat.sampler import StructureSampler
+    from ppmat.sampler.structure_sampler import StructureSampler as Implementation
+
+    assert StructureSampler is Implementation
+
+
+def test_structure_sampler_rejects_num_atoms_for_composition_only_model():
+    from ppmat.sampler import StructureSampler
+
+    sampler = StructureSampler.__new__(StructureSampler)
+    sampler.model = type(
+        "CompositionOnlyModel", (), {"supports_num_atoms_sampling": False}
+    )()
+
+    with pytest.raises(NotImplementedError, match="sample_by_chemical_formula"):
+        sampler.sample_by_num_atoms(4)
+
+
+def test_structure_sampler_uses_configured_sample_parameters():
+    from ppmat.sampler import StructureSampler
+
+    class Model:
+        def sample(self, data, **kwargs):
+            return {"result": data, "sample_params": kwargs}
+
+    sampler = StructureSampler.__new__(StructureSampler)
+    sampler.model = Model()
+    sampler.sample_config = {"model_sample_params": {"num_inference_steps": 7}}
+    sampler.post_transforms = None
+
+    result = sampler.sample({"value": 1})
+
+    assert result["sample_params"] == {"num_inference_steps": 7}
+
+
+@pytest.mark.parametrize(
+    "document_name",
+    [
+        "README.md",
+        "README_PYPI.md",
+        "README_zh.md",
+        "README_ja.md",
+        "Install.md",
+        "Install_cn.md",
+    ],
+)
+def test_user_facing_docs_use_current_one_click_inference_contract(document_name):
+    document = (ROOT / document_name).read_text()
+
+    assert "--config=" not in document
+    assert "--checkpoint=" not in document
+    assert "--save_path='result_diffnmr" not in document
+    assert "--model_name='infgcn_qm9'" in document
+    assert "--mol_input='electronic_structure/configs/infgcn/example/methane.mol'" in (
+        document
+    )
+    assert "--model_name='diffnmr_msdnmr_nless15'" in document
+    assert "--model_name='sfin_haadf_enhance'" in document
+    assert "--input_path='path/to/noisy_image.png'" in document
+
+
+def test_homepage_language_badges_follow_project_badges():
+    readme = (ROOT / "README.md").read_text()
+
+    project_badges_end = readme.index("</p>", readme.index("PyPI version"))
+    language_badges_start = readme.index('<a href="README.md">', project_badges_end)
+    assert language_badges_start > project_badges_end
+    assert '<a href="README_zh.md">' in readme
+    assert '<a href="README_ja.md">' in readme
diff --git a/test/test_one_click_model_loading.py b/test/test_one_click_model_loading.py
new file mode 100644
index 00000000..54881602
--- /dev/null
+++ b/test/test_one_click_model_loading.py
@@ -0,0 +1,1160 @@
+from __future__ import annotations
+
+import ast
+import importlib.util
+import inspect
+import re
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+from omegaconf import OmegaConf
+
+ROOT = Path(__file__).resolve().parents[1]
+INFGCN_CONFIG_DIR = ROOT / "electronic_structure/configs/infgcn"
+INFGCN_MODEL_NAMES = [
+    "infgcn_md17_benzene",
+    "infgcn_md17_ethane",
+    "infgcn_md17_ethanol",
+    "infgcn_md17_malonaldehyde",
+    "infgcn_md17_phenol",
+    "infgcn_md17_resorcinol",
+    "infgcn_mp",
+    "infgcn_omol25_mc_5k_trimmed",
+    "infgcn_qm9",
+]
+
+
+def _models_module_ast():
+    return ast.parse((ROOT / "ppmat/models/__init__.py").read_text())
+
+
+def _literal_assign(name: str):
+    for node in _models_module_ast().body:
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and target.id == name:
+                    return ast.literal_eval(node.value)
+    raise AssertionError(f"{name} assignment not found")
+
+
+def test_infgcn_and_diffnmr_are_registered_for_one_click_loading():
+    registry = _literal_assign("MODEL_REGISTRY")
+
+    for model_name in [*INFGCN_MODEL_NAMES, "diffnmr_msdnmr_nless15"]:
+        assert model_name in registry
+        assert registry[model_name].startswith("https://paddle-org.bj.bcebos.com/")
+        assert registry[model_name].endswith(".zip")
+        assert registry[model_name].endswith(f"{model_name}.zip")
+
+
+def test_models_init_keeps_one_click_surface_minimal():
+    source = (ROOT / "ppmat/models/__init__.py").read_text()
+
+    forbidden_names = [
+        "MODEL_CONFIG_REGISTRY",
+        "MODEL_SUPPORT_REGISTRY",
+        "_repo_root",
+        "_resolve_repo_path",
+        "get_model_config_path_from_name",
+        "get_model_package_path_from_name",
+        "get_model_file_path_from_package",
+        "get_model_config_path_from_package",
+    ]
+    for name in forbidden_names:
+        assert name not in source
+
+
+def test_model_package_helpers_resolve_standard_zip_layout(tmp_path):
+    from ppmat.utils.model_package import get_model_config_path
+    from ppmat.utils.model_package import resolve_model_package_dir
+
+    cache_dir = tmp_path / "infgcn_qm9"
+    package_dir = cache_dir / "infgcn_qm9"
+    checkpoints_dir = package_dir / "checkpoints"
+    checkpoints_dir.mkdir(parents=True)
+    config_path = package_dir / "infgcn_qm9.yaml"
+    config_path.write_text("Model: {}\n")
+
+    resolved_dir = resolve_model_package_dir("infgcn_qm9", str(cache_dir))
+    assert Path(resolved_dir) == package_dir
+    assert Path(get_model_config_path("infgcn_qm9", resolved_dir)) == config_path
+
+
+def test_model_package_helpers_reject_nonstandard_recursive_layout(tmp_path):
+    import pytest
+
+    from ppmat.utils.model_package import resolve_model_package_dir
+
+    nested_dir = tmp_path / "download" / "arbitrary" / "nested"
+    nested_dir.mkdir(parents=True)
+    (nested_dir / "infgcn_qm9.yaml").write_text("Model: {}\n")
+
+    with pytest.raises(FileNotFoundError, match="Invalid package"):
+        resolve_model_package_dir("infgcn_qm9", str(tmp_path / "download"))
+
+
+def test_model_package_helpers_reject_unrelated_single_yaml(tmp_path):
+    import pytest
+
+    from ppmat.utils.model_package import get_model_config_path
+
+    package_dir = tmp_path / "infgcn_qm9"
+    package_dir.mkdir()
+    (package_dir / "unrelated.yaml").write_text("Model: {}\n")
+
+    with pytest.raises(FileNotFoundError, match="infgcn_qm9.yaml"):
+        get_model_config_path("infgcn_qm9", str(package_dir))
+
+
+def test_model_package_helpers_have_focused_module():
+    package_source = (ROOT / "ppmat/utils/model_package.py").read_text()
+    generic_io_source = (ROOT / "ppmat/utils/io.py").read_text()
+
+    assert "def resolve_model_package_dir" in package_source
+    assert "def get_model_config_path" in package_source
+    assert "model_package" not in generic_io_source
+    assert "find_config_file_in_package" not in generic_io_source
+
+
+def test_diffnmr_assets_do_not_expand_shared_download_api():
+    download_source = (ROOT / "ppmat/utils/download.py").read_text()
+    sampler_source = (ROOT / "ppmat/sampler/diffnmr.py").read_text()
+
+    assert "get_assets_path_from_url" not in download_source
+    assert "DIFFNMR_ASSETS_HOME" not in sampler_source
+    assert "_download_diffnmr_asset" not in sampler_source
+
+
+def test_build_model_from_name_requires_standard_package():
+    import ppmat.models as models
+
+    source = inspect.getsource(models.build_model_from_name)
+
+    assert "resolve_model_package_dir(model_name, extracted_path)" in source
+    assert "get_model_config_path(model_name, path)" in source
+    assert "os.walk" not in source
+
+
+def test_diffnmr_train_smiles_uses_existing_datadir_cache(tmp_path):
+    import numpy as np
+
+    from ppmat.datasets.msd_nmr_dataset import get_train_smiles
+
+    cache_dir = tmp_path / "msd_nmr_nless15_cache" / "train"
+    cache_dir.mkdir(parents=True)
+    smiles_path = cache_dir / "train_smiles_no_h.npy"
+    expected_smiles = np.array(["CCO", "CO"])
+    np.save(smiles_path, expected_smiles)
+
+    cfg = {
+        "datadir": str(tmp_path),
+        "data_flag": "n<15",
+        "build_graph_cfg": {"__init_params__": {"remove_h": True}},
+    }
+    dataset_infos = SimpleNamespace(atom_decoder=["C", "N", "O", "F"])
+
+    train_smiles = get_train_smiles(cfg, dataloader=[], dataset_infos=dataset_infos)
+
+    np.testing.assert_array_equal(train_smiles, expected_smiles)
+
+
+def test_diffnmr_dataset_infos_can_skip_train_smiles():
+    from ppmat.datasets.msd_nmr_dataset import MSDnmrinfos
+
+    cfg = {
+        "data_flag": "n<15",
+        "build_graph_cfg": {"__init_params__": {"remove_h": True}},
+        "load_train_smiles": False,
+    }
+
+    dataset_infos = MSDnmrinfos(
+        dataloaders=SimpleNamespace(train_dataloader=None),
+        cfg=cfg,
+    )
+
+    assert dataset_infos.train_smiles is None
+    assert dataset_infos.atom_decoder == ["C", "N", "O", "F", "P", "S", "Cl", "Br", "I"]
+
+
+def test_infgcn_predict_uses_config_defaults_for_cli_options():
+    source = (ROOT / "electronic_structure/predict.py").read_text()
+    field_source = (ROOT / "ppmat/predictor/field_predictor.py").read_text()
+    cfg = OmegaConf.to_container(
+        OmegaConf.load(INFGCN_CONFIG_DIR / "infgcn_qm9.yaml"),
+        resolve=True,
+    )
+
+    assert "FieldPredictor" in source
+    assert "def apply_predict_config" in field_source
+    assert cfg["Predict"]["grid_batch_size"] == 20000
+    assert cfg["Predict"]["output_dir"] == "output/infgcn_qm9/vis_val0"
+    assert cfg["Predict"]["save_pred_cube"] is True
+    assert cfg["Predict"]["save_true_cube"] is True
+    assert cfg["Predict"]["cube_dir"] == "output/infgcn_qm9/cubes"
+
+
+def test_infgcn_predict_config_fills_unset_cli_options():
+    from electronic_structure.predict import apply_predict_config
+
+    args = SimpleNamespace(
+        split=None,
+        index=None,
+        data_root=None,
+        split_file=None,
+        atom_file=None,
+        output_dir=None,
+        grid_batch_size=123,
+        skip_vis=None,
+        save_true_cube=None,
+        save_pred_cube=None,
+        save_html=None,
+        cube_dir=None,
+        show_plot=None,
+        mol_pattern=None,
+        mol_grid_shape=None,
+        mol_grid_padding=None,
+        mol_true_cube_dir=None,
+    )
+    cfg = {
+        "Predict": {
+            "split": "validation",
+            "index": 3,
+            "output_dir": "from_config",
+            "grid_batch_size": 456,
+            "save_pred_cube": True,
+        }
+    }
+
+    apply_predict_config(args, cfg)
+
+    assert args.split == "validation"
+    assert args.index == 3
+    assert args.output_dir == "from_config"
+    assert args.grid_batch_size == 123
+    assert args.save_pred_cube is True
+    assert args.save_true_cube is False
+
+
+def test_infgcn_predict_cli_accepts_one_click_model_arguments():
+    source = (ROOT / "electronic_structure/predict.py").read_text()
+    cli_source = (ROOT / "ppmat/utils/inference_cli.py").read_text()
+
+    assert "add_model_loading_arguments(parser)" in source
+    assert '"--model_name"' in cli_source
+    assert '"--weights_name"' in cli_source
+    assert "FieldPredictor(" in source
+
+
+def test_infgcn_mol_atom_mapping_requires_explicit_existing_file(tmp_path):
+    import pytest
+
+    from ppmat.predictor.field_predictor import resolve_atom_file_for_mol
+
+    with pytest.raises(FileNotFoundError, match="--atom_file"):
+        resolve_atom_file_for_mol(
+            str(tmp_path / "cli_atoms.json"),
+            str(tmp_path / "config_atoms.json"),
+        )
+
+
+def test_infgcn_mol_atom_mapping_prefers_cli_path(tmp_path):
+    from ppmat.predictor.field_predictor import resolve_atom_file_for_mol
+
+    cli_path = tmp_path / "cli_atoms.json"
+    config_path = tmp_path / "config_atoms.json"
+    cli_path.write_text("[]")
+    config_path.write_text("[]")
+
+    assert resolve_atom_file_for_mol(str(cli_path), str(config_path)) == cli_path
+
+
+def test_field_predictor_is_shared_predictor_entrypoint():
+    import ppmat.predictor as predictor
+
+    field_source = (ROOT / "ppmat/predictor/field_predictor.py").read_text()
+    entry_source = (ROOT / "electronic_structure/predict.py").read_text()
+
+    assert not (ROOT / "ppmat/predictor/field.py").exists()
+    assert hasattr(predictor, "FieldPredictor")
+    assert "class FieldPredictor" in field_source
+    assert "from ppmat.predictor import FieldPredictor" in entry_source
+    assert "from ppmat.predictor.field_predictor import apply_predict_config" in (
+        entry_source
+    )
+    assert "from ppmat.models import MODEL_REGISTRY" not in entry_source
+    assert "from ppmat.datasets import DensityDataset" not in entry_source
+
+
+def test_field_predictor_reuses_base_and_keeps_helpers_outside_predictor():
+    field_source = (ROOT / "ppmat/predictor/field_predictor.py").read_text()
+    io_source = (ROOT / "ppmat/utils/io.py").read_text()
+    field_io_source = (ROOT / "ppmat/predictor/field_io.py").read_text()
+    visualization_source = (ROOT / "ppmat/utils/visualization.py").read_text()
+
+    assert "from ppmat.predictor.base import BasePredictor" in field_source
+    assert "class FieldPredictor(BasePredictor):" in field_source
+    assert "self._load_model()" in field_source
+    assert "def _load_model(self):" in field_source
+    assert not (ROOT / "ppmat/utils/field_io.py").exists()
+    assert (ROOT / "ppmat/predictor/field_io.py").exists()
+    assert not (ROOT / "ppmat/utils/field_visualization.py").exists()
+
+    for helper_name in [
+        "draw_volume",
+        "safe_write_image",
+        "maybe_downsample_volume",
+        "read_cube_density",
+        "write_cube",
+        "prepare_cube_info",
+    ]:
+        assert f"def {helper_name}" not in field_source
+
+    for helper_name in ["read_cube_density", "write_cube", "prepare_cube_info"]:
+        assert f"def {helper_name}" in field_io_source
+        assert f"def {helper_name}" not in io_source
+
+    for helper_name in ["draw_volume", "safe_write_image", "maybe_downsample_volume"]:
+        assert f"def {helper_name}" in visualization_source
+
+    top_level_imports = "\n".join(
+        line
+        for line in visualization_source.splitlines()
+        if line.startswith("import ") or line.startswith("from ")
+    )
+    assert "import imageio" in top_level_imports
+    assert "import matplotlib.pyplot as plt" in top_level_imports
+    assert "import networkx as nx" in top_level_imports
+    assert "import plotly.graph_objects as go" in top_level_imports
+    assert "import rdkit" in top_level_imports
+    assert "def _rdkit_modules" not in visualization_source
+    assert "def _matplotlib_pyplot" not in visualization_source
+    assert "def _networkx" not in visualization_source
+    assert "def _imageio" not in visualization_source
+
+    assert "def _save_cubes" in field_source
+    assert "def _save_visualizations" in field_source
+    assert "FieldPredictor._save_cubes(" in field_source
+    assert "FieldPredictor._save_visualizations(" in field_source
+
+
+def test_field_cube_io_round_trip(tmp_path):
+    import numpy as np
+
+    from ppmat.predictor.field_io import read_cube_density
+    from ppmat.predictor.field_io import write_cube
+
+    cube_path = tmp_path / "density.cube"
+    density = np.arange(8, dtype=np.float32)
+    info = {
+        "shape": [2, 2, 2],
+        "cell": np.eye(3, dtype=np.float32) * 2,
+        "origin": np.asarray([0.5, 1.0, 1.5], dtype=np.float32),
+    }
+    with cube_path.open("w") as file_obj:
+        write_cube(
+            file_obj,
+            atom_type=np.asarray([6]),
+            atom_coord=np.asarray([[0.0, 0.0, 0.0]], dtype=np.float32),
+            density=density,
+            info=info,
+        )
+
+    loaded_density, grid_coord, loaded_info = read_cube_density(cube_path)
+
+    np.testing.assert_allclose(loaded_density.numpy(), density)
+    np.testing.assert_allclose(loaded_info["origin"].numpy(), info["origin"])
+    np.testing.assert_allclose(grid_coord.numpy()[0], info["origin"])
+    assert loaded_info["shape"] == info["shape"]
+
+
+def test_prepare_cube_info_uses_explicit_grid():
+    import numpy as np
+    import paddle
+
+    from ppmat.predictor.field_io import prepare_cube_info
+
+    origin = np.asarray([0.5, 1.0, 1.5], dtype=np.float32)
+    axes = np.stack(
+        np.meshgrid(
+            np.arange(2, dtype=np.float32),
+            np.arange(2, dtype=np.float32),
+            np.arange(2, dtype=np.float32),
+            indexing="ij",
+        ),
+        axis=-1,
+    )
+    grid = paddle.to_tensor((axes + origin).reshape(1, -1, 3))
+
+    info = prepare_cube_info({"shape": [2, 2, 2]}, grid)
+
+    np.testing.assert_allclose(info["origin"], origin)
+    np.testing.assert_allclose(info["cell"], np.eye(3) * 2)
+
+
+def test_electronic_structure_models_use_builtin_scatter():
+    for relative_path in [
+        "ppmat/models/infgcn/infgcn.py",
+        "ppmat/models/mateno/mateno.py",
+    ]:
+        source = (ROOT / relative_path).read_text()
+        assert "from paddle_scatter import scatter" not in source
+        assert "from ppmat.utils.scatter import scatter" in source
+
+
+def test_all_infgcn_configs_are_parseable_and_complete():
+    config_paths = sorted(INFGCN_CONFIG_DIR.glob("*.yaml"))
+    assert [path.name for path in config_paths] == [
+        "infgcn_md17_benzene.yaml",
+        "infgcn_md17_ethane.yaml",
+        "infgcn_md17_ethanol.yaml",
+        "infgcn_md17_malonaldehyde.yaml",
+        "infgcn_md17_phenol.yaml",
+        "infgcn_md17_resorcinol.yaml",
+        "infgcn_mp.yaml",
+        "infgcn_omol25_MC_5k_trimmed.yaml",
+        "infgcn_qm9.yaml",
+    ]
+
+    required_model_params = {
+        "n_atom_type",
+        "num_radial",
+        "num_spherical",
+        "radial_embed_size",
+        "radial_hidden_size",
+        "cutoff",
+        "grid_cutoff",
+    }
+
+    for config_path in config_paths:
+        cfg = OmegaConf.to_container(OmegaConf.load(config_path), resolve=True)
+        assert cfg["Model"]["__class_name__"] == "InfGCN", config_path.name
+        assert required_model_params.issubset(
+            cfg["Model"]["__init_params__"]
+        ), config_path.name
+
+        dataset_cfg = cfg["Dataset"]
+        for split in ["train", "val", "test"]:
+            split_cfg = dataset_cfg[split]
+            dataset = split_cfg["dataset"]
+            assert dataset["__class_name__"] in {
+                "DensityDataset",
+                "SmallDensityDataset",
+            }, config_path.name
+            assert "root" in dataset["__init_params__"], config_path.name
+            assert "sampler" in split_cfg, config_path.name
+            assert "loader" in split_cfg, config_path.name
+            assert isinstance(
+                split_cfg["loader"]["use_shared_memory"], bool
+            ), config_path.name
+            assert split_cfg["loader"]["collate_fn"] in {
+                "DensityCollator",
+                "DensityVoxelCollator",
+            }, config_path.name
+
+        predict_cfg = cfg["Predict"]
+        for key in [
+            "split",
+            "index",
+            "output_dir",
+            "grid_batch_size",
+            "save_true_cube",
+            "save_pred_cube",
+            "save_html",
+            "cube_dir",
+            "mol_grid_shape",
+            "mol_grid_padding",
+        ]:
+            assert key in predict_cfg, config_path.name
+
+
+def test_infgcn_readme_commands_and_config_links_are_clean():
+    readme_path = INFGCN_CONFIG_DIR / "README.md"
+    readme = readme_path.read_text()
+    example_mol = INFGCN_CONFIG_DIR / "example/methane.mol"
+
+    assert "--model_name infgcn_qm9" in readme
+    assert "--weights_name best.pdparams" in readme
+    assert "--mol_input electronic_structure/configs/infgcn/example/methane.mol" in (
+        readme
+    )
+    assert "--atom_file electronic_structure/configs/qm9.json" in readme
+    assert example_mol.is_file()
+    assert "conda run" not in readme
+    assert "/home/" not in readme
+    assert ".pt" not in readme
+    assert "_t_2026" not in readme
+    assert "_s_42.zip" not in readme
+
+    hrefs = re.findall(r'href="([^"]*configs/infgcn/[^"]+\.yaml)"', readme)
+    assert hrefs
+    for href in hrefs:
+        assert (readme_path.parent / href).resolve().exists(), href
+
+
+def test_infgcn_bundled_molecule_builds_inference_grid():
+    from ppmat.predictor.field_predictor import build_mol_sample
+    from ppmat.predictor.field_predictor import load_atom_mapping
+
+    example_mol = INFGCN_CONFIG_DIR / "example/methane.mol"
+    atom_file = ROOT / "electronic_structure/configs/qm9.json"
+    atom_name2idx, _ = load_atom_mapping(atom_file)
+
+    graph, density, grid_coord, info = build_mol_sample(
+        example_mol,
+        atom_name2idx,
+        mol_grid_shape=[8, 8, 8],
+        mol_grid_padding=6.0,
+    )
+
+    assert graph.x.shape == [5]
+    assert graph.pos.shape == [5, 3]
+    assert density is None
+    assert grid_coord.shape == [512, 3]
+    assert info["shape"] == [8, 8, 8]
+
+
+def test_diffnmr_sample_readme_documents_one_click_sample_command():
+    readme = (ROOT / "spectrum_elucidation/configs/diffnmr/README.md").read_text()
+    sample_csv = ROOT / "spectrum_elucidation/configs/diffnmr/example/sample.csv"
+
+    assert "--model_name='diffnmr_msdnmr_nless15'" in readme
+    assert "--weights_name='best.pdparams'" in readme
+    assert "bundled one-row validation example" in readme
+    assert "Sampler.data.dataset.__init_params__.path" in readme
+    assert (
+        "Sampler.data.dataset.__init_params__.path='./data/MSD_nmr/test.csv'" in readme
+    )
+    assert "### Sampling Sample" not in readme
+    assert "--checkpoint_path='./checkpoints'" in readme
+    assert sample_csv.exists()
+    assert sample_csv.read_text().splitlines()[0] == "smiles,tokenized_input,atom_count"
+    assert (
+        sample_csv.read_text()
+        .splitlines()[1]
+        .startswith('CSc1ccc(C(C)C(=O)O)cc1F,"{""1HNMR"":')
+    )
+
+
+def test_diffnmr_package_sample_defaults_to_bundled_example():
+    config = OmegaConf.to_container(
+        OmegaConf.load(ROOT / "spectrum_elucidation/configs/diffnmr/DiffNMR.yaml"),
+        resolve=False,
+    )
+
+    assert config["Sampler"]["name"] == "diffnmr"
+    assert config["Sampler"]["retrieval_database_path"] is None
+    sampler_params = config["Sampler"]["data"]["dataset"]["__init_params__"]
+    assert sampler_params["path"] == "./example/sample.csv"
+    assert sampler_params["vocab_peakwidth_path"] == (
+        "./spectrum_elucidation/vocab/nless15/H1_statistic/delta_distribution.csv"
+    )
+    assert sampler_params["vocab_split_path"] == (
+        "./spectrum_elucidation/vocab/nless15/H1_statistic/"
+        "split_type_distribution.csv"
+    )
+    assert sampler_params["cache_path"] == "./output/diffnmr_example_cache"
+    assert sampler_params["overwrite"] is True
+    assert config["Sampler"]["data"]["sampler"]["__init_params__"]["batch_size"] == 1
+    assert config["Sampler"]["sample_batch_iters"] == 1
+    assert config["Sampler"]["visual_num"] == 1
+    assert config["Sampler"]["chains_to_save"] == 0
+
+
+def test_molecular_sampler_entrypoint_keeps_diffnmr_imports_lazy():
+    source = (ROOT / "ppmat/sampler/molecular_sampler.py").read_text()
+
+    forbidden_snippets = [
+        "import paddle",
+        "from ppmat.datasets",
+        "from ppmat.metrics",
+        "from ppmat.models.diffnmr",
+        "from ppmat.schedulers",
+        "DiffNMRStreamingAdapter",
+        "ExtraMolecularFeatures",
+        "MolecularVisualization",
+        "scheduling_diffnmr",
+        "graphs_from_mol",
+    ]
+    for snippet in forbidden_snippets:
+        assert snippet not in source
+    assert "importlib.import_module" in source
+    assert "ppmat.sampler.diffnmr:DiffNMRSampler" in source
+    assert "MODEL_NAME_TO_SAMPLER" in source
+
+
+def test_molecular_sampler_module_source_load_does_not_load_diffnmr():
+    sys.modules.pop("ppmat.models.diffnmr.diffnmr", None)
+    sys.modules.pop("ppmat.sampler.diffnmr", None)
+
+    module_path = ROOT / "ppmat/sampler/molecular_sampler.py"
+    spec = importlib.util.spec_from_file_location(
+        "molecular_sampler_under_test", module_path
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    assert module.SAMPLER_REGISTRY["diffnmr"].endswith("diffnmr:DiffNMRSampler")
+    assert "ppmat.models.diffnmr.diffnmr" not in sys.modules
+    assert "ppmat.sampler.diffnmr" not in sys.modules
+
+
+def test_molecular_sampler_dispatches_diffnmr_config(tmp_path, monkeypatch):
+    import ppmat.sampler.molecular_sampler as molecular_sampler
+
+    class FakeSampler:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    def fake_import_module(module_name):
+        assert module_name == "fake_sampler_module"
+        return SimpleNamespace(FakeSampler=FakeSampler)
+
+    monkeypatch.setattr(
+        molecular_sampler,
+        "SAMPLER_REGISTRY",
+        {"diffnmr": "fake_sampler_module:FakeSampler"},
+    )
+    monkeypatch.setattr(
+        molecular_sampler.importlib, "import_module", fake_import_module
+    )
+
+    config_path = tmp_path / "DiffNMR.yaml"
+    checkpoint_path = tmp_path / "checkpoints"
+    checkpoint_path.mkdir()
+    config_path.write_text(
+        "Model:\n" "  __class_name__: DiffNMR\n" "Sampler:\n" "  name: diffnmr\n"
+    )
+
+    sampler = molecular_sampler.MolecularSampler(
+        config_path=str(config_path),
+        checkpoint_path=str(checkpoint_path),
+        config_overrides=["Sampler.name=diffnmr"],
+    )
+
+    assert isinstance(sampler, FakeSampler)
+    assert sampler.kwargs["config_path"] == str(config_path)
+    assert sampler.kwargs["checkpoint_path"] == str(checkpoint_path)
+    assert sampler.kwargs["config_overrides"] == ["Sampler.name=diffnmr"]
+
+
+def test_molecular_sampler_dispatches_registered_diffnmr_name(monkeypatch):
+    import ppmat.sampler.molecular_sampler as molecular_sampler
+
+    class FakeSampler:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    monkeypatch.setattr(
+        molecular_sampler,
+        "SAMPLER_REGISTRY",
+        {"diffnmr": "fake_sampler_module:FakeSampler"},
+    )
+    monkeypatch.setattr(
+        molecular_sampler.importlib,
+        "import_module",
+        lambda module_name: SimpleNamespace(FakeSampler=FakeSampler),
+    )
+
+    sampler = molecular_sampler.MolecularSampler(model_name="diffnmr_msdnmr_nless15")
+
+    assert isinstance(sampler, FakeSampler)
+    assert sampler.kwargs["model_name"] == "diffnmr_msdnmr_nless15"
+
+
+def test_molecular_sampler_updates_visualization_output_dir_for_save_path(tmp_path):
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    sampler = object.__new__(DiffNMRSampler)
+    sampler.sample_config = {"data": {}}
+    sampler.output_dir = "old_output"
+    sampler.visualization_tools = SimpleNamespace(result_path="old_output/graph/")
+    sampler.model = SimpleNamespace(eval=lambda: None)
+    sampler.flag_retrieval_sampling = False
+    sampler.num_candidates = 1
+    sampler.metric_dict_sample = {"Accuracy"}
+
+    save_path = tmp_path / "sample_output"
+    sampler.sample_epoch = lambda *args, **kwargs: {"Accuracy": 1.0}
+
+    result = sampler.sample_by_dataloader(
+        save_path=str(save_path),
+        data_loader=[],
+    )
+
+    assert result == {"Accuracy": 1.0}
+    assert sampler.output_dir == str(save_path)
+    assert sampler.visualization_tools.result_path == str(save_path / "graph")
+
+
+def test_molecular_sampler_compute_metric_reuses_sample_metrics(tmp_path):
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    sampler = object.__new__(DiffNMRSampler)
+    sampler.sample_config = {}
+    sampler.output_dir = "old_output"
+    sampler.visualization_tools = None
+    sampler.sample_by_dataloader = lambda save_path: {"Accuracy": 0.5}
+
+    assert sampler.compute_metric(save_path=str(tmp_path)) == {"Accuracy": 0.5}
+    assert sampler.output_dir == str(tmp_path)
+
+
+def test_molecular_sampler_clamps_keep_chain_to_batch_size():
+    import paddle
+
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    sampler = object.__new__(DiffNMRSampler)
+    assert sampler._clamp_keep_chain(5, 1) == 1
+    assert sampler._clamp_keep_chain(0, 1) == 0
+    assert sampler._clamp_keep_chain(3, paddle.to_tensor([2, 2], dtype="int64")) == 2
+
+
+def test_molecular_sampler_accumulates_conditions_across_batches():
+    import numpy as np
+
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    samples = {"batch_condition": []}
+    first_batch = [np.full((2, 3), branch, dtype=np.float32) for branch in range(4)]
+    second_batch = [
+        np.full((1, 3), branch + 10, dtype=np.float32) for branch in range(4)
+    ]
+
+    DiffNMRSampler._append_batch_conditions(samples, first_batch)
+    DiffNMRSampler._append_batch_conditions(samples, second_batch)
+
+    assert len(samples["batch_condition"]) == 4
+    for branch, tensor in enumerate(samples["batch_condition"]):
+        expected = np.concatenate([first_batch[branch], second_batch[branch]], axis=0)
+        np.testing.assert_allclose(tensor.numpy(), expected)
+
+
+def test_molecular_sampler_rejects_inconsistent_condition_branches():
+    import numpy as np
+    import pytest
+
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    samples = {
+        "batch_condition": [np.zeros((1, 2), dtype=np.float32) for _ in range(4)]
+    }
+    with pytest.raises(ValueError, match="Inconsistent number"):
+        DiffNMRSampler._append_batch_conditions(
+            samples,
+            [np.zeros((1, 2), dtype=np.float32) for _ in range(2)],
+        )
+
+
+def test_molecular_sampler_sample_epoch_preserves_two_batch_conditions(monkeypatch):
+    import numpy as np
+    import paddle
+
+    import ppmat.sampler.diffnmr as diffnmr_sampler
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    class DenseData:
+        X = paddle.zeros([1, 1, 1], dtype="float32")
+        E = paddle.zeros([1, 1, 1, 1], dtype="float32")
+
+        def mask(self, node_mask):
+            return self
+
+    class Streaming:
+        def __init__(self):
+            self.samples = None
+
+        def update_step(self, result, **kwargs):
+            self.samples = result["samples"]
+
+        def compute_epoch(self, **kwargs):
+            return {}
+
+    monkeypatch.setattr(
+        diffnmr_sampler.diffgraphformer_utils,
+        "to_dense",
+        lambda *args, **kwargs: (
+            DenseData(),
+            paddle.ones([1, 1], dtype="bool"),
+        ),
+    )
+
+    sampler = object.__new__(DiffNMRSampler)
+    sampler.model = SimpleNamespace(
+        eval=lambda: None,
+        conditioning_mode="spectrum",
+        dataset_info=SimpleNamespace(atom_decoder=["C"]),
+    )
+    sampler.sample_batch_iters = 2
+    sampler.visual_num = 0
+    sampler.chains_left_to_save = 0
+    sampler.number_chain_steps = 1
+    sampler.flag_use_formula = False
+    sampler.flag_retrieval_initialization = False
+    sampler.clip = None
+    sampler.molecular_vectors = None
+    sampler.smiles_list = None
+    sampler.streaming = Streaming()
+    sampler.rank = 0
+    sampler.output_dir = "unused"
+    sampler.sample_batch = lambda **kwargs: ([["pred"]], [["true"]])
+
+    def make_batch(offset):
+        return {
+            "graph": SimpleNamespace(
+                node_feat={"feat": np.zeros((1, 1), dtype=np.float32)},
+                edges=np.zeros((0, 2), dtype=np.int64),
+                edge_feat={"feat": np.zeros((0, 1), dtype=np.float32)},
+                graph_node_id=np.zeros(1, dtype=np.int64),
+            ),
+            "property": {
+                "atom_count": np.asarray([1], dtype=np.int64),
+                "y": np.asarray([[0]], dtype=np.float32),
+            },
+            "spectrum": {
+                "H_nmr": np.full((1, 2), offset, dtype=np.float32),
+                "num_H_peak": np.asarray([offset], dtype=np.float32),
+                "C_nmr": np.full((1, 3), offset + 1, dtype=np.float32),
+                "num_C_peak": np.asarray([offset + 1], dtype=np.float32),
+            },
+        }
+
+    sampler.sample_epoch(
+        [make_batch(1), make_batch(5)],
+        epoch_id=0,
+        num_candidates=1,
+    )
+
+    conditions = sampler.streaming.samples["batch_condition"]
+    assert [tensor.shape[0] for tensor in conditions] == [2, 2, 2, 2]
+    np.testing.assert_allclose(conditions[0].numpy()[:, 0], [1, 5])
+    assert sampler.streaming.samples["n_all"] == 2
+    assert len(sampler.streaming.samples["pred"]) == 2
+
+
+def test_diffnmr_spectrum_dispatch_uses_explicit_model_capability():
+    import paddle
+
+    from ppmat.schedulers.scheduling_diffnmr import _encode_spectrum_condition
+
+    model_source = (ROOT / "ppmat/models/diffnmr/diffnmr.py").read_text()
+    sampler_source = (ROOT / "ppmat/sampler/diffnmr.py").read_text()
+    scheduler_source = (ROOT / "ppmat/schedulers/scheduling_diffnmr.py").read_text()
+    assert 'self.conditioning_mode = "spectrum"' in model_source
+    assert 'getattr(self.model, "conditioning_mode", None)' in sampler_source
+    assert 'getattr(model, "conditioning_mode", None)' in scheduler_source
+    assert "isinstance(model, DiffNMR)" not in scheduler_source
+    assert "model.__class__ is DiffNMR" not in scheduler_source
+
+    expected_embedding = paddle.ones([1, 4])
+    expected_tokens = paddle.ones([1, 2, 4])
+    expected_mask = paddle.ones([1, 2])
+    condition = [object(), object(), object(), object()]
+
+    class Encoder:
+        def __call__(self, value):
+            assert value is condition
+            return expected_embedding, (expected_tokens, expected_mask)
+
+    model = SimpleNamespace(flag_onlyH=False, encoder=Encoder())
+    embedding, tokens = _encode_spectrum_condition(model, condition)
+
+    assert embedding is expected_embedding
+    assert tokens is expected_tokens
+
+
+def test_diffnmr_config_uses_standard_checkpoint_paths():
+    config_path = ROOT / "spectrum_elucidation/configs/diffnmr/DiffNMR.yaml"
+    source = config_path.read_text()
+    cfg = OmegaConf.to_container(OmegaConf.load(config_path), resolve=True)
+
+    assert "./pretrained/" not in source
+    assert cfg["Sampler"]["pretrained_model_path"] == "./checkpoints/best.pdparams"
+    assert cfg["Model"]["__init_params__"]["encoder_cfg"]["pretrained_path"].startswith(
+        "./checkpoints/"
+    )
+    assert cfg["Model"]["__init_params__"]["decoder_cfg"]["pretrained_path"].startswith(
+        "./checkpoints/"
+    )
+    assert cfg["CLIP"]["__init_params__"]["spectrum_encoder"][
+        "pretrained_model_path"
+    ].startswith("./checkpoints/")
+    assert cfg["CLIP"]["__init_params__"]["graph_encoder"][
+        "pretrained_model_path"
+    ].startswith("./checkpoints/")
+
+
+def test_molecular_sampler_resolves_diffnmr_checkpoint_paths(tmp_path):
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    package_dir = tmp_path / "diffnmr_msdnmr_nless15"
+    package_ckpt_dir = package_dir / "checkpoints"
+    package_assets_dir = package_dir / "assets"
+    package_vocab_dir = package_assets_dir / "vocab/nless15/H1_statistic"
+    package_retrieval_dir = package_assets_dir / "retrieval_database"
+    package_ckpt_dir.mkdir(parents=True)
+    package_vocab_dir.mkdir(parents=True)
+    package_retrieval_dir.mkdir(parents=True)
+    package_weight = package_ckpt_dir / "DiffNMR_NMRNet_nless15_best.pdparams"
+    package_vocab = package_vocab_dir / "delta_distribution.csv"
+    package_retrieval = (
+        package_retrieval_dir
+        / "msd_nmr_nless15_retrieval_molecular_representations.csv"
+    )
+    package_weight.write_bytes(b"fake")
+    package_vocab.write_text("Value,Count\n0.03,1\n")
+    package_retrieval.write_text("smiles,mol_rep\n")
+
+    package_config = {
+        "Model": {
+            "__init_params__": {
+                "encoder_cfg": {
+                    "pretrained_path": (
+                        "./checkpoints/DiffNMR_NMRNet_nless15_best.pdparams"
+                    )
+                }
+            }
+        },
+        "Dataset": {
+            "train": {
+                "dataset": {
+                    "__init_params__": {
+                        "vocab_peakwidth_path": (
+                            "./assets/vocab/nless15/H1_statistic/"
+                            "delta_distribution.csv"
+                        ),
+                    }
+                }
+            }
+        },
+        "Sampler": {
+            "retrieval_database_path": (
+                "./assets/retrieval_database/"
+                "msd_nmr_nless15_retrieval_molecular_representations.csv"
+            )
+        },
+    }
+
+    DiffNMRSampler._resolve_package_paths(
+        package_config,
+        config_base_dir=str(package_dir),
+        checkpoint_dir=None,
+    )
+
+    assert package_config["Model"]["__init_params__"]["encoder_cfg"][
+        "pretrained_path"
+    ] == str(package_weight)
+    assert package_config["Dataset"]["train"]["dataset"]["__init_params__"][
+        "vocab_peakwidth_path"
+    ] == str(package_vocab)
+    assert package_config["Sampler"]["retrieval_database_path"] == str(
+        package_retrieval
+    )
+
+    custom_ckpt_dir = tmp_path / "custom_checkpoints"
+    custom_ckpt_dir.mkdir()
+    custom_weight = custom_ckpt_dir / "DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+    custom_weight.write_bytes(b"fake")
+    custom_config = {
+        "CLIP": {
+            "__init_params__": {
+                "graph_encoder": {
+                    "pretrained_model_path": (
+                        "./checkpoints/DiffNMR_DiffGraphFormer_nless15_best.pdparams"
+                    )
+                }
+            }
+        }
+    }
+
+    DiffNMRSampler._resolve_package_paths(
+        custom_config,
+        config_base_dir=str(tmp_path / "config_dir"),
+        checkpoint_dir=str(custom_ckpt_dir),
+    )
+
+    assert custom_config["CLIP"]["__init_params__"]["graph_encoder"][
+        "pretrained_model_path"
+    ] == str(custom_weight)
+
+
+def test_molecular_sampler_requires_packaged_diffnmr_vocab(tmp_path):
+    import pytest
+
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    package_dir = tmp_path / "diffnmr_msdnmr_nless15"
+    package_dir.mkdir()
+    config = {
+        "Sampler": {
+            "data": {
+                "dataset": {
+                    "__init_params__": {
+                        "vocab_peakwidth_path": (
+                            "./assets/vocab/nless15/H1_statistic/"
+                            "delta_distribution.csv"
+                        ),
+                        "vocab_split_path": (
+                            "./assets/vocab/nless15/H1_statistic/"
+                            "split_type_distribution.csv"
+                        ),
+                    }
+                }
+            }
+        }
+    }
+
+    with pytest.raises(FileNotFoundError, match="vocab_peakwidth_path"):
+        DiffNMRSampler._resolve_package_paths(
+            config,
+            config_base_dir=str(package_dir),
+            checkpoint_dir=None,
+        )
+
+
+def test_molecular_sampler_requires_retrieval_database_only_when_enabled(tmp_path):
+    import pytest
+
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    package_dir = tmp_path / "diffnmr_msdnmr_nless15"
+    package_dir.mkdir()
+    config = {
+        "Sampler": {
+            "flag_retrieval_sampling": True,
+            "retrieval_database_path": "./missing/retrieval.csv",
+        }
+    }
+
+    with pytest.raises(FileNotFoundError, match="retrieval_database_path"):
+        DiffNMRSampler._resolve_package_paths(
+            config,
+            config_base_dir=str(package_dir),
+            checkpoint_dir=None,
+        )
+
+    config["Sampler"]["flag_retrieval_sampling"] = False
+    DiffNMRSampler._resolve_package_paths(
+        config,
+        config_base_dir=str(package_dir),
+        checkpoint_dir=None,
+    )
+
+
+def test_molecular_sampler_allows_zero_saved_chains(monkeypatch):
+    import paddle
+
+    import ppmat.sampler.diffnmr as diffnmr_sampler
+    from ppmat.sampler.diffnmr import DiffNMRSampler
+
+    class FakeData:
+        def __init__(self, X, E, y=None):
+            self.X = X
+            self.E = E
+            self.y = y
+
+        def mask(self, node_mask, collapse=False):
+            if collapse:
+                return FakeData(
+                    paddle.argmax(self.X, axis=-1),
+                    paddle.argmax(self.E, axis=-1),
+                    self.y,
+                )
+            return self
+
+    class FakeModel:
+        T = 1
+        limit_dist = None
+
+    def fake_noise(limit_dist, node_mask):
+        del limit_dist
+        batch_size, n_max = node_mask.shape
+        return FakeData(
+            paddle.ones([batch_size, n_max, 1], dtype="float32"),
+            paddle.ones([batch_size, n_max, n_max, 1], dtype="float32"),
+            paddle.zeros([batch_size, 1], dtype="float32"),
+        )
+
+    def fake_step(model, **kwargs):
+        del model
+        batch_size = kwargs["X_t"].shape[0]
+        n_max = kwargs["X_t"].shape[1]
+        sampled = FakeData(
+            paddle.ones([batch_size, n_max, 1], dtype="float32"),
+            paddle.ones([batch_size, n_max, n_max, 1], dtype="float32"),
+            paddle.zeros([batch_size, 1], dtype="float32"),
+        )
+        discrete = FakeData(
+            paddle.zeros([batch_size, n_max], dtype="int64"),
+            paddle.zeros([batch_size, n_max, n_max], dtype="int64"),
+        )
+        return sampled, discrete
+
+    monkeypatch.setattr(
+        diffnmr_sampler.scheduling_diffnmr,
+        "sample_discrete_feature_noise",
+        fake_noise,
+    )
+    monkeypatch.setattr(diffnmr_sampler.scheduling_diffnmr, "step", fake_step)
+
+    sampler = object.__new__(DiffNMRSampler)
+    sampler.visualization_tools = None
+
+    mol_list, mol_true = sampler.sample_batch(
+        model=FakeModel(),
+        batch_id=0,
+        batch_size=1,
+        batch_condition=[],
+        number_chain_steps=1,
+        keep_chain=0,
+        visual_num=0,
+        batch_X=paddle.ones([1, 1, 1], dtype="float32"),
+        batch_E=paddle.ones([1, 1, 1, 1], dtype="float32"),
+        batch_y=paddle.zeros([1, 1], dtype="float32"),
+        iter_idx=0,
+        num_nodes=paddle.to_tensor([1], dtype="int64"),
+    )
+
+    assert len(mol_list) == 1
+    assert len(mol_true) == 1
+
+
+def test_diffnmr_sample_entrypoint_supports_config_overrides():
+    source = (ROOT / "spectrum_elucidation/sample.py").read_text()
+    sampler_source = (ROOT / "ppmat/sampler/diffnmr.py").read_text()
+
+    assert "parse_known_args()" in source
+    assert "config_overrides=config_overrides" in source
+    assert "config_overrides: Optional[List[str]] = None" in sampler_source
+    assert "OmegaConf.merge(config, cli_config)" in sampler_source
+    assert "_apply_package_support_files" not in sampler_source
+    assert "_replace_with_package_file" not in sampler_source
+
+
+def test_diffnmr_train_entrypoint_builds_training_statistics_for_eval_and_test():
+    source = (ROOT / "spectrum_elucidation/train.py").read_text()
+
+    assert "if do_train or do_eval or do_test:" in source
+    assert "At least one of Global.do_train" in source
+    assert (
+        "dataloaders.train_dataloader"
+        in (ROOT / "ppmat/datasets/msd_nmr_dataset.py").read_text()
+    )
+
+
+def test_diffnmr_uses_molecular_sampler_from_sampler_package():
+    source = (ROOT / "spectrum_elucidation/sample.py").read_text()
+    sampler_path = ROOT / "ppmat/sampler/molecular_sampler.py"
+    diffnmr_sampler_path = ROOT / "ppmat/sampler/diffnmr.py"
+    diffnmr_sampler_source = diffnmr_sampler_path.read_text()
+    legacy_sample_dir = ROOT / "ppmat/sample"
+
+    assert sampler_path.exists()
+    assert diffnmr_sampler_path.exists()
+    assert not legacy_sample_dir.exists()
+    assert "from ppmat.sampler import MolecularSampler" in source
+    assert "class MolecularSampler" in sampler_path.read_text()
+    assert "class DiffNMRSampler" in diffnmr_sampler_source
+    assert 'setattr(self.model, "clip"' not in diffnmr_sampler_source
+    assert 'setattr(self.model, "streaming_adapter"' not in diffnmr_sampler_source
diff --git a/test/test_registered_model_end_to_end.py b/test/test_registered_model_end_to_end.py
new file mode 100644
index 00000000..85828cf8
--- /dev/null
+++ b/test/test_registered_model_end_to_end.py
@@ -0,0 +1,226 @@
+"""Opt-in end-to-end tests for every registered model variant.
+
+These tests complement package-loading checks by exercising the public task
+interfaces and asserting reader-visible output artifacts.
+
+Run with:
+
+    RUN_REGISTERED_MODEL_E2E=1 \
+      pytest test/test_registered_model_end_to_end.py -q
+"""
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import paddle
+import pytest
+from PIL import Image
+
+from ppmat.models import MODEL_REGISTRY
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("RUN_REGISTERED_MODEL_E2E") != "1",
+    reason="set RUN_REGISTERED_MODEL_E2E=1 to run published-model workflows",
+)
+
+ROOT = Path(__file__).resolve().parents[1]
+
+CRYSTAL_PROPERTY_MODELS = tuple(
+    sorted(
+        name
+        for name in MODEL_REGISTRY
+        if name.startswith(("comformer_", "dimenetpp_", "megnet_"))
+    )
+)
+MOLECULAR_PROPERTY_MODELS = tuple(
+    sorted(name for name in MODEL_REGISTRY if name.startswith("spherenet_qm9_"))
+)
+POTENTIAL_MODELS = (
+    "chgnet_mptrj",
+    "mattersim_1M",
+    "mattersim_5M",
+    *sorted(name for name in MODEL_REGISTRY if name.startswith("spherenet_md17_")),
+)
+STRUCTURE_MODELS = tuple(
+    sorted(
+        name for name in MODEL_REGISTRY if name.startswith(("diffcsp_", "mattergen_"))
+    )
+)
+SPECTRUM_MODELS = tuple(
+    sorted(name for name in MODEL_REGISTRY if name.startswith("sfin_"))
+)
+FIELD_MODELS = tuple(
+    sorted(name for name in MODEL_REGISTRY if name.startswith("infgcn_"))
+)
+
+CONDITION_VALUES = {
+    "chemical_system": "Mo-Si",
+    "dft_band_gap": 0.897,
+    "dft_bulk_modulus": 100.0,
+    "dft_mag_density": 0.1,
+    "energy_above_hull": 0.05,
+    "hhi_score": 0.5,
+    "ml_bulk_modulus": 100.0,
+    "space_group": 225,
+    "spacegroup": 225,
+}
+
+
+@pytest.mark.parametrize("model_name", CRYSTAL_PROPERTY_MODELS)
+def test_crystal_property_prediction(model_name, tmp_path):
+    from property_prediction.predict import PropertyPredictor
+
+    output_path = tmp_path / f"{model_name}.csv"
+    predictor = PropertyPredictor(model_name=model_name)
+    result = predictor.from_cif_file(
+        str(ROOT / "property_prediction/example_data/cifs/mp-18767-LiMnO2.cif"),
+        str(output_path),
+    )
+
+    assert result
+    assert output_path.is_file()
+
+
+@pytest.mark.parametrize("model_name", MOLECULAR_PROPERTY_MODELS)
+def test_molecular_property_prediction(model_name, tmp_path):
+    from property_prediction.predict import PropertyPredictor
+
+    output_path = tmp_path / f"{model_name}.csv"
+    predictor = PropertyPredictor(model_name=model_name)
+    result = predictor.from_xyz_file(
+        str(ROOT / "property_prediction/example_data/molecules/isoguvacine.xyz"),
+        str(output_path),
+    )
+
+    assert result
+    assert output_path.is_file()
+
+
+@pytest.mark.parametrize("model_name", POTENTIAL_MODELS)
+def test_interatomic_potential_prediction(model_name, tmp_path):
+    from interatomic_potentials.predict import PotentialPredictor
+
+    output_path = tmp_path / f"{model_name}.csv"
+    predictor = PotentialPredictor(model_name=model_name)
+    if model_name.startswith("spherenet_md17_"):
+        result = predictor.from_xyz_file(
+            str(ROOT / "interatomic_potentials/example_data/xyz/md17_aspirin.xyz"),
+            str(output_path),
+        )
+    else:
+        result = predictor.from_cif_file(
+            str(ROOT / "interatomic_potentials/example_data/cifs/mp-18767-LiMnO2.cif"),
+            str(output_path),
+        )
+
+    assert result
+    assert output_path.is_file()
+
+
+@pytest.mark.parametrize("model_name", STRUCTURE_MODELS)
+def test_structure_generation(model_name, tmp_path):
+    from ppmat.sampler import StructureSampler
+
+    paddle.seed(42)
+    output_dir = tmp_path / model_name
+    sampler = StructureSampler(model_name=model_name)
+    sample_params = {"num_inference_steps": 2}
+
+    if model_name == "diffcsp_mp20":
+        result = sampler.sample_by_chemical_formula(
+            "LiMnO2",
+            save_path=str(output_dir),
+            sample_params=sample_params,
+        )
+    elif getattr(sampler.model, "condition_names", None):
+        conditions = {
+            name: CONDITION_VALUES[name] for name in sampler.model.condition_names
+        }
+        result = sampler.sample_by_condition(
+            4,
+            conditions,
+            save_path=str(output_dir),
+            sample_params=sample_params,
+        )
+    else:
+        result = sampler.sample_by_num_atoms(
+            4,
+            save_path=str(output_dir),
+            sample_params=sample_params,
+        )
+
+    assert result["result"]
+    assert list(output_dir.glob("*.cif"))
+
+
+@pytest.mark.parametrize("model_name", SPECTRUM_MODELS)
+def test_spectrum_enhancement_prediction(model_name, tmp_path):
+    from spectrum_enhancement.predict import SpectrumPredictor
+
+    input_path = tmp_path / "input.png"
+    Image.fromarray(np.arange(4096, dtype=np.uint8).reshape(64, 64)).save(input_path)
+    output_dir = tmp_path / model_name
+
+    predictor = SpectrumPredictor(model_name=model_name)
+    saved_paths = predictor.from_image_path(
+        str(input_path),
+        str(output_dir),
+    )
+
+    assert len(saved_paths) == 1
+    assert Path(saved_paths[0]).is_file()
+
+
+@pytest.mark.parametrize("model_name", FIELD_MODELS)
+def test_electronic_structure_prediction(model_name, tmp_path):
+    from electronic_structure.predict import build_parser
+    from ppmat.predictor import FieldPredictor
+
+    if model_name.startswith("infgcn_md17_"):
+        atom_file = tmp_path / "md17_atoms.json"
+        atom_file.write_text(
+            json.dumps(
+                [
+                    {"name": "C", "atom_num": 6},
+                    {"name": "H", "atom_num": 1},
+                    {"name": "O", "atom_num": 8},
+                ]
+            )
+        )
+    elif model_name == "infgcn_qm9":
+        atom_file = ROOT / "electronic_structure/configs/qm9.json"
+    else:
+        atom_file = ROOT / "electronic_structure/configs/crystal.json"
+
+    output_dir = tmp_path / model_name
+    args = build_parser().parse_args(
+        [
+            "--model_name",
+            model_name,
+            "--weights_name",
+            "best.pdparams",
+            "--mol_input",
+            str(ROOT / "electronic_structure/configs/infgcn/example/methane.mol"),
+            "--atom_file",
+            str(atom_file),
+            "--mol_grid_shape",
+            "8",
+            "--grid_batch_size",
+            "128",
+            "--output_dir",
+            str(output_dir),
+            "--cube_dir",
+            str(output_dir),
+            "--save_pred_cube",
+            "--skip_vis",
+        ]
+    )
+    predictor = FieldPredictor(
+        model_name=args.model_name,
+        weights_name=args.weights_name,
+    )
+    predictor.predict(args)
+
+    assert (output_dir / "methane_pred.cube").is_file()
diff --git a/test/test_registered_model_packages_integration.py b/test/test_registered_model_packages_integration.py
new file mode 100644
index 00000000..c74d67de
--- /dev/null
+++ b/test/test_registered_model_packages_integration.py
@@ -0,0 +1,119 @@
+"""Opt-in tests for the large model packages published in MODEL_REGISTRY.
+
+Run after uploading release archives:
+
+    RUN_MODEL_PACKAGE_INTEGRATION=1 \
+      pytest test/test_registered_model_packages_integration.py -q
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import paddle
+import pytest
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("RUN_MODEL_PACKAGE_INTEGRATION") != "1",
+    reason="set RUN_MODEL_PACKAGE_INTEGRATION=1 to download published packages",
+)
+
+INFGCN_MODEL_NAMES = (
+    "infgcn_md17_benzene",
+    "infgcn_md17_ethane",
+    "infgcn_md17_ethanol",
+    "infgcn_md17_malonaldehyde",
+    "infgcn_md17_phenol",
+    "infgcn_md17_resorcinol",
+    "infgcn_mp",
+    "infgcn_omol25_mc_5k_trimmed",
+    "infgcn_qm9",
+)
+
+
+def _assert_raw_state_dict(checkpoint_path, model):
+    checkpoint = paddle.load(str(checkpoint_path))
+    assert isinstance(checkpoint, dict)
+    assert "model" not in checkpoint
+    assert "state_dict" not in checkpoint
+
+    model_state = model.state_dict()
+    assert set(checkpoint) == set(model_state)
+    for name, expected in checkpoint.items():
+        np.testing.assert_allclose(
+            model_state[name].numpy(),
+            expected.numpy(),
+            rtol=0,
+            atol=0,
+            err_msg=name,
+        )
+
+
+@pytest.mark.parametrize("model_name", INFGCN_MODEL_NAMES)
+def test_published_infgcn_packages_load_exact_weights(model_name):
+    from ppmat.models import MODEL_REGISTRY
+    from ppmat.models import build_model_from_name
+    from ppmat.utils import download
+    from ppmat.utils.model_package import resolve_model_package_dir
+
+    extracted_path = download.get_weights_path_from_url(MODEL_REGISTRY[model_name])
+    package_dir = Path(resolve_model_package_dir(model_name, extracted_path))
+    checkpoint_path = package_dir / "checkpoints" / "best.pdparams"
+    assert checkpoint_path.is_file()
+
+    model, config = build_model_from_name(model_name, "best.pdparams")
+
+    assert config["Model"]["__class_name__"] == "InfGCN"
+    _assert_raw_state_dict(checkpoint_path, model)
+
+
+def test_published_infgcn_qm9_runs_documented_inference(tmp_path):
+    model_name = "infgcn_qm9"
+    output_dir = tmp_path / "infgcn"
+    command = [
+        sys.executable,
+        "electronic_structure/predict.py",
+        "--model_name",
+        model_name,
+        "--weights_name",
+        "best.pdparams",
+        "--mol_input",
+        "electronic_structure/configs/infgcn/example/methane.mol",
+        "--atom_file",
+        "electronic_structure/configs/qm9.json",
+        "--mol_grid_shape",
+        "8",
+        "--grid_batch_size",
+        "128",
+        "--output_dir",
+        str(output_dir),
+        "--cube_dir",
+        str(output_dir),
+        "--skip_vis",
+    ]
+    subprocess.run(command, check=True)
+    assert (output_dir / "methane_pred.cube").is_file()
+
+
+def test_published_diffnmr_package_loads_and_samples(tmp_path):
+    from ppmat.models import MODEL_REGISTRY
+    from ppmat.sampler import MolecularSampler
+    from ppmat.utils import download
+    from ppmat.utils.model_package import resolve_model_package_dir
+
+    model_name = "diffnmr_msdnmr_nless15"
+    extracted_path = download.get_weights_path_from_url(MODEL_REGISTRY[model_name])
+    package_dir = Path(resolve_model_package_dir(model_name, extracted_path))
+    checkpoint_path = package_dir / "checkpoints" / "best.pdparams"
+    assert checkpoint_path.is_file()
+
+    sampler = MolecularSampler(
+        model_name=model_name,
+        weights_name="best.pdparams",
+    )
+    _assert_raw_state_dict(checkpoint_path, sampler.model)
+
+    result = sampler.sample_by_dataloader(save_path=str(tmp_path / "sample"))
+    assert result["Total Number"] == 1
diff --git a/test/test_scatter.py b/test/test_scatter.py
index 0d4b92ac..5f162eee 100644
--- a/test/test_scatter.py
+++ b/test/test_scatter.py
@@ -16,6 +16,7 @@
 import paddle
 
 from ppmat.utils.scatter import scatter_argmin
+from ppmat.utils.scatter import scatter_sum
 
 
 def test_scatter_argmin_handles_unsorted_and_empty_groups():
@@ -43,3 +44,18 @@ def test_scatter_argmin_handles_empty_input():
     result = scatter_argmin(values, groups, dim_size=3)
 
     np.testing.assert_array_equal(result.numpy(), [-1, -1, -1])
+
+
+def test_scatter_sum_dim_zero_supports_matrix_values():
+    values = paddle.arange(24, dtype="float32").reshape([4, 2, 3])
+    groups = paddle.to_tensor([0, 1, 0, 1], dtype="int64")
+
+    result = scatter_sum(values, groups, dim=0, dim_size=2)
+
+    expected = np.stack(
+        [
+            values.numpy()[[0, 2]].sum(axis=0),
+            values.numpy()[[1, 3]].sum(axis=0),
+        ]
+    )
+    np.testing.assert_allclose(result.numpy(), expected)
diff --git a/test/test_train_output_dir_timestamp.py b/test/test_train_output_dir_timestamp.py
new file mode 100644
index 00000000..e03f8858
--- /dev/null
+++ b/test/test_train_output_dir_timestamp.py
@@ -0,0 +1,38 @@
+import datetime
+import re
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from ppmat.utils.io import append_timestamp_to_output_dir
+
+
+def test_append_timestamp_to_output_dir_uses_seed_and_timestamp():
+    config = OmegaConf.create({"Trainer": {"output_dir": "./output/demo", "seed": 7}})
+
+    append_timestamp_to_output_dir(
+        config,
+        now=datetime.datetime(2026, 6, 29, 12, 34, 56),
+    )
+
+    assert config["Trainer"]["output_dir"] == "./output/demo_t_20260629_123456_s_7"
+
+
+def test_output_dir_helper_lives_in_io_module():
+    assert not Path("ppmat/utils/output_dir.py").exists()
+
+
+def test_train_entrypoints_append_timestamp_by_default():
+    train_scripts = [
+        "property_prediction/train.py",
+        "electronic_structure/train.py",
+        "structure_generation/train.py",
+        "spectrum_elucidation/train.py",
+        "interatomic_potentials/train.py",
+        "spectrum_enhancement/train.py",
+    ]
+
+    for script in train_scripts:
+        source = Path(script).read_text()
+        assert re.search(r"append_timestamp_to_output_dir\((config|cfg)\)", source)
+        assert "--append_timestamp" not in source