Update app.py

app.py

@@ -110,7 +110,7 @@ def load_models():
                     import torch.serialization
                     
                     # 添加安全全局变量
-                    torch.serialization.add_safe_globals([getattr(np.core.multiarray, 'scalar')])
+                    torch.serialization.add_safe_globals([getattr(np._core.multiarray, 'scalar')])  # 修复警告
                     state = torch.load(pth_path, map_location=device, weights_only=True)
                     logger.info("使用 weights_only=True 和安全全局变量成功加载模型")
                 except:
@@ -261,24 +261,6 @@ def predict_from_drawing(mol_data):
     # 使用预测函数
     return predict_all(smiles)
 
-# 创建绘图输入界面
-drawing_input = gr.Sketchpad(
-    label="绘制分子结构",
-    type="mol",
-    height=300,
-    brush_radius=10
-)
-
-# 创建输出组件
-outputs = [
-    gr.Text(label="Elastic"), 
-    gr.Image(type="pil", label="Elastic attention visualization"),
-    gr.Text(label="Plastic"), 
-    gr.Image(type="pil", label="Plastic attention visualization"),
-    gr.Text(label="Brittle"), 
-    gr.Image(type="pil", label="Brittle attention visualization")
-]
-
 # 使用TabbedInterface组织两种输入方式
 with gr.Blocks(title="CrystalGAT") as demo:
     gr.Markdown("# CrystalGAT")
@@ -305,6 +287,7 @@ with gr.Blocks(title="CrystalGAT") as demo:
     with gr.Tab("绘制分子"):
         gr.Markdown("### 绘制分子结构")
         gr.Markdown("使用下方的绘图工具绘制分子结构，然后点击预测按钮")
+        drawing_input = gr.MoleculeEditor(label="绘制分子结构")
         drawing_display = gr.Image(label="分子结构预览", interactive=False)
         drawing_output = gr.Text(label="生成的SMILES")
         submit_btn2 = gr.Button("预测", variant="primary")
@@ -331,7 +314,7 @@ with gr.Blocks(title="CrystalGAT") as demo:
     
     # 绘图输入路径
     drawing_input.change(
-        fn=draw_to_smiles,
+        fn=lambda mol: draw_to_smiles(mol) if mol else (None, None),
         inputs=drawing_input,
         outputs=[drawing_output, drawing_display]
     )