Replace gradio_molecule3d with direct 3Dmol.js: full style/bg/label/H control

app.py

@@ -183,18 +183,7 @@ def create_xyz_zip(xyz_strings):
             
     return tmp_zip.name
 
-from gradio_molecule3d import Molecule3D
-
-# 3Dmol.js representation configs
-STYLE_REPS = {
-    "Ball and Stick": [
-        {"model": 0, "style": "stick", "color": "whiteCarbon", "radius": 0.2},
-        {"model": 0, "style": "sphere", "color": "whiteCarbon", "scale": 0.3}
-    ],
-    "Licorice": [
-        {"model": 0, "style": "stick", "color": "whiteCarbon", "radius": 0.4}
-    ]
-}
+# ── 3Dmol.js SETUP (loaded via head tag) ──────────────────────
 
 def xyz_to_pdb(xyz_str):
     """Convert XYZ string to PDB format for the viewer."""
@@ -288,12 +277,7 @@ def generate(num_molecules, size_mode, fixed_size, diffusion_steps, seed):
             xyz_strings.append(xyz_str)
             summary_rows.append(parse_composition(xyz_str))
             
-        # 5. Output generation – save PDB files for the viewer (XYZ for download)
-        file_paths = []
-        for i, xyz_str in enumerate(xyz_strings):
-            # Default to PDB for viewer as it's most compatible
-            file_paths.append(save_to_format(xyz_str, i, "pdb"))
-        
+        # 5. Output generation – save zip for bulk download
         zip_path = create_xyz_zip(xyz_strings)
         
         # Prepare table with "Name" column
@@ -304,21 +288,20 @@ def generate(num_molecules, size_mode, fixed_size, diffusion_steps, seed):
         cols = ["Name"] + [c for c in df_out.columns if c != "Name"]
         df_out = df_out[cols]
         
-        choices = [f"Molecule {i+1}" for i in range(len(file_paths))]
+        choices = [f"Molecule {i+1}" for i in range(len(xyz_strings))]
         
         return (
-            file_paths[0],                                    # viewer
-            gr.update(choices=choices, value=choices[0]),      # selector
-            file_paths,                                       # state (paths)
-            xyz_strings,                                      # state (raw)
-            zip_path,                                         # download
-            df_out,                                           # table
+            xyz_strings[0],                                    # current_xyz (triggers JS)
+            gr.update(choices=choices, value=choices[0]),       # selector
+            xyz_strings,                                       # raw_xyz_state
+            zip_path,                                          # download_all
+            df_out,                                            # table
         )
         
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return None, gr.update(choices=[], value=None), [], [], None, None
+        return "", gr.update(choices=[], value=None), [], None, None
         
     finally:
         # Restore original T to be safe
@@ -328,8 +311,75 @@ def generate(num_molecules, size_mode, fixed_size, diffusion_steps, seed):
              elif hasattr(TASK, 'T'):
                 TASK.T = original_T
 
+# ── 3Dmol.js JavaScript ───────────────────────────────────────
+THREEDMOL_HEAD = """
+<script src="https://3Dmol.org/build/3Dmol-min.js"></script>
+<style>
+  #mol3d-container { width: 100%; height: 480px; position: relative; border: 1px solid #ccc; border-radius: 8px; overflow: hidden; }
+</style>
+<script>
+  var _viewer = null;
+  var _currentXYZ = null;
+
+  function initViewer() {
+    var el = document.getElementById('mol3d-container');
+    if (!el) return;
+    if (_viewer) { _viewer.clear(); }
+    else { _viewer = $3Dmol.createViewer(el, {backgroundColor: '0xffffff'}); }
+  }
+
+  function loadMolecule(xyzStr, style, bg, showLabels, hideH) {
+    if (!_viewer) initViewer();
+    if (!_viewer) return;
+    _currentXYZ = xyzStr;
+    _viewer.clear();
+    if (!xyzStr || xyzStr.length < 5) return;
+    _viewer.addModel(xyzStr, 'xyz');
+    applyStyle(style, hideH);
+    if (showLabels) {
+      _viewer.addPropertyLabels('elem', {}, {fontSize: 12, fontColor: 'black', backgroundOpacity: 0.3, backgroundColor: 'white', alignment: 'center'});
+    }
+    _viewer.setBackgroundColor(bg);
+    _viewer.zoomTo();
+    _viewer.render();
+  }
+
+  function applyStyle(style, hideH) {
+    if (!_viewer) return;
+    _viewer.setStyle({}, {});
+    var sel = hideH ? {elem: ['C','N','O','S','P','F','Cl','Br','I','B','Si','Se','As','Al','Hg','Bi']} : {};
+    if (hideH) { _viewer.setStyle({elem: 'H'}, {}); }
+    switch(style) {
+      case 'Ball and Stick':
+        _viewer.setStyle(sel, {stick: {radius: 0.15, colorscheme: 'Jmol'}, sphere: {scale: 0.25, colorscheme: 'Jmol'}});
+        break;
+      case 'Licorice':
+        _viewer.setStyle(sel, {stick: {radius: 0.3, colorscheme: 'Jmol'}});
+        break;
+      case 'Sphere':
+        _viewer.setStyle(sel, {sphere: {colorscheme: 'Jmol'}});
+        break;
+      case 'Stick':
+        _viewer.setStyle(sel, {stick: {colorscheme: 'Jmol'}});
+        break;
+      default:
+        _viewer.setStyle(sel, {stick: {radius: 0.15, colorscheme: 'Jmol'}, sphere: {scale: 0.25, colorscheme: 'Jmol'}});
+    }
+    _viewer.render();
+  }
+
+  function refreshViewer(style, bg, showLabels, hideH) {
+    if (!_viewer || !_currentXYZ) return;
+    loadMolecule(_currentXYZ, style, bg, showLabels, hideH);
+  }
+</script>
+"""
+
+# Background color mapping
+BG_COLORS = {"White": "0xffffff", "Black": "0x000000", "Grey": "0x333333"}
+
 # ── GRADIO UI ──────────────────────────────────────────────────
