Upload 2 files

app.py

@@ -1,13 +1,38 @@
 import gradio as gr # type: ignore
 import os
+from gradio_molecule3d import Molecule3D
 
 # Import your custom modules from the /scripts folder
 from scripts.download import download_and_clean_pdb
 from scripts.generator import run_broteinshake_generator
 from scripts.refine import polish_design
+from scripts.visualize import create_design_plot
 
 # --- HELPER FUNCTIONS ---
 
+def get_all_sequences(fasta_file: str) -> str:
+    """Get all designed sequences from FASTA file."""
+    sequences = []
+    with open(fasta_file, 'r') as f:
+        lines = [line.strip() for line in f.readlines() if line.strip()]
+    
+    for i in range(0, len(lines), 2):
+        if i + 1 >= len(lines):
+            break
+        header = lines[i]
+        sequence = lines[i+1]
+        
+        # Skip the original native sequence (first entry)
+        if "sample" not in header:
+            continue
+        
+        sequences.append(f"{header}\n{sequence}")
+    
+    if sequences:
+        return "\n\n".join(sequences)
+    else:
+        raise ValueError(f"No valid designs found in {fasta_file}")
+
 def extract_best_sequence(fasta_file: str) -> str:
     """Extract the best sequence (lowest score) from FASTA file."""
     best_score = float('inf')
@@ -54,17 +79,32 @@ def run_part1(pdb_id, fixed_chains, variable_chains):
         # This creates the .fa files you need for the ESM Atlas
         run_broteinshake_generator(pdb_path, fixed_chains, variable_chains)
         
-        # Extract and display only the best sequence
+        # Get all sequences and the best one
         fa_file = os.path.join("generated", pdb_id.lower(), "seqs", f"{pdb_id.lower()}_clones.fa")
+        all_sequences = get_all_sequences(fa_file)
         best_sequence = extract_best_sequence(fa_file)
         
+        # Generate the dashboard plot
+        evolution_plot = create_design_plot(fa_file)
+        
         # Parse score from header for status message
         score_part = [p for p in best_sequence.split('\n')[0].split(',') if 'score' in p][0]
         best_score = float(score_part.split('=')[1])
+        
+        # Count number of designs
+        num_designs = len([s for s in all_sequences.split('\n\n') if s.strip()])
+        
+        # Format status with best sequence
+        status_message = (
+            f"✅ Design Complete! {num_designs} designs generated.\n\n"
+            f"⭐ Lead Candidate (Best Score: {best_score:.4f}):\n"
+            f"{best_sequence}\n\n"
+            f"📋 Copy the lead sequence above and fold it at [esmatlas.com](https://esmatlas.com/resources?action=fold)."
+        )
             
-        return best_sequence, f"Part 1 Complete. Best design (score={best_score:.4f}) shown above. Copy the sequence and fold it at esmatlas.com."
+        return all_sequences, evolution_plot, status_message
     except Exception as e:
-        return "", f"Error in Part 1: {str(e)}"
+        return "", None, f"❌ Error in Part 1: {str(e)}"
 
 def run_part2(pdb_id, uploaded_esm_file):
     """Aligns the ESM-folded PDB to the target structure."""
@@ -72,8 +112,27 @@ def run_part2(pdb_id, uploaded_esm_file):
         if uploaded_esm_file is None:
             return None, "⚠️ Please upload the PDB from ESM Atlas first."
         
+        # Validate inputs
+        if not pdb_id or not pdb_id.strip():
+            return None, "❌ Error: PDB ID is required."
+        
+        # Get file path
+        file_path = uploaded_esm_file.name if hasattr(uploaded_esm_file, 'name') else uploaded_esm_file
+        if not file_path or not os.path.exists(file_path):
+            return None, f"❌ Error: Uploaded file not found: {file_path}"
+        
         # Call the new lightweight Biopython-only script
-        final_pdb_path, rmsd_val = polish_design(pdb_id, uploaded_esm_file.name)
+        print(f"🔍 Starting alignment for PDB ID: {pdb_id}, File: {file_path}")
+        final_pdb_path, rmsd_val = polish_design(pdb_id, file_path)
+        
+        # Validate output
+        if not final_pdb_path or not os.path.exists(final_pdb_path):
+            return None, f"❌ Error: Alignment failed - output file not created: {final_pdb_path}"
+        
+        if rmsd_val is None:
+            return None, "❌ Error: Alignment failed - RMSD calculation returned None"
+        
+        print(f"✅ Alignment successful: RMSD = {rmsd_val:.3f} Å, Output: {final_pdb_path}")
         
         # Update the report to focus on Backbone Alignment
         report = (
@@ -82,12 +141,61 @@ def run_part2(pdb_id, uploaded_esm_file):
             f"🧬 Status: High-Precision Backbone Match"
         )
         return final_pdb_path, report
+    except FileNotFoundError as e:
+        error_msg = f"❌ File Error: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+    except ValueError as e:
+        error_msg = f"❌ Validation Error: {str(e)}"
+        print(error_msg)
+        return None, error_msg
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        error_msg = f"❌ Unexpected Error: {str(e)}\n\nTraceback:\n{type(e).__name__}"
+        print(error_msg)
+        import traceback
+        traceback.print_exc()
+        return None, error_msg
 
 # --- GRADIO INTERFACE ---
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# 1. Define the Biohub Theme
+biohub_theme = gr.themes.Soft(
+    primary_hue="emerald",
+    secondary_hue="slate",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
+).set(
+    # Custom CSS variable overrides for that "Glass" look
+    block_title_text_weight="600",
+    block_label_text_size="*text_sm",
+    input_background_fill="*neutral_50",
+    button_primary_background_fill="linear-gradient(90deg, #10b981, #34d399)",
+    button_primary_text_color="white",
+    border_color_primary="*neutral_200",
+)
+
+# 2. Inject Custom CSS for the "Spice"
+biohub_css = """
+.gradio-container {
+    background: radial-gradient(circle at top right, #f8fafc, #f1f5f9); /* Subtle depth */
+}
+#biohub-header {
+    text-align: center;
+    padding: 20px;
+    background: #064e3b;
+    color: white;
+    border-radius: 15px;
+    margin-bottom: 20px;
+    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
+}
+.gr-box {
+    border-radius: 12px !important;
+    border: 1px solid #e2e8f0 !important;
+}
+footer {display: none !important;} /* Cleans up the portfolio look */
+"""
+
+with gr.Blocks(theme=biohub_theme, css=biohub_css) as demo:
     gr.Markdown("# broteinshake")
     gr.Markdown("### AI-Driven Structural Mimicry & BBB-Shuttle Design")
     
@@ -100,24 +208,77 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 f_chains = gr.Textbox(label="Fixed Chains (Lock)", value="A")
                 v_chains = gr.Textbox(label="Variable Chains (Key)", value="B")
             
-            gen_btn = gr.Button("Generate Optimized Sequences", variant="primary")
-            fa_output = gr.Code(label="Best Design Sequence (.fasta format)", language="markdown")
-            status1 = gr.Textbox(label="System Status", interactive=False)
+            gen_btn = gr.Button("🚀 Generate Optimized Sequences", variant="primary")
+            
+            # Stack components vertically
+            fa_output = gr.Code(
+                label="Designed Sequences", 
+                language="markdown",
+                lines=10,
+                max_lines=20  # Prevents infinite growth when showing all 20 candidates
+            )
+            plot_output = gr.Plot(
+                label="Design Evolution Dashboard"
+            )
+            
+            gr.Markdown("### System Status")
+            status1 = gr.Markdown()
             
-            gen_btn.click(run_part1, inputs=[pdb_input, f_chains, v_chains], outputs=[fa_output, status1])
+            gen_btn.click(run_part1, inputs=[pdb_input, f_chains, v_chains], 
+                         outputs=[fa_output, plot_output, status1])
 
         # TAB 2: STRUCTURAL VALIDATION
         with gr.Tab("2. Structural Validation"):
-            gr.Markdown("### Align Designed Shuttle to Human Receptor")
+            gr.Markdown("### 🧬 Final Structure Preview")
+            
+            # Updated REPS for local dev visibility
+            REPS = [
+                {
+                    "model": 0,
+                    "style": "cartoon",
+                    "color": "spectrum",
+                    "opacity": 1.0
+                    # Removing "chain": "A" is the key fix here!
+                }
+            ]
+            
+            # This component ACTUALLY supports .pdb natively
+            protein_view = Molecule3D(label="3D Structure Viewer (Refined Shuttle)", reps=REPS)
+            
             with gr.Row():
                 esm_upload = gr.File(label="Upload ESM-Folded PDB", file_types=[".pdb"])
                 refined_download = gr.File(label="Download Aligned Lead (.pdb)")
             
-            # Change the button text to match our new approach
             validate_btn = gr.Button("✨ Run Structural Alignment", variant="primary")
             status2 = gr.Textbox(label="Validation Report", interactive=False, lines=5)
-            
-            validate_btn.click(run_part2, inputs=[pdb_input, esm_upload], outputs=[refined_download, status2])
+
+            # Update viewer - REPS are set on component creation, just return the file path
+            def run_validation_with_view(pdb_id, file):
+                try:
+                    # This is your existing alignment/refine logic
+                    final_pdb, report = run_part2(pdb_id, file)
+                    
+                    if final_pdb is not None and os.path.exists(final_pdb):
+                        # Verify we're using the refined shuttle (Refined_Shuttle.pdb)
+                        if "Refined_Shuttle.pdb" in final_pdb or os.path.basename(final_pdb) == "Refined_Shuttle.pdb":
+                            print(f"🎯 Visualizing refined shuttle: {final_pdb}")
+                        else:
+                            print(f"⚠️ Warning: Expected Refined_Shuttle.pdb but got: {final_pdb}")
+                    
+                    # Molecule3D expects just the file path string (REPS are already set on component)
+                    return final_pdb, report, final_pdb
+                except Exception as e:
+                    error_msg = f"❌ Error generating 3D view: {str(e)}"
+                    print(error_msg)
+                    import traceback
+                    traceback.print_exc()
+                    return None, error_msg, None
+
+            validate_btn.click(
+                run_validation_with_view, 
+                inputs=[pdb_input, esm_upload], 
+                outputs=[refined_download, status2, protein_view]
+            )
 
 # Launch the app
 if __name__ == "__main__":

requirements.txt

@@ -1,4 +1,5 @@
-gradio
+gradio==5.50.0
+gradio_molecule3d
 accelerate==1.12.0
 backports.zstd
 biopython==1.86
@@ -19,7 +20,6 @@ MarkupSafe
 mpmath
 networkx
 numpy==2.2.6
-openmm
 packaging==25.0
 pillow
 psutil==7.2.1