初次提交

sci_mcp/material_mcp/mattergen_gen/material_gen_tools.py

@@ -0,0 +1,240 @@
+
+import ast
+import json
+import logging
+
+# Configure logger
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+handler = logging.StreamHandler()
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+handler.setFormatter(formatter)
+logger.addHandler(handler)
+
+import tempfile
+import os
+import datetime
+import asyncio
+import zipfile
+import shutil
+import re
+import multiprocessing
+from multiprocessing import Process, Queue
+from pathlib import Path
+from typing import Literal, Dict, Any, Tuple, Union, Optional, List
+import logging
+# 设置多进程启动方法为spawn，解决CUDA初始化错误
+try:
+    multiprocessing.set_start_method('spawn', force=True)
+except RuntimeError:
+    # 如果已经设置过启动方法，会抛出RuntimeError
+    pass
+
+from ase.optimize import FIRE
+from ase.filters import FrechetCellFilter
+from ase.atoms import Atoms
+from ase.io import read, write
+from pymatgen.core.structure import Structure
+from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
+from pymatgen.io.cif import CifWriter
+
+# 导入路径已更新
+from ...core.llm_tools import llm_tool
+from .mattergen_wrapper import *
+
+# 使用mattergen_wrapper
+import sys
+import os
+
+
+
+def convert_values(data_str):
+    """
+    将字符串转换为字典
+
+    Args:
+        data_str: JSON字符串
+
+    Returns:
+        解析后的数据，如果解析失败则返回原字符串
+    """
+    try:
+        data = json.loads(data_str)
+    except json.JSONDecodeError:
+        return data_str  # 如果无法解析为JSON，返回原字符串
+
+    return data
+
+
+def preprocess_property(property_name: str, property_value: Union[str, float, int]) -> Tuple[str, Any]:
+    """
+    Preprocess a property value based on its name, converting it to the appropriate type.
+
+    Args:
+        property_name: Name of the property
+        property_value: Value of the property (can be string, float, or int)
+
+    Returns:
+        Tuple of (property_name, processed_value)
+
+    Raises:
+        ValueError: If the property value is invalid for the given property name
+    """
+    valid_properties = [
+        "dft_mag_density", "dft_bulk_modulus", "dft_shear_modulus",
+        "energy_above_hull", "formation_energy_per_atom", "space_group",
+        "hhi_score", "ml_bulk_modulus", "chemical_system", "dft_band_gap"
+    ]
+
+    if property_name not in valid_properties:
+        raise ValueError(f"Invalid property_name: {property_name}. Must be one of: {', '.join(valid_properties)}")
+
+    # Process property_value if it's a string
+    if isinstance(property_value, str):
+        try:
+            # Try to convert string to float for numeric properties
+            if property_name != "chemical_system":
+                property_value = float(property_value)
+        except ValueError:
+            # If conversion fails, keep as string (for chemical_system)
+            pass
+
+    # Handle special cases for properties that need specific types
+    if property_name == "chemical_system":
+        if isinstance(property_value, (int, float)):
+            logger.warning(f"Converting numeric property_value {property_value} to string for chemical_system property")
+            property_value = str(property_value)
+    elif property_name == "space_group" :
+        space_group = property_value
+        if space_group < 1 or space_group > 230:
+            raise ValueError(f"Invalid space_group value: {space_group}. Must be an integer between 1 and 230.")
+
+    return property_name, property_value
+
+
+def main(
+    output_path: str,
+    pretrained_name: PRETRAINED_MODEL_NAME | None = None,
+    model_path: str | None = None,
+    batch_size: int = 2,
+    num_batches: int = 1,
+    config_overrides: list[str] | None = None,
+    checkpoint_epoch: Literal["best", "last"] | int = "last",
+    properties_to_condition_on: TargetProperty | None = None,
+    sampling_config_path: str | None = None,
+    sampling_config_name: str = "default",
+    sampling_config_overrides: list[str] | None = None,
+    record_trajectories: bool = True,
+    diffusion_guidance_factor: float | None = None,
+    strict_checkpoint_loading: bool = True,
+    target_compositions: list[dict[str, int]] | None = None,
+):
+    """
+    Evaluate diffusion model against molecular metrics.
+
+    Args:
+        model_path: Path to DiffusionLightningModule checkpoint directory.
+        output_path: Path to output directory.
+        config_overrides: Overrides for the model config, e.g., `model.num_layers=3 model.hidden_dim=128`.
+        properties_to_condition_on: Property value to draw conditional sampling with respect to. When this value is an empty dictionary (default), unconditional samples are drawn.
+        sampling_config_path: Path to the sampling config file. (default: None, in which case we use `DEFAULT_SAMPLING_CONFIG_PATH` from explorers.common.utils.utils.py)
+        sampling_config_name: Name of the sampling config (corresponds to `{sampling_config_path}/{sampling_config_name}.yaml` on disk). (default: default)
+        sampling_config_overrides: Overrides for the sampling config, e.g., `condition_loader_partial.batch_size=32`.
+        load_epoch: Epoch to load from the checkpoint. If None, the best epoch is loaded. (default: None)
+        record: Whether to record the trajectories of the generated structures. (default: True)
+        strict_checkpoint_loading: Whether to raise an exception when not all parameters from the checkpoint can be matched to the model.
+        target_compositions: List of dictionaries with target compositions to condition on. Each dictionary should have the form `{element: number_of_atoms}`. If None, the target compositions are not conditioned on.
+           Only supported for models trained for crystal structure prediction (CSP) (default: None)
+
+    NOTE: When specifying dictionary values via the CLI, make sure there is no whitespace between the key and value, e.g., `--properties_to_condition_on={key1:value1}`.
+    """
+    assert (
+        pretrained_name is not None or model_path is not None
+    ), "Either pretrained_name or model_path must be provided."
+    assert (
+        pretrained_name is None or model_path is None
+    ), "Only one of pretrained_name or model_path can be provided."
+
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+
+    sampling_config_overrides = sampling_config_overrides or []
+    config_overrides = config_overrides or []
+    properties_to_condition_on = properties_to_condition_on or {}
+    target_compositions = target_compositions or []
+
+    if pretrained_name is not None:
+        checkpoint_info = MatterGenCheckpointInfo.from_hf_hub(
+            pretrained_name, config_overrides=config_overrides
+        )
+    else:
+        checkpoint_info = MatterGenCheckpointInfo(
+            model_path=Path(model_path).resolve(),
+            load_epoch=checkpoint_epoch,
+            config_overrides=config_overrides,
+            strict_checkpoint_loading=strict_checkpoint_loading,
+        )
+    _sampling_config_path = Path(sampling_config_path) if sampling_config_path is not None else None
+    generator = CrystalGenerator(
+        checkpoint_info=checkpoint_info,
+        properties_to_condition_on=properties_to_condition_on,
+        batch_size=batch_size,
+        num_batches=num_batches,
+        sampling_config_name=sampling_config_name,
+        sampling_config_path=_sampling_config_path,
+        sampling_config_overrides=sampling_config_overrides,
+        record_trajectories=record_trajectories,
+        diffusion_guidance_factor=(
+            diffusion_guidance_factor if diffusion_guidance_factor is not None else 0.0
+        ),
+        target_compositions_dict=target_compositions,
+    )
+    generator.generate(output_dir=Path(output_path))
+
+
+@llm_tool(name="generate_material_MatterGen", description="Generate crystal structures with optional property constraints using MatterGen model")
+def generate_material_MatterGen(
+    properties: Optional[Dict[str, Union[float, str, Dict[str, Union[float, str]]]]] = None,
+    batch_size: int = 2,
+    num_batches: int = 1,
+    diffusion_guidance_factor: float = 2.0
+) -> str:
+    """
+    Generate crystal structures with optional property constraints.
+
+    This unified function can generate materials in three modes:
+    1. Unconditional generation (no properties specified)
+    2. Single property conditional generation (one property specified)
+    3. Multi-property conditional generation (multiple properties specified)
+
+    Args:
+        properties: Optional property constraints. Can be:
+                   - None or empty dict for unconditional generation
+                   - Dict with single key-value pair for single property conditioning
+                   - Dict with multiple key-value pairs for multi-property conditioning
+                   Valid property names include: "dft_band_gap", "chemical_system", etc.
+        batch_size: Number of structures per batch
+        num_batches: Number of batches to generate
+        diffusion_guidance_factor: Controls adherence to target properties
+
+    Returns:
+        Descriptive text with generated crystal structures in CIF format
+    """
+
+
+    # 导入MatterGenService
+    from .mattergen_service import MatterGenService
+    logger.info("子进程成功导入MatterGenService")
+
+    # 获取MatterGenService实例
+    service = MatterGenService.get_instance()
+    logger.info("子进程成功获取MatterGenService实例")
+
+    # 使用服务生成材料
+    logger.info("子进程开始调用generate方法...")
+    result = service.generate(properties=properties, batch_size=batch_size, num_batches=num_batches, diffusion_guidance_factor=diffusion_guidance_factor)
+    logger.info("子进程generate方法调用完成")
+    if "Error generating structures" in result:
+        return f"Error: Invalid properties {properties}."
+    else:
+        return result