-with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft(), head=THREEDMOL_HEAD) as demo:
     gr.Markdown(
         """
         # 🧪 MolCraftDiffusion - Unconditional Generation
@@ -369,20 +419,27 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
                 )
                 
                 seed = gr.Number(value=42, label="Random Seed", precision=0)
+            
+            with gr.Accordion("🎨 Viewer Options", open=True):
+                mol_style = gr.Radio(
+                    ["Ball and Stick", "Licorice", "Sphere", "Stick"], 
+                    value="Ball and Stick", 
+                    label="3D Style"
+                )
+                bg_color = gr.Radio(
+                    ["White", "Black", "Grey"],
+                    value="White",
+                    label="Background"
+                )
+                show_labels = gr.Checkbox(label="Show atom labels", value=False)
+                hide_h = gr.Checkbox(label="Hide hydrogens", value=False)
+                dl_format = gr.Radio(
+                    ["PDB", "XYZ"], 
+                    value="PDB", 
+                    label="Download Format"
+                )
                 
-                with gr.Accordion("🎨 Visual & Download Options", open=True):
-                    mol_style = gr.Radio(
-                        ["Ball and Stick", "Licorice"], 
-                        value="Ball and Stick", 
-                        label="3D Display Style"
-                    )
-                    dl_format = gr.Radio(
-                        ["PDB", "XYZ"], 
-                        value="PDB", 
-                        label="Individual Download Format"
-                    )
-                
-                btn = gr.Button("🚀 Generate Molecules", variant="primary", size="lg")
+            btn = gr.Button("🚀 Generate Molecules", variant="primary", size="lg")
         
         with gr.Column(scale=2):
             gr.Markdown("### 🧬 Generated Molecules")
@@ -393,8 +450,13 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
                 )
                 single_dl = gr.File(label="Download Selection", scale=2)
             
+            # 3Dmol.js viewer container
+            viewer_html = gr.HTML(
+                value='<div id="mol3d-container" style="width:100%;height:480px;"></div>',
+                label="3D Viewer"
+            )
+            
             with gr.Row():
-                viewer = Molecule3D(label="3D Viewer", reps=STYLE_REPS["Ball and Stick"], scale=2)
                 preview = gr.Image(label="2D Preview (PNG)", scale=1)
             
             with gr.Row():
@@ -402,14 +464,13 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
                 download_all = gr.File(label="Download All (.zip)")
 
     # Hidden states
-    mol_files_state = gr.State([])
-    raw_xyz_state = gr.State([])
+    raw_xyz_state = gr.State([])  # List of raw XYZ strings
+    current_xyz = gr.Textbox(visible=False)  # Current molecule XYZ for JS
 
     def mol_to_png(xyz_str):
         """Generate a 2D PNG preview using RDKit."""
         from rdkit.Chem import Draw
         try:
-            # Simple XYZ -> RDKit Mol (no connectivity info in XYZ, so we guess)
             lines = xyz_str.strip().splitlines()
             if len(lines) < 3: return None
             
@@ -422,7 +483,6 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
                     mol.AddAtom(atom)
                     pos.append([float(parts[1]), float(parts[2]), float(parts[3])])
             
-            # Simple distance-based connectivity for drawing
             adj = build_adjacency_matrix(np.array(pos), [a.GetAtomicNum() for a in mol.GetAtoms()], scale=1.2)
             for i in range(len(pos)):
                 for j in range(i+1, len(pos)):
@@ -430,11 +490,8 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
                         mol.AddBond(i, j, Chem.BondType.SINGLE)
             
             rd_mol = mol.GetMol()
-            # Compute 2D coordinates for drawing
-            from rdkit.Chem import AllChem
             AllChem.Compute2DCoords(rd_mol)
             
-            # Draw to PNG
             img = Draw.MolToImage(rd_mol, size=(400, 400))
             tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
             img.save(tmp_img.name)
@@ -443,53 +500,87 @@ with gr.Blocks(title="MolCraftDiffusion", theme=gr.themes.Soft()) as demo:
             print(f"Drawing error: {e}")
             return None
 
-    def update_view(choice, file_paths, raw_xyzs, style, fmt):
-        """Update viewer, preview and individual download based on selection/style/format."""
-        if not file_paths or not raw_xyzs or choice is None:
-            return gr.update(value=None), None, None
-        
+    def select_molecule(choice, raw_xyzs, fmt):
+        """Update download file and current XYZ when molecule selection changes."""
+        if not raw_xyzs or choice is None:
+            return None, None, ""
         try:
             idx = int(choice.split()[-1]) - 1
-            if not (0 <= idx < len(file_paths)):
-                return gr.update(value=None), None, None
+            if not (0 <= idx < len(raw_xyzs)):
+                return None, None, ""
         except (ValueError, IndexError):
-            return gr.update(value=None), None, None
-            
-        # 1. Update Viewer (always use PDB for viewer compatibility)
-        new_reps = STYLE_REPS.get(style, STYLE_REPS["Ball and Stick"])
-        viewer_update = gr.update(value=file_paths[idx], reps=new_reps)
+            return None, None, ""
         
-        # 2. Update Download File (respect user's format preference)
         dl_path = save_to_format(raw_xyzs[idx], idx, fmt)
-        
-        # 3. Generate PNG Preview
         png_path = mol_to_png(raw_xyzs[idx])
-        
-        return viewer_update, dl_path, png_path
+        return dl_path, png_path, raw_xyzs[idx]
+
+    # JS to call loadMolecule when current_xyz changes
+    LOAD_JS = """
+    (xyz, style, bg, labels, hideH) => {
+        var bgHex = {'White':'0xffffff','Black':'0x000000','Grey':'0x333333'}[bg] || '0xffffff';
+        setTimeout(() => { loadMolecule(xyz, style, bgHex, labels, hideH); }, 100);
+        return [xyz, style, bg, labels, hideH];
+    }
+    """
 
-    # Listeners
+    REFRESH_JS = """
+    (style, bg, labels, hideH) => {
+        var bgHex = {'White':'0xffffff','Black':'0x000000','Grey':'0x333333'}[bg] || '0xffffff';
+        refreshViewer(style, bgHex, labels, hideH);
+        return [style, bg, labels, hideH];
+    }
+    """
+
+    # When molecule selection changes: update download + preview, then trigger JS
     mol_selector.change(
-        fn=update_view,
-        inputs=[mol_selector, mol_files_state, raw_xyz_state, mol_style, dl_format],
-        outputs=[viewer, single_dl, preview],
-    )
-    
-    mol_style.change(
-        fn=update_view,
-        inputs=[mol_selector, mol_files_state, raw_xyz_state, mol_style, dl_format],
-        outputs=[viewer, single_dl, preview],
+        fn=select_molecule,
+        inputs=[mol_selector, raw_xyz_state, dl_format],
+        outputs=[single_dl, preview, current_xyz],
+    ).then(
+        fn=None,
+        inputs=[current_xyz, mol_style, bg_color, show_labels, hide_h],
+        outputs=None,
+        js=LOAD_JS,
     )
 
+    # When style/bg/labels/hideH changes: refresh viewer via JS only
+    for ctrl in [mol_style, bg_color, show_labels, hide_h]:
+        ctrl.change(
+            fn=None,
+            inputs=[mol_style, bg_color, show_labels, hide_h],
+            outputs=None,
+            js=REFRESH_JS,
+        )
+
+    # When format changes: just update the download file
+    def update_dl(choice, raw_xyzs, fmt):
+        if not raw_xyzs or choice is None:
+            return None
+        try:
+            idx = int(choice.split()[-1]) - 1
+            if 0 <= idx < len(raw_xyzs):
+                return save_to_format(raw_xyzs[idx], idx, fmt)
+        except (ValueError, IndexError):
+            pass
+        return None
+
     dl_format.change(
-        fn=update_view,
-        inputs=[mol_selector, mol_files_state, raw_xyz_state, mol_style, dl_format],
-        outputs=[viewer, single_dl, preview],
+        fn=update_dl,
+        inputs=[mol_selector, raw_xyz_state, dl_format],
+        outputs=[single_dl],
     )
 
+    # Generate button
     btn.click(
         fn=generate, 
         inputs=[num_mol, size_mode, fixed_size, diffusion_steps, seed],
-        outputs=[viewer, mol_selector, mol_files_state, raw_xyz_state, download_all, table]
+        outputs=[current_xyz, mol_selector, raw_xyz_state, download_all, table]
+    ).then(
+        fn=None,
+        inputs=[current_xyz, mol_style, bg_color, show_labels, hide_h],
+        outputs=None,
+        js=LOAD_JS,
     )
     
     gr.Markdown("---")

requirements.txt

@@ -41,5 +41,4 @@ psutil
 # App / Deployment
 gradio==4.42.0
 gradio_client==1.3.0
-gradio_molecule3d
 huggingface-hub<0.25.0