Upload app.py with huggingface_hub

app.py

@@ -81,7 +81,6 @@ def save_interesting_log(smiles, properties, suggested_properties):
         }
         writer.writerow(log_data)
 
-@spaces.GPU
 def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
     model, device = model_state
     
@@ -99,61 +98,54 @@ def generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_ti
     num_nodes = None if num_nodes == 0 else num_nodes
     
     for _ in range(repeating_time):
-        try:
-            # def generate_func():
-                # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-                # model.to(device)
-                # print('Before generation, move model to', device)
-                # return generated_molecule, img_list
-            # generated_molecule, img_list = generate_func()
-
-            model.to(device)
-            generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
-
-            # Create GIF if img_list is available
-            gif_path = None
-            if img_list and len(img_list) > 0:
-                imgs = [np.array(pil_img) for pil_img in img_list]
-                imgs.extend([imgs[-1]] * 10)
-                gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
-                imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
-            
-            if generated_molecule is not None:
-                mol = Chem.MolFromSmiles(generated_molecule)
-                if mol is not None:
-                    standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
-                    is_novel = standardized_smiles not in knwon_smiles['SMILES'].values
-                    novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
-                    img = Draw.MolToImage(mol)
-                    
-                    # Evaluate the generated molecule
-                    suggested_properties = {}
-                    for prop, evaluator in evaluators.items():
-                        suggested_properties[prop] = evaluator([standardized_smiles])[0]
-                    
-                    suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
-                    
-                    return (
-                        f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
-                        f"**{nan_message}**\n\n"
-                        f"**{novelty_status}**\n\n"
-                        f"**Suggested Properties:**\n{suggested_properties_text}",
-                        img,
-                        gif_path,
-                        properties,  # Add this
-                        suggested_properties  # Add this
-                    )
-            else:
+        # try:
+        model.to(device)
+        generated_molecule, img_list = model.generate(properties, device=device, guide_scale=guidance_scale, num_nodes=num_nodes, number_chain_steps=num_chain_steps)
+
+        # Create GIF if img_list is available
+        gif_path = None
+        if img_list and len(img_list) > 0:
+            imgs = [np.array(pil_img) for pil_img in img_list]
+            imgs.extend([imgs[-1]] * 10)
+            gif_path = os.path.join(temp_dir, f"polymer_gen_{random.randint(0, 999999)}.gif")
+            imageio.mimsave(gif_path, imgs, format='GIF', fps=fps, loop=0)
+        
+        if generated_molecule is not None:
+            mol = Chem.MolFromSmiles(generated_molecule)
+            if mol is not None:
+                standardized_smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
+                is_novel = standardized_smiles not in knwon_smiles['SMILES'].values
+                novelty_status = "Novel (Not in Labeled Set)" if is_novel else "Not Novel (Exists in Labeled Set)"
+                img = Draw.MolToImage(mol)
+                
+                # Evaluate the generated molecule
+                suggested_properties = {}
+                for prop, evaluator in evaluators.items():
+                    suggested_properties[prop] = evaluator([standardized_smiles])[0]
+                
+                suggested_properties_text = "\n".join([f"**Suggested {prop}:** {value:.2f}" for prop, value in suggested_properties.items()])
+                
                 return (
-                    f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
-                    None,
+                    f"**Generated polymer SMILES:** `{standardized_smiles}`\n\n"
+                    f"**{nan_message}**\n\n"
+                    f"**{novelty_status}**\n\n"
+                    f"**Suggested Properties:**\n{suggested_properties_text}",
+                    img,
                     gif_path,
-                    properties,
-                    None,
+                    properties,  # Add this
+                    suggested_properties  # Add this
                 )
-        except Exception as e:
-            print(f"Error in generation: {e}")
-            continue
+        else:
+            return (
+                f"**Generation failed:** Could not generate a valid molecule.\n\n**{nan_message}**",
+                None,
+                gif_path,
+                properties,
+                None,
+            )
+        # except Exception as e:
+        #     print(f"Error in generation: {e}")
+        #     continue
     
     return f"**Generation failed:** Could not generate a valid molecule after {repeating_time} attempts.\n\n**{nan_message}**", None, None
 
@@ -180,6 +172,7 @@ def numpy_to_python(obj):
     else:
         return obj
         
+@spaces.GPU
 def on_generate(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps):
     result = generate_graph(CH4, CO2, H2, N2, O2, guidance_scale, num_nodes, repeating_time, model_state, num_chain_steps, fps)
     # Check if the generation was successful
@@ -276,7 +269,7 @@ with gr.Blocks(title="Polymer Design with GraphDiT") as iface:
         ```
         """)
 
-    model_state = gr.State(lambda: load_model("model_all"))
+    model_state = gr.State(lambda: load_model("model_labeled"))
 
     with gr.Row():
         CH4_input = gr.Slider(0, property_ranges['CH4'][1], value=2.5, label=f"CH₄ (Barrier) [0-{property_ranges['CH4'][1]:.1f}]")