调整mattergen-wrapper位置

mars_toolkit/core/config.py

@@ -27,6 +27,7 @@ class Config:
 
     # MatterGen
     MATTERGENMODEL_ROOT = '/home/ubuntu/50T/lzy/mars-mcp/pretrained_models/mattergen_ckpt'
+    MATTERGEN_ROOT='/home/ubuntu/50T/lzy/mars-mcp/mattergen'
     MATTERGENMODEL_RESULT_PATH = 'results/'
 
     # Dify

mars_toolkit/core/mattergen_wrapper.py

@@ -0,0 +1,26 @@
+"""
+This is a wrapper module that provides access to the mattergen modules
+by modifying the Python path at runtime.
+"""
+import sys
+import os
+from pathlib import Path
+from .config import config
+# Add the mattergen directory to the Python path
+mattergen_dir = config.MATTERGEN_ROOT
+sys.path.insert(0, mattergen_dir)
+
+# Import the necessary modules from the mattergen package
+try:
+    from mattergen import generator
+    from mattergen.common.data import chemgraph
+    from mattergen.common.data.types import TargetProperty
+    from mattergen.common.utils.eval_utils import MatterGenCheckpointInfo
+    from mattergen.common.utils.data_classes import PRETRAINED_MODEL_NAME
+except ImportError as e:
+    print(f"Error importing mattergen modules: {e}")
+    print(f"Python path: {sys.path}")
+    raise
+CrystalGenerator = generator.CrystalGenerator
+# Re-export the modules
+__all__ = ['generator', 'chemgraph', 'TargetProperty', 'MatterGenCheckpointInfo', 'PRETRAINED_MODEL_NAME','CrystalGenerator']