Update app.py

app.py

@@ -9,13 +9,23 @@ from pymatgen.io.ase import AseAtomsAdaptor
 from ase.io import write as ase_write
 import tempfile
 import time
+# 设置PyTorch使用的线程数
+torch.set_num_threads(2)
+def load_quantized_model(model_path):
+    model = MatterGPTWrapper.from_pretrained(model_path)
+    model.to('cpu')
+    model.eval()
+    quantized_model = torch.quantization.quantize_dynamic(
+        model, {torch.nn.Linear}, dtype=torch.qint8
+    )
+    return quantized_model
 
 
-# Load the model
+# Load and quantize the model
 model_path = "./"
-model = MatterGPTWrapper.from_pretrained(model_path)
-model.to("cpu")
-model.eval()
+quantized_model = load_quantized_model(model_path)
+quantized_model.to("cpu")
+quantized_model.eval()
 # Load the tokenizer
 tokenizer_path = "Voc_prior"
 tokenizer = SimpleTokenizer(tokenizer_path)
@@ -27,22 +37,23 @@ try:
 except Exception as e:
     backend = SLICES(relax_model=None)
 
-def generate_slices_optimized(model, tokenizer, formation_energy, band_gap, max_length, temperature, do_sample, top_k, top_p):
+
+
+def generate_slices_quantized(quantized_model, tokenizer, formation_energy, band_gap, max_length, temperature, do_sample, top_k, top_p):
     condition = torch.tensor([[float(formation_energy), float(band_gap)]], dtype=torch.float32)
     context = '>'
     x = torch.tensor([[tokenizer.stoi[context]]], dtype=torch.long)
     
     with torch.no_grad():
-        generated = model.generate(x, prop=condition, max_length=max_length, 
-                                   temperature=temperature, do_sample=do_sample, 
-                                   top_k=top_k, top_p=top_p)
+        generated = quantized_model.generate(x, prop=condition, max_length=max_length, 
+                                             temperature=temperature, do_sample=do_sample, 
+                                             top_k=top_k, top_p=top_p)
     
     return tokenizer.decode(generated[0].tolist())
 
 def generate_slices(formation_energy, band_gap):
-    return generate_slices_optimized(model, tokenizer, formation_energy, band_gap, 
-                                     model.config.block_size, 1.2, True, 0, 0.9)
-
+    return generate_slices_quantized(quantized_model, tokenizer, formation_energy, band_gap, 
+                                     quantized_model.config.block_size, 1.2, True, 0, 0.9)
 def wrap_structure(structure):
     """Wrap all atoms back into the unit cell."""
     for i, site in enumerate(structure):
@@ -109,7 +120,7 @@ with gr.Blocks() as iface:
     with gr.Row():
         with gr.Column():
             gr.Image("Figure1.png", label="De novo crystal generation by MatterGPT targeting desired Eg, Ef", width=1000, height=300)
-            gr.Markdown("**Enter desired properties to inversely design materials (encoded in SLICES), then decode it into crystal structure.**")
+            gr.Markdown("**Enter desired properties to inversely design materials (encoded in SLICES), then decode it into crystal structure. Take 1-2 mins to finish with 2 cpus**")
     
     with gr.Row():
         with gr.Column(scale=2):