sci_mcp/material_mcp/mattergen_gen/mattergen_service.py

@@ -0,0 +1,466 @@
+"""
+MatterGen service for mars_toolkit.
+
+This module provides a service for generating crystal structures using MatterGen.
+The service initializes the CrystalGenerator once and reuses it for multiple
+generation requests, improving performance.
+"""
+
+import datetime
+import os
+import logging
+import json
+from pathlib import Path
+import re
+from typing import Dict, Any, Optional, Union, List
+import threading
+import torch
+
+from .mattergen_wrapper import *
+from ...core.config import material_config
+
+logger = logging.getLogger(__name__)
+
+
+
+def format_cif_content(content):
+    """
+    Format CIF content by removing unnecessary headers and organizing each CIF file.
+
+    Args:
+        content: String containing CIF content, possibly with PK headers
+
+    Returns:
+        Formatted string with each CIF file properly labeled and formatted
+    """
+    # 如果内容为空，直接返回空字符串
+    if not content or content.strip() == '':
+        return ''
+
+    # 删除从PK开始到第一个_chemical_formula_structural之前的所有内容
+    content = re.sub(r'PK.*?(?=_chemical_formula_structural)', '', content, flags=re.DOTALL)
+
+    # 删除从PK开始到字符串结束且没有_chemical_formula_structural的内容
+    content = re.sub(r'PK[^_]*$', '', content, flags=re.DOTALL)
+    content = re.sub(r'PK.*?(?!.*_chemical_formula_structural)$', '', content, flags=re.DOTALL)
+
+    # 使用_chemical_formula_structural作为分隔符来分割不同的CIF文件
+    # 但我们需要保留这个字段在每个CIF文件中
+    cif_blocks = []
+
+    # 查找所有_chemical_formula_structural的位置
+    formula_positions = [m.start() for m in re.finditer(r'_chemical_formula_structural', content)]
+
+    # 如果没有找到任何_chemical_formula_structural，返回空字符串
+    if not formula_positions:
+        return ''
+
+    # 分割CIF块
+    for i in range(len(formula_positions)):
+        start_pos = formula_positions[i]
+        # 如果是最后一个块，结束位置是字符串末尾
+        end_pos = formula_positions[i+1] if i < len(formula_positions)-1 else len(content)
+
+        cif_block = content[start_pos:end_pos].strip()
+
+        # 提取formula值
+        formula_match = re.search(r'_chemical_formula_structural\s+(\S+)', cif_block)
+        if formula_match:
+            formula = formula_match.group(1)
+            cif_blocks.append((formula, cif_block))
+
+    # 格式化输出
+    result = []
+    for i, (formula, cif_content) in enumerate(cif_blocks, 1):
+        formatted = f"[cif {i} begin]\ndata_{formula}\n{cif_content}\n[cif {i} end]\n"
+        result.append(formatted)
+
+    return "\n".join(result)
+
+def extract_cif_file_from_zip(cifs_zip_path: str):
+    """
+    Extract CIF files from a zip archive, extract formula from each CIF file,
+    and save each CIF file with its formula as the filename.
+
+    Args:
+        cifs_zip_path: Path to the zip file
+
+    Returns:
+        list: List of tuples containing (index, formula, cif_path)
+    """
+    result_dict = {}
+    if os.path.exists(cifs_zip_path):
+            with open(cifs_zip_path, 'rb') as f:
+                result_dict['cif_content'] = f.read()
+    cifs_content = format_cif_content(result_dict.get('cif_content', b'').decode('utf-8', errors='replace') if isinstance(result_dict.get('cif_content', b''), bytes) else str(result_dict.get('cif_content', '')))
+    pattern = r'\[cif (\d+) begin\]\n(.*?)\n\[cif \1 end\]'
+    matches = re.findall(pattern, cifs_content, re.DOTALL)
+
+    # 处理每个匹配项，提取formula并保存CIF文件
+    saved_files = []
+    for idx, cif_content in matches:
+        # 提取data_{formula}中的formula
+        formula_match = re.search(r'data_([^\s]+)', cif_content)
+        if formula_match:
+            formula = formula_match.group(1)
+            # 构建保存路径
+            cif_path = os.path.join(material_config.TEMP_ROOT, f"{formula}.cif")
+            # 保存CIF文件
+            with open(cif_path, 'w') as f:
+                f.write(cif_content)
+            saved_files.append((idx, formula, cif_path))
+
+    return saved_files
+
+
+class MatterGenService:
+    """
+    Service for generating crystal structures using MatterGen.
+
+    This service initializes the CrystalGenerator once and reuses it for multiple
+    generation requests, improving performance.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+
+    # 模型到GPU ID的映射
+    MODEL_TO_GPU = {
+        "mattergen_base": "0",                          # 基础模型使用GPU 0
+        "dft_mag_density": "1",                         # 磁密度模型使用GPU 1
+        "dft_bulk_modulus": "2",                        # 体积模量模型使用GPU 2
+        "dft_shear_modulus": "3",                       # 剪切模量模型使用GPU 3
+        "energy_above_hull": "4",                       # 能量模型使用GPU 4
+        "formation_energy_per_atom": "5",               # 形成能模型使用GPU 5
+        "space_group": "6",                             # 空间群模型使用GPU 6
+        "hhi_score": "7",                               # HHI评分模型使用GPU 7
+        "ml_bulk_modulus": "0",                         # ML体积模量模型使用GPU 0
+        "chemical_system": "1",                         # 化学系统模型使用GPU 1
+        "dft_band_gap": "2",                            # 带隙模型使用GPU 2
+        "dft_mag_density_hhi_score": "3",               # 多属性模型使用GPU 3
+        "chemical_system_energy_above_hull": "4"        # 多属性模型使用GPU 4
+    }
+
+    @classmethod
+    def get_instance(cls):
+        """
+        Get the singleton instance of MatterGenService.
+
+        Returns:
+            MatterGenService: The singleton instance.
+        """
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        """
+        Initialize the MatterGenService.
+
+        This initializes the base generator without any property conditioning.
+        Specific generators for different property conditions will be initialized
+        on demand.
+        """
+        self._generators = {}
+        self._output_dir = material_config.TEMP_ROOT
+
+        # 确保输出目录存在
+        if not os.path.exists(self._output_dir):
+            os.makedirs(self._output_dir)
+
+        # 初始化基础生成器（无条件生成）
+        self._init_base_generator()
+
+    def _init_base_generator(self):
+        """
+        Initialize the base generator for unconditional generation.
+        """
+        model_path = os.path.join(material_config.MATTERGENMODEL_ROOT, "mattergen_base")
+
+        if not os.path.exists(model_path):
+            logger.warning(f"Base model directory not found at {model_path}. MatterGen service may not work properly.")
+            return
+
+        logger.info(f"Initializing base MatterGen generator from {model_path}")
+
+        try:
+            checkpoint_info = MatterGenCheckpointInfo(
+                model_path=Path(model_path).resolve(),
+                load_epoch="last",
+                config_overrides=[],
+                strict_checkpoint_loading=True,
+            )
+
+            generator = CrystalGenerator(
+                checkpoint_info=checkpoint_info,
+                properties_to_condition_on=None,
+                batch_size=2,  # 默认值，可在生成时覆盖
+                num_batches=1,  # 默认值，可在生成时覆盖
+                sampling_config_name="default",
+                sampling_config_path=None,
+                sampling_config_overrides=[],
+                record_trajectories=True,
+                diffusion_guidance_factor=0.0,
+                target_compositions_dict=[],
+            )
+
+            self._generators["base"] = generator
+            logger.info("Base MatterGen generator initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize base MatterGen generator: {e}")
+
+    def _get_or_create_generator(
+        self,
+        properties: Optional[Dict[str, Any]] = None,
+        batch_size: int = 2,
+        num_batches: int = 1,
+        diffusion_guidance_factor: float = 2.0
+    ):
+        """
+        Get or create a generator for the specified properties.
+
+        Args:
+            properties: Optional property constraints
+            batch_size: Number of structures per batch
+            num_batches: Number of batches to generate
+            diffusion_guidance_factor: Controls adherence to target properties
+
+        Returns:
+            tuple: (generator, generator_key, properties_to_condition_on, gpu_id)
+        """
+        # 如果没有属性约束，使用基础生成器
+        if not properties:
+            if "base" not in self._generators:
+                self._init_base_generator()
+            gpu_id = self.MODEL_TO_GPU.get("mattergen_base", "0")  # 默认使用GPU 0
+            return self._generators.get("base"), "base", None, gpu_id
+
+        # 处理属性约束
+        properties_to_condition_on = {}
+        for property_name, property_value in properties.items():
+            properties_to_condition_on[property_name] = property_value
+
+        # 确定模型目录
+        if len(properties) == 1:
+            # 单属性条件
+            property_name = list(properties.keys())[0]
+            property_to_model = {
+                "dft_mag_density": "dft_mag_density",
+                "dft_bulk_modulus": "dft_bulk_modulus",
+                "dft_shear_modulus": "dft_shear_modulus",
+                "energy_above_hull": "energy_above_hull",
+                "formation_energy_per_atom": "formation_energy_per_atom",
+                "space_group": "space_group",
+                "hhi_score": "hhi_score",
+                "ml_bulk_modulus": "ml_bulk_modulus",
+                "chemical_system": "chemical_system",
+                "dft_band_gap": "dft_band_gap"
+            }
+            model_dir = property_to_model.get(property_name, property_name)
+            generator_key = f"single_{property_name}"
+        else:
+            # 多属性条件
+            property_keys = set(properties.keys())
+            if property_keys == {"dft_mag_density", "hhi_score"}:
+                model_dir = "dft_mag_density_hhi_score"
+                generator_key = "multi_dft_mag_density_hhi_score"
+            elif property_keys == {"chemical_system", "energy_above_hull"}:
+                model_dir = "chemical_system_energy_above_hull"
+                generator_key = "multi_chemical_system_energy_above_hull"
+            else:
+                # 如果没有特定的多属性模型，使用第一个属性的模型
+                first_property = list(properties.keys())[0]
+                model_dir = first_property
+                generator_key = f"multi_{first_property}_etc"
+
+        # 获取对应的GPU ID
+        gpu_id = self.MODEL_TO_GPU.get(model_dir, "0")  # 默认使用GPU 0
+
+        # 构建完整的模型路径
+        model_path = os.path.join(material_config.MATTERGENMODEL_ROOT, model_dir)
+
+        # 检查模型目录是否存在
+        if not os.path.exists(model_path):
+            # 如果特定模型不存在，回退到基础模型
+            logger.warning(f"Model directory for {model_dir} not found. Using base model instead.")
+            model_path = os.path.join(material_config.MATTERGENMODEL_ROOT, "mattergen_base")
+            generator_key = "base"
+
+        # 检查是否已经有这个生成器
+        if generator_key in self._generators:
+            # 更新生成器的参数
+            generator = self._generators[generator_key]
+            generator.batch_size = batch_size
+            generator.num_batches = num_batches
+            generator.diffusion_guidance_factor = diffusion_guidance_factor if properties else 0.0
+            return generator, generator_key, properties_to_condition_on, gpu_id
+
+        # 创建新的生成器
+        try:
+            logger.info(f"Initializing new MatterGen generator for {generator_key} from {model_path}")
+
+            checkpoint_info = MatterGenCheckpointInfo(
+                model_path=Path(model_path).resolve(),
+                load_epoch="last",
+                config_overrides=[],
+                strict_checkpoint_loading=True,
+            )
+
+            generator = CrystalGenerator(
+                checkpoint_info=checkpoint_info,
+                properties_to_condition_on=properties_to_condition_on,
+                batch_size=batch_size,
+                num_batches=num_batches,
+                sampling_config_name="default",
+                sampling_config_path=None,
+                sampling_config_overrides=[],
+                record_trajectories=True,
+                diffusion_guidance_factor=diffusion_guidance_factor if properties else 0.0,
+                target_compositions_dict=[],
+            )
+
+            self._generators[generator_key] = generator
+            logger.info(f"MatterGen generator for {generator_key} initialized successfully")
+            return generator, generator_key, properties_to_condition_on, gpu_id
+        except Exception as e:
+            logger.error(f"Failed to initialize MatterGen generator for {generator_key}: {e}")
+            # 回退到基础生成器
+            if "base" not in self._generators:
+                self._init_base_generator()
+            base_gpu_id = self.MODEL_TO_GPU.get("mattergen_base", "0")
+            return self._generators.get("base"), "base", None, base_gpu_id
+
+    def generate(
+        self,
+        properties: Optional[Dict[str, Union[float, str, Dict[str, Union[float, str]]]]] = None,
+        batch_size: int = 2,
+        num_batches: int = 1,
+        diffusion_guidance_factor: float = 2.0
+    ) -> str:
+        """
+        Generate crystal structures with optional property constraints.
+
+        Args:
+            properties: Optional property constraints
+            batch_size: Number of structures per batch
+            num_batches: Number of batches to generate
+            diffusion_guidance_factor: Controls adherence to target properties
+
+        Returns:
+            str: Descriptive text with generated crystal structures in CIF format
+        """
+
+
+        # 处理字符串输入（如果提供）
+        if isinstance(properties, str):
+            try:
+                properties = json.loads(properties)
+            except json.JSONDecodeError:
+                raise ValueError(f"Invalid properties JSON string: {properties}")
+
+        # 如果为None，默认为空字典
+        properties = properties or {}
+
+        # 获取或创建生成器和GPU ID
+        generator, generator_key, properties_to_condition_on, gpu_id = self._get_or_create_generator(
+            properties, batch_size, num_batches, diffusion_guidance_factor
+        )
+        print("gpu_id",gpu_id)
+        if generator is None:
+            return "Error: Failed to initialize MatterGen generator"
+
+        # 使用torch.cuda.set_device()直接设置当前GPU
+        try:
+            # 将字符串类型的gpu_id转换为整数
+            cuda_device_id = int(gpu_id)
+            torch.cuda.set_device(cuda_device_id)
+            logger.info(f"Setting CUDA device to GPU {cuda_device_id} for model {generator_key}")
+            print(f"Using GPU {cuda_device_id} (CUDA device index) for model {generator_key}")
+        except Exception as e:
+            logger.warning(f"Error setting CUDA device: {e}. Falling back to default device.")
+
+        # 生成结构
+        try:
+
+            output_dir= Path(self._output_dir+f'/{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}')
+            Path.mkdir(output_dir, parents=True, exist_ok=True)
+            generator.generate(output_dir=output_dir)
+        except Exception as e:
+            logger.error(f"Error generating structures: {e}")
+            return f"Error generating structures: {e}"
+
+        # 创建字典存储文件内容
+        result_dict = {}
+
+        # 定义文件路径
+        cif_zip_path = os.path.join(str(output_dir), f"generated_crystals_cif.zip")
+        xyz_file_path = os.path.join(str(output_dir), f"generated_crystals.extxyz")
+        trajectories_zip_path = os.path.join(str(output_dir), f"generated_trajectories.zip")
+
+        # 读取CIF压缩文件
+        if os.path.exists(cif_zip_path):
+            with open(cif_zip_path, 'rb') as f:
+                result_dict['cif_content'] = f.read()
+
+        # 根据生成类型创建描述性提示
+        if not properties:
+            generation_type = "unconditional"
+            title = "Generated Material Structures"
+            description = "These structures were generated unconditionally, meaning no specific properties were targeted."
+            property_description = "unconditionally"
+        elif len(properties) == 1:
+            generation_type = "single_property"
+            property_name = list(properties.keys())[0]
+            property_value = properties[property_name]
+            title = f"Generated Material Structures Conditioned on {property_name} = {property_value}"
+            description = f"These structures were generated with property conditioning, targeting a {property_name} value of {property_value}."
+            property_description = f"conditioned on {property_name} = {property_value}"
+        else:
+            generation_type = "multi_property"
+            title = "Generated Material Structures Conditioned on Multiple Properties"
+            description = "These structures were generated with multi-property conditioning, targeting the specified property values."
+            property_description = f"conditioned on multiple properties: {', '.join([f'{name} = {value}' for name, value in properties.items()])}"
+
+        # 创建完整的提示
+        prompt = f"""
+# {title}
+
+This data contains {batch_size * num_batches} crystal structures generated by the MatterGen model, {property_description}.
+
+    {'' if generation_type == 'unconditional' else f'''
+    A diffusion guidance factor of {diffusion_guidance_factor} was used, which controls how strongly
+    the generation adheres to the specified property values. Higher values produce samples that more
+    closely match the target properties but may reduce diversity.
+    '''}
+
+## CIF Files (Crystallographic Information Files)
+
+- Standard format for crystallographic structures
+- Contains unit cell parameters, atomic positions, and symmetry information
+- Used by crystallographic software and visualization tools
+
+```
+{format_cif_content(result_dict.get('cif_content', b'').decode('utf-8', errors='replace') if isinstance(result_dict.get('cif_content', b''), bytes) else str(result_dict.get('cif_content', '')))}
+```
+
+{description}
+You can use these structures for materials discovery, property prediction, or further analysis.
+"""
+        # print("prompt",prompt)
+        # 清理文件（读取后删除）
+        # try:
+        #     if os.path.exists(cif_zip_path):
+        #         os.remove(cif_zip_path)
+        #     if os.path.exists(xyz_file_path):
+        #         os.remove(xyz_file_path)
+        #     if os.path.exists(trajectories_zip_path):
+        #         os.remove(trajectories_zip_path)
+        # except Exception as e:
+        #     logger.warning(f"Error cleaning up files: {e}")
+
+        # GPU设备已经在生成前由torch.cuda.set_device()设置，不需要额外清理
+        logger.info(f"Generation completed on GPU for model {generator_key}")
+
+        return prompt

sci_mcp/material_mcp/mattergen_gen/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 ...core.config import material_config
+# Add the mattergen directory to the Python path
+mattergen_dir = material_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']