Convert XYZ to PDB for compatibility with gradio_molecule3d viewer

app.py

@@ -175,12 +175,41 @@ MOL_REPS = [
     {"model": 0, "style": "sphere", "color": "whiteCarbon", "opacity": 0.3},
 ]
 
-def save_xyz_to_file(xyz_str, idx):
-    """Save XYZ string to a temp file and return the path."""
+def xyz_to_pdb(xyz_str):
+    """Convert XYZ string to PDB format for the viewer."""
+    lines = xyz_str.strip().splitlines()
+    atoms = lines[2:]  # Skip header
+    pdb_lines = []
+    
+    for i, line in enumerate(atoms):
+        parts = line.split()
+        if len(parts) >= 4:
+            elem = parts[0]
+            x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
+            # Simple HETATM record
+            # Col 1-6: Record name "HETATM"
+            # Col 7-11: Serial number
+            # Col 13-16: Atom name
+            # Col 18-20: Residue name "UNK"
+            # Col 23-26: Residue sequence number "1"
+            # Col 31-38: X coord
+            # Col 39-46: Y coord
+            # Col 47-54: Z coord
+            # Col 55-60: Occupancy (1.00)
+            # Col 61-66: Temp factor (0.00)
+            # Col 77-78: Element symbol
+            pdb_lines.append(
+                f"HETATM{i+1:>5}  {elem:<3} UNK     1    {x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           {elem:>2}"
+            )
+    return "\n".join(pdb_lines)
+
+def save_pdb_to_file(xyz_str, idx):
+    """Save converted PDB string to a temp file and return the path."""
+    pdb_str = xyz_to_pdb(xyz_str)
     tmp = tempfile.NamedTemporaryFile(
-        suffix=".xyz", prefix=f"mol_{idx:03d}_", delete=False, mode="w"
+        suffix=".pdb", prefix=f"mol_{idx:03d}_", delete=False, mode="w"
     )
-    tmp.write(xyz_str)
+    tmp.write(pdb_str)
     tmp.close()
     return tmp.name
 
@@ -242,10 +271,10 @@ 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 XYZ files for the 3D viewer
+        # 5. Output generation – save PDB files for the viewer (XYZ for download)
         file_paths = []
         for i, xyz_str in enumerate(xyz_strings):
-            file_paths.append(save_xyz_to_file(xyz_str, i))
+            file_paths.append(save_pdb_to_file(xyz_str, i))
         
         zip_path = create_xyz_zip(xyz_strings)
         df_out = pd.DataFrame(summary_rows)