remove cap

app.py

@@ -16,7 +16,6 @@ from transformers.models.esmfold2.protein_utils import output_to_pdb, prepare_pr
 # reasonable on a shared ZeroGPU slot. Swap for "biohub/ESMFold2" for MSA-conditioned
 # accuracy on hard targets (slower, and needs an MSA pipeline this app doesn't provide).
 MODEL_NAME = "biohub/ESMFold2-Fast"
-MAX_TOTAL_RESIDUES = 600  # keeps inference within the ZeroGPU time/memory budget
 MAX_ENTITY_ROWS = 8
 VALID_AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWYXBZUO")
 ENTITY_TYPES = ["protein", "dna", "rna"]
@@ -118,8 +117,44 @@ def prepare_multi_chain_protein_features(sequences: list[str], device) -> dict[s
     return {k: v.to(device) for k, v in combined.items()}
 
 
+def harvest_entities(etypes: list[str], copiess: list, seqs: list[str]) -> list[dict]:
+    """Read the currently displayed value of every row's components directly.
+
+    Used instead of syncing a ``gr.State`` copy on every keystroke, so editing a field never
+    depends on a change/input event having already round-tripped to the server — the value the
+    user sees is always exactly the value Add/Remove/Fold act on.
+    """
+    return [
+        {"type": t, "copies": max(1, int(c or 1)), "sequence": s or ""}
+        for t, c, s in zip(etypes, copiess, seqs)
+    ]
+
+
+def add_entity_live(etypes: list[str], copiess: list, seqs: list[str]) -> list[dict]:
+    entities = harvest_entities(etypes, copiess, seqs)
+    if len(entities) < MAX_ENTITY_ROWS:
+        entities.append(dict(DEFAULT_ENTITY))
+    return entities
+
+
+def remove_entity_live(etypes: list[str], copiess: list, seqs: list[str], remove_idx: int) -> list[dict]:
+    entities = harvest_entities(etypes, copiess, seqs)
+    if len(entities) <= 1:
+        return entities
+    return [e for i, e in enumerate(entities) if i != remove_idx]
+
+
+def load_example(sequence: str) -> list[dict]:
+    return [{"type": "protein", "copies": 1, "sequence": sequence}]
+
+
 @spaces.GPU(duration=180)
-def fold_entities(entities: list[dict], num_loops: int, num_sampling_steps: int, progress=gr.Progress()):
+def fold_from_rows(*args, progress=gr.Progress()):
+    *row_values, num_loops, num_sampling_steps = args
+    num_rows = len(row_values) // 3
+    etypes, copiess, seqs = row_values[:num_rows], row_values[num_rows : 2 * num_rows], row_values[2 * num_rows :]
+    entities = harvest_entities(etypes, copiess, seqs)
+
     chain_sequences: list[str] = []
     for entity in entities:
         sequence = clean_sequence(entity.get("sequence", ""))
@@ -140,12 +175,6 @@ def fold_entities(entities: list[dict], num_loops: int, num_sampling_steps: int,
         raise gr.Error("Please enter at least one protein sequence.")
 
     total_residues = sum(len(s) for s in chain_sequences)
-    if total_residues > MAX_TOTAL_RESIDUES:
-        raise gr.Error(
-            f"Total length is {total_residues} residues across {len(chain_sequences)} chain(s); this "
-            f"Space is limited to {MAX_TOTAL_RESIDUES} residues to keep inference within the ZeroGPU "
-            "time budget."
-        )
 
     progress(0.1, desc="Moving model to GPU...")
     model.to("cuda")
@@ -186,28 +215,6 @@ def fold_entities(entities: list[dict], num_loops: int, num_sampling_steps: int,
     return viewer_html, "  \n".join(summary_lines), pdb_path
 
 
-def add_entity(entities: list[dict]) -> list[dict]:
-    if len(entities) >= MAX_ENTITY_ROWS:
-        return entities
-    return entities + [dict(DEFAULT_ENTITY)]
-
-
-def remove_entity(entities: list[dict], idx: int) -> list[dict]:
-    if len(entities) <= 1:
-        return entities
-    return [e for i, e in enumerate(entities) if i != idx]
-
-
-def update_entity_field(entities: list[dict], idx: int, field: str, value) -> list[dict]:
-    entities = [dict(e) for e in entities]
-    entities[idx][field] = value
-    return entities
-
-
-def load_example(sequence: str) -> list[dict]:
-    return [{"type": "protein", "copies": 1, "sequence": sequence}]
-
-
 with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
     gr.Markdown(
         f"""
@@ -226,48 +233,10 @@ with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
                 example1_btn = gr.Button("Load: Ubiquitin", size="sm")
                 example2_btn = gr.Button("Load: CAP protein", size="sm")
 
-            @gr.render(inputs=entities_state)
-            def render_entities(entities):
-                for idx, entity in enumerate(entities):
-                    with gr.Row():
-                        etype = gr.Dropdown(ENTITY_TYPES, value=entity["type"], label="Type", scale=2)
-                        copies = gr.Number(
-                            value=entity["copies"], label="Copies", precision=0, minimum=1, scale=1
-                        )
-                        seq = gr.Textbox(
-                            value=entity["sequence"],
-                            label="Sequence",
-                            lines=4,
-                            max_lines=4,
-                            scale=6,
-                            placeholder="Amino acid sequence...",
-                        )
-                        remove_btn = gr.Button("×", scale=0, min_width=40, visible=len(entities) > 1)
-
-                    etype.change(
-                        lambda val, ents, i=idx: update_entity_field(ents, i, "type", val),
-                        inputs=[etype, entities_state],
-                        outputs=entities_state,
-                    )
-                    copies.input(
-                        lambda val, ents, i=idx: update_entity_field(ents, i, "copies", int(val or 1)),
-                        inputs=[copies, entities_state],
-                        outputs=entities_state,
-                    )
-                    seq.change(
-                        lambda val, ents, i=idx: update_entity_field(ents, i, "sequence", val),
-                        inputs=[seq, entities_state],
-                        outputs=entities_state,
-                    )
-                    remove_btn.click(
-                        lambda ents, i=idx: remove_entity(ents, i),
-                        inputs=entities_state,
-                        outputs=entities_state,
-                    )
-
-            add_btn.click(add_entity, inputs=entities_state, outputs=entities_state)
-            example1_btn.click(lambda: load_example(EXAMPLE_UBIQUITIN), outputs=entities_state)
-            example2_btn.click(lambda: load_example(EXAMPLE_CAP), outputs=entities_state)
+            # Empty placeholder, positioned here (above the settings/fold/download components
+            # below); filled in further down, once fold_btn etc. exist for render_entities to
+            # close over — see the ``with rows_container:`` block after this gr.Row().
+            rows_container = gr.Column()
 
             with gr.Accordion("Advanced settings", open=False):
                 num_loops = gr.Slider(1, 3, value=1, step=1, label="Recycling loops")
@@ -281,11 +250,58 @@ with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
         with gr.Column(scale=2):
             structure_output = gr.HTML(label="Predicted structure")
 
-    fold_btn.click(
-        fn=fold_entities,
-        inputs=[entities_state, num_loops, num_sampling_steps],
-        outputs=[structure_output, summary_output, pdb_output],
-    )
+    example1_btn.click(lambda: load_example(EXAMPLE_UBIQUITIN), outputs=entities_state)
+    example2_btn.click(lambda: load_example(EXAMPLE_CAP), outputs=entities_state)
+
+    with rows_container:
+
+        @gr.render(inputs=entities_state)
+        def render_entities(entities):
+            etype_boxes, copies_boxes, seq_boxes, remove_btns = [], [], [], []
+            for idx, entity in enumerate(entities):
+                with gr.Row():
+                    etype = gr.Dropdown(ENTITY_TYPES, value=entity["type"], label="Type", scale=2)
+                    copies = gr.Number(
+                        value=entity["copies"],
+                        label="Copies",
+                        precision=0,
+                        minimum=1,
+                        step=1,
+                        scale=1,
+                    )
+                    seq = gr.Textbox(
+                        value=entity["sequence"],
+                        label="Sequence",
+                        lines=4,
+                        max_lines=4,
+                        scale=6,
+                        placeholder="Amino acid sequence...",
+                    )
+                    remove_btn = gr.Button("×", scale=0, min_width=40, visible=len(entities) > 1)
+                etype_boxes.append(etype)
+                copies_boxes.append(copies)
+                seq_boxes.append(seq)
+                remove_btns.append(remove_btn)
+
+            # No change/input handlers on the row fields themselves: Add/Remove/Fold all read
+            # the live, currently-displayed value of every row directly (below), so nothing
+            # depends on a change event having already round-tripped to the server first.
+            live_inputs = [*etype_boxes, *copies_boxes, *seq_boxes]
+
+            add_btn.click(add_entity_live, inputs=live_inputs, outputs=entities_state)
+            for idx, remove_btn in enumerate(remove_btns):
+                remove_btn.click(
+                    lambda *vals, i=idx, n=len(entities): remove_entity_live(
+                        vals[:n], vals[n : 2 * n], vals[2 * n :], i
+                    ),
+                    inputs=live_inputs,
+                    outputs=entities_state,
+                )
+            fold_btn.click(
+                fn=fold_from_rows,
+                inputs=[*live_inputs, num_loops, num_sampling_steps],
+                outputs=[structure_output, summary_output, pdb_output],
+            )
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(theme=gr.themes.Soft())