merle commited on
Commit
efae0f6
·
1 Parent(s): 5e87e3d

Upgrade to Gradio 6.9.0 with full API migration

Browse files

- README: sdk_version 3.36.1 -> 6.9.0, keep python_version 3.10
- requirements: torch 2.5.0, dgl 2.2.1, pydantic>=2.0, numpy<2
(e3nn stays at 0.3.3 for checkpoint compatibility)
- app.py Gradio migration:
- .style() calls -> direct constructor params (Row equal_height)
- gr.update() -> component constructors (Slider, Textbox)
- gr.TabItem -> gr.Tab
- theme moved from Blocks() to launch()
- Removed redundant demo.queue()
- File upload .name -> helper for path extraction
- Switched to gr.themes.Soft() (ParityError/Interstellar incompatible)
- Removed dead commented-out code blocks
- General cleanup

Made-with: Cursor

Files changed (4) hide show
  1. .gitignore +1 -0
  2. README.md +1 -1
  3. app.py +87 -207
  4. requirements.txt +4 -4
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ .venv/
README.md CHANGED
@@ -5,7 +5,7 @@ thumbnail: http://files.ipd.uw.edu/pub/sequence_diffusion/figs/diffusion_landsca
5
  colorFrom: blue
6
  colorTo: purple
7
  sdk: gradio
8
- sdk_version: 3.36.1
9
  python_version: "3.10"
10
  app_file: app.py
11
  pinned: false
 
5
  colorFrom: blue
6
  colorTo: purple
7
  sdk: gradio
8
+ sdk_version: 6.9.0
9
  python_version: "3.10"
10
  app_file: app.py
11
  pinned: false
app.py CHANGED
@@ -32,7 +32,6 @@ plt.rcParams.update({'font.size': 13})
32
  with open('./tmp/args.json','r') as f:
33
  args = json.load(f)
34
 
35
- # manually set checkpoint to load
36
  args['checkpoint'] = None
37
  args['dump_trb'] = False
38
  args['dump_args'] = True
@@ -43,25 +42,27 @@ args['loop_bias'] = 0.0
43
  args['helix_bias'] = 0.0
44
 
45
 
 
 
 
 
 
46
 
47
- def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias,
48
- secondary_structure, aa_bias, aa_bias_potential,
49
- #target_charge, target_ph, charge_potential,
50
  num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
51
  contigs, pssm, seq_mask, str_mask, rewrite_pdb):
52
-
53
  dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
54
  og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'
55
 
56
  model_args = copy.deepcopy(args)
57
 
58
- # make sampler
59
  S = HuggingFace_sampler(args=model_args)
60
 
61
- # get random prefix
62
  S.out_prefix = './tmp/'+secrets.token_hex(nbytes=10).upper()
63
 
64
- # set args
65
  S.args['checkpoint'] = None
66
  S.args['dump_trb'] = False
67
  S.args['dump_args'] = True
@@ -74,8 +75,6 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
74
  S.args['potential_scale'] = None
75
  S.args['aa_composition'] = None
76
 
77
-
78
- # get sequence if entered and make sure all chars are valid
79
  alt_aa_dict = {'B':['D','N'],'J':['I','L'],'U':['C'],'Z':['E','Q'],'O':['K']}
80
  if sequence not in ['',None]:
81
  L = len(sequence)
@@ -92,10 +91,9 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
92
  else:
93
  S.args['contigs'] = [f'{seq_len}']
94
  L = int(seq_len)
95
-
96
- print('DEBUG: ',rewrite_pdb)
97
  if rewrite_pdb not in ['',None]:
98
- S.args['pdb'] = rewrite_pdb.name
99
 
100
  if seq_mask not in ['',None]:
101
  S.args['inpaint_seq'] = [seq_mask]
@@ -113,9 +111,7 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
113
  else:
114
  dseq = L - len(secondary_structure)
115
  secondary_structure += secondary_structure[-1]*dseq
116
-
117
 
118
- # potentials
119
  potential_list = []
120
  potential_bias_list = []
121
 
@@ -125,49 +121,33 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
125
  if aa_bias_potential in ['',None]:
126
  aa_bias_potential = 3
127
  potential_bias_list.append(str(aa_bias_potential))
128
- '''
129
- if target_charge not in ['',None]:
130
- potential_list.append('charge')
131
- if charge_potential in ['',None]:
132
- charge_potential = 1
133
- potential_bias_list.append(str(charge_potential))
134
- S.args['target_charge'] = float(target_charge)
135
- if target_ph in ['',None]:
136
- target_ph = 7.4
137
- S.args['target_pH'] = float(target_ph)
138
- '''
139
-
140
  if hydrophobic_target_score not in ['',None]:
141
  potential_list.append('hydrophobic')
142
  S.args['hydrophobic_score'] = float(hydrophobic_target_score)
143
  if hydrophobic_potential in ['',None]:
144
  hydrophobic_potential = 3
145
  potential_bias_list.append(str(hydrophobic_potential))
146
-
147
  if pssm not in ['',None]:
148
  potential_list.append('PSSM')
149
  potential_bias_list.append('5')
150
- S.args['PSSM'] = pssm.name
151
-
152
 
153
  if len(potential_list) > 0:
154
  S.args['potentials'] = ','.join(potential_list)
155
  S.args['potential_scale'] = ','.join(potential_bias_list)
156
 
157
-
158
- # normalise secondary_structure bias from range 0-0.3
159
  S.args['secondary_structure'] = secondary_structure
160
  S.args['helix_bias'] = helix_bias
161
  S.args['strand_bias'] = strand_bias
162
  S.args['loop_bias'] = loop_bias
163
-
164
- # set T
165
  if num_steps in ['',None]:
166
  S.args['T'] = 20
167
  else:
168
  S.args['T'] = int(num_steps)
169
 
170
- # noise
171
  if 'normal' in noise:
172
  S.args['sample_distribution'] = noise
173
  S.args['sample_distribution_gmm_means'] = [0]
@@ -181,8 +161,6 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
181
  S.args['sample_distribution_gmm_means'] = [-1,0,1]
182
  S.args['sample_distribution_gmm_variances'] = [1,1,1]
183
 
184
-
185
-
186
  if secondary_structure not in ['',None] or helix_bias+strand_bias+loop_bias > 0:
187
  S.args['checkpoint'] = dssp_checkpoint
188
  S.args['d_t1d'] = 29
@@ -191,24 +169,21 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
191
  S.args['checkpoint'] = og_checkpoint
192
  S.args['d_t1d'] = 24
193
  print('using og checkpoint')
194
-
195
 
196
  for k,v in S.args.items():
197
  print(f"{k} --> {v}")
198
-
199
- # init S
200
  S.model_init()
201
  S.diffuser_init()
202
  S.setup()
203
 
204
- # sampling loop
205
  plddt_data = []
206
  for j in range(S.max_t):
207
  print(f'on step {j}')
208
  output_seq, output_pdb, plddt = S.take_step_get_outputs(j)
209
  plddt_data.append(plddt)
210
  yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
211
-
212
  output_seq, output_pdb, plddt = S.get_outputs()
213
  yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
214
 
@@ -224,47 +199,27 @@ def get_plddt_plot(plddt_data, max_t):
224
  return fig
225
 
226
  def display_pdb(path_to_pdb):
227
- '''
228
- #function to display pdb in py3dmol
229
- '''
230
  pdb = open(path_to_pdb, "r").read()
231
-
232
  view = py3Dmol.view(width=500, height=500)
233
  view.addModel(pdb, "pdb")
234
- view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})#'linear', 'min': 0, 'max': 1, 'colors': ["#ff9ef0","#a903fc",]}}})
235
  view.zoomTo()
236
  output = view._make_html().replace("'", '"')
237
- print(view._make_html())
238
- x = f"""<!DOCTYPE html><html></center> {output} </center></html>""" # do not use ' in this input
239
-
240
  return f"""<iframe height="500px" width="100%" name="result" allow="midi; geolocation; microphone; camera;
241
  display-capture; encrypted-media;" sandbox="allow-modals allow-forms
242
  allow-scripts allow-same-origin allow-popups
243
  allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
244
  allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
245
 
246
- '''
247
- return f"""<iframe style="width: 100%; height:700px" name="result" allow="midi; geolocation; microphone; camera;
248
- display-capture; encrypted-media;" sandbox="allow-modals allow-forms
249
- allow-scripts allow-same-origin allow-popups
250
- allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
251
- allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
252
- '''
253
-
254
-
255
 
256
  # MOTIF SCAFFOLDING
257
  def get_motif_preview(pdb_id, contigs):
258
- '''
259
- #function to display selected motif in py3dmol
260
- '''
261
  input_pdb = fetch_pdb(pdb_id=pdb_id.lower())
262
 
263
- # rewrite pdb
264
  parse = parse_pdb(input_pdb)
265
- #output_name = './rewrite_'+input_pdb.split('/')[-1]
266
- #writepdb(output_name, torch.tensor(parse_og['xyz']),torch.tensor(parse_og['seq']))
267
- #parse = parse_pdb(output_name)
268
  output_name = input_pdb
269
 
270
  pdb = open(output_name, "r").read()
@@ -276,11 +231,7 @@ def get_motif_preview(pdb_id, contigs):
276
  else:
277
  contigs = [contigs]
278
 
279
- print('DEBUG: ',contigs)
280
-
281
  pdb_map = get_mappings(ContigMap(parse,contigs))
282
- print('DEBUG: ',pdb_map)
283
- print('DEBUG: ',pdb_map['con_ref_idx0'])
284
  roi = [x[1]-1 for x in pdb_map['con_ref_pdb_idx']]
285
 
286
  colormap = {0:'#D3D3D3', 1:'#F74CFF'}
@@ -288,9 +239,8 @@ def get_motif_preview(pdb_id, contigs):
288
  view.setStyle({"cartoon": {"colorscheme": {"prop": "resi", "map": colors}}})
289
  view.zoomTo()
290
  output = view._make_html().replace("'", '"')
291
- print(view._make_html())
292
- x = f"""<!DOCTYPE html><html></center> {output} </center></html>""" # do not use ' in this input
293
-
294
  return f"""<iframe height="500px" width="100%" name="result" allow="midi; geolocation; microphone; camera;
295
  display-capture; encrypted-media;" sandbox="allow-modals allow-forms
296
  allow-scripts allow-same-origin allow-popups
@@ -306,25 +256,21 @@ def fetch_pdb(pdb_id=None):
306
 
307
  # MSA AND PSSM GUIDANCE
308
  def save_pssm(file_upload):
309
- filename = file_upload.name
310
- orig_name = file_upload.orig_name
311
  if filename.split('.')[-1] in ['fasta', 'a3m']:
312
  return msa_to_pssm(file_upload)
313
  return filename
314
 
315
  def msa_to_pssm(msa_file):
316
- # Define the lookup table for converting amino acids to indices
317
  aa_to_index = {'A': 0, 'R': 1, 'N': 2, 'D': 3, 'C': 4, 'Q': 5, 'E': 6, 'G': 7, 'H': 8, 'I': 9, 'L': 10,
318
  'K': 11, 'M': 12, 'F': 13, 'P': 14, 'S': 15, 'T': 16, 'W': 17, 'Y': 18, 'V': 19, 'X': 20, '-': 21}
319
- # Open the FASTA file and read the sequences
320
- records = list(SeqIO.parse(msa_file.name, "fasta"))
321
 
322
- assert len(records) >= 1, "MSA must contain more than one protein sequecne."
323
 
324
  first_seq = str(records[0].seq)
325
  aligned_seqs = [first_seq]
326
- # print(aligned_seqs)
327
- # Perform sequence alignment using the Needleman-Wunsch algorithm
328
  aligner = Align.PairwiseAligner()
329
  aligner.open_gap_score = -0.7
330
  aligner.extend_gap_score = -0.3
@@ -343,13 +289,9 @@ def msa_to_pssm(msa_file):
343
  al1_fin += al1[i]
344
  al2_fin += al2[i]
345
  aligned_seqs.append(str(al2_fin))
346
- # Get the length of the aligned sequences
347
  aligned_seq_length = len(first_seq)
348
- # Initialize the position scoring matrix
349
  matrix = np.zeros((22, aligned_seq_length))
350
- # Iterate through the aligned sequences and count the amino acids at each position
351
  for seq in aligned_seqs:
352
- #print(seq)
353
  for i in range(aligned_seq_length):
354
  if i == len(seq):
355
  break
@@ -359,19 +301,18 @@ def msa_to_pssm(msa_file):
359
  else:
360
  aa_index = aa_to_index[amino_acid.upper()]
361
  matrix[aa_index, i] += 1
362
- # Normalize the counts to get the frequency of each amino acid at each position
363
  matrix /= len(aligned_seqs)
364
  print(len(aligned_seqs))
365
  matrix[20:,]=0
366
 
367
- outdir = ".".join(msa_file.name.split('.')[:-1]) + ".csv"
 
368
  np.savetxt(outdir, matrix[:21,:].T, delimiter=",")
369
  return outdir
370
 
371
  def get_pssm(fasta_msa, input_pssm):
372
-
373
  if input_pssm not in ['',None]:
374
- outdir = input_pssm.name
375
  else:
376
  outdir = save_pssm(fasta_msa)
377
 
@@ -382,110 +323,96 @@ def get_pssm(fasta_msa, input_pssm):
382
  return fig, outdir
383
 
384
 
385
- #toggle options
386
  def toggle_seq_input(choice):
387
  if choice == "protein length":
388
- return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
389
  elif choice == "custom sequence":
390
- return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
391
 
392
  def toggle_secondary_structure(choice):
393
  if choice == "sliders":
394
- return gr.update(visible=True, value=None),gr.update(visible=True, value=None),gr.update(visible=True, value=None),gr.update(visible=False, value=None)
395
  elif choice == "explicit":
396
- return gr.update(visible=False, value=None),gr.update(visible=False, value=None),gr.update(visible=False, value=None),gr.update(visible=True, value=None)
397
 
398
 
399
- # Define the Gradio interface
400
- with gr.Blocks(theme='ParityError/Interstellar') as demo:
401
-
402
- gr.Markdown(f"""# Protein Generation via Diffusion in Sequence Space""")
403
 
404
  with gr.Row():
405
  with gr.Column(min_width=500):
406
- gr.Markdown(f"""
407
  ## How does it work?\n
408
- !!! [PAPER](https://www.nature.com/articles/s41587-024-02395-w) !!!
409
  Protein sequence and structure co-generation is a long outstanding problem in the field of protein design. By implementing [ddpm](https://arxiv.org/abs/2006.11239) style diffusion over protein seqeuence space we generate protein sequence and structure pairs. Starting with [RoseTTAFold](https://www.science.org/doi/10.1126/science.abj8754), a protein structure prediction network, we finetuned it to predict sequence and structure given a partially noised sequence. By applying losses to both the predicted sequence and structure the model is forced to generate meaningful pairs. Diffusing in sequence space makes it easy to implement potentials to guide the diffusive process toward particular amino acid composition, net charge, and more! Furthermore, you can sample proteins from a family of sequences or even train a small sequence to function classifier to guide generation toward desired sequences.
410
  ![fig1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/diffusion_landscape.png)
411
-
412
  ## How to use it?\n
413
  A user can either design a custom input sequence to diffuse from or specify a length below. To scaffold a sequence use the following format where X represent residues to diffuse: XXXXXXXXSCIENCESCIENCEXXXXXXXXXXXXXXXXXXX. You can even design a protein with your name XXXXXXXXXXXXNAMEHEREXXXXXXXXXXXXX!
414
-
415
  ### Acknowledgements\n
416
  Thank you to Simon Dürr and the Hugging Face team for setting us up with a community GPU grant!
417
  """)
418
-
419
  gr.Markdown("""
420
  ## Model in Action
421
  ![gif1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/seqdiff_anim_720p.gif)
422
  """)
423
 
424
- #with gr.Row().style(equal_height=False):
425
- with gr.Row():
426
  with gr.Column():
427
  with gr.Tabs():
428
- with gr.TabItem("Inputs"):
429
- gr.Markdown("""## INPUTS""")
430
  gr.Markdown("""#### Start Sequence
431
  Specify the protein length for complete unconditional generation, or scaffold a motif (or your name) using the custom sequence input""")
432
  seq_opt = gr.Radio(["protein length","custom sequence"], label="How would you like to specify the starting sequence?", value='protein length')
433
 
434
  sequence = gr.Textbox(label="custom sequence", lines=1, placeholder='AMINO ACIDS: A,C,D,E,F,G,H,I,K,L,M,N,P,Q,R,S,T,V,W,Y\n MASK TOKEN: X', visible=False)
435
  seq_len = gr.Slider(minimum=5.0, maximum=250.0, label="protein length", value=100, visible=True)
436
-
437
  seq_opt.change(fn=toggle_seq_input,
438
  inputs=[seq_opt],
439
- outputs=[seq_len, sequence],
440
- queue=False)
441
 
442
- gr.Markdown("""### Optional Parameters""")
443
  with gr.Accordion(label='Secondary Structure',open=True):
444
- gr.Markdown("""Try changing the sliders or inputing explicit secondary structure conditioning for each residue""")
445
  sec_str_opt = gr.Radio(["sliders","explicit"], label="How would you like to specify secondary structure?", value='sliders')
446
 
447
  secondary_structure = gr.Textbox(label="secondary structure", lines=1, placeholder='HELIX = H STRAND = S LOOP = L MASK = X(must be the same length as input sequence)', visible=False)
448
-
449
  with gr.Column():
450
  helix_bias = gr.Slider(minimum=0.0, maximum=0.05, label="helix bias", visible=True)
451
  strand_bias = gr.Slider(minimum=0.0, maximum=0.05, label="strand bias", visible=True)
452
  loop_bias = gr.Slider(minimum=0.0, maximum=0.20, label="loop bias", visible=True)
453
-
454
  sec_str_opt.change(fn=toggle_secondary_structure,
455
  inputs=[sec_str_opt],
456
- outputs=[helix_bias,strand_bias,loop_bias,secondary_structure],
457
- queue=False)
458
-
459
  with gr.Accordion(label='Amino Acid Compositional Bias',open=False):
460
- gr.Markdown("""Bias sequence composition for particular amino acids by specifying the one letter code followed by the fraction to bias. This can be input as a list for example: W0.2,E0.1""")
461
  with gr.Row():
462
  aa_bias = gr.Textbox(label="aa bias", lines=1, placeholder='specify one letter AA and fraction to bias, for example W0.1 or M0.1,K0.1' )
463
  aa_bias_potential = gr.Textbox(label="aa bias scale", lines=1, placeholder='AA Bias potential scale (recomended range 1.0-5.0)')
464
-
465
- '''
466
- with gr.Accordion(label='Charge Bias',open=False):
467
- gr.Markdown("""Bias for a specified net charge at a particular pH using the boxes below""")
468
- with gr.Row():
469
- target_charge = gr.Textbox(label="net charge", lines=1, placeholder='net charge to target')
470
- target_ph = gr.Textbox(label="pH", lines=1, placeholder='pH at which net charge is desired')
471
- charge_potential = gr.Textbox(label="charge potential scale", lines=1, placeholder='charge potential scale (recomended range 1.0-5.0)')
472
- '''
473
 
474
  with gr.Accordion(label='Hydrophobic Bias',open=False):
475
- gr.Markdown("""Bias for or against hydrophobic composition, to get more soluble proteins, bias away with a negative target score (ex. -5)""")
476
  with gr.Row():
477
  hydrophobic_target_score = gr.Textbox(label="hydrophobic score", lines=1, placeholder='hydrophobic score to target (negative score is good for solublility)')
478
  hydrophobic_potential = gr.Textbox(label="hydrophobic potential scale", lines=1, placeholder='hydrophobic potential scale (recomended range 1.0-2.0)')
479
-
480
  with gr.Accordion(label='Diffusion Params',open=False):
481
- gr.Markdown("""Increasing T to more steps can be helpful for harder design challenges, sampling from different distributions can change the sequence and structural composition""")
482
  with gr.Row():
483
  num_steps = gr.Textbox(label="T", lines=1, placeholder='number of diffusion steps (25 or less will speed things up)')
484
- noise = gr.Dropdown(['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'], label='noise type', value='normal')
485
-
486
- with gr.TabItem("Motif Selection"):
487
 
488
- gr.Markdown("""### Motif Selection Preview""")
489
  gr.Markdown('Contigs explained: to grab residues (seq and str) on a pdb chain you will provide the chain letter followed by a range of residues as indexed in the pdb file for example (A3-10) is the syntax to select residues 3-10 on chain A (the chain always needs to be specified). To add diffused residues to either side of this motif you can specify a range or discrete value without a chain letter infront. To add 15 residues before the motif and 20-30 residues (randomly sampled) after use the following syntax: 15,A3-10,20-30 commas are used to separate regions selected from the pdb and designed (diffused) resiudes which will be added. ')
490
  pdb_id_code = gr.Textbox(label="PDB ID", lines=1, placeholder='INPUT PDB ID TO FETCH (ex. 1DPX)', visible=True)
491
  contigs = gr.Textbox(label="contigs", lines=1, placeholder='specify contigs to grab particular residues from pdb ()', visible=True)
@@ -497,8 +424,8 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
497
  rewrite_pdb = gr.File(label='PDB file')
498
  preview_btn = gr.Button("Preview Motif")
499
 
500
- with gr.TabItem("MSA to PSSM"):
501
- gr.Markdown("""### MSA to PSSM Generation""")
502
  gr.Markdown('input either an MSA or PSSM to guide the model toward generating samples within your family of interest')
503
  with gr.Row():
504
  fasta_msa = gr.File(label='MSA')
@@ -510,92 +437,45 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
510
 
511
  btn = gr.Button("GENERATE")
512
 
513
- #with gr.Row():
514
  with gr.Column():
515
- gr.Markdown("""## OUTPUTS""")
516
- gr.Markdown("""#### Confidence score for generated structure at each timestep""")
517
  plddt_plot = gr.Plot(label='plddt at step t')
518
- gr.Markdown("""#### Output protein sequnece""")
519
  output_seq = gr.Textbox(label="sequence")
520
- gr.Markdown("""#### Download PDB file""")
521
  output_pdb = gr.File(label="PDB file")
522
- gr.Markdown("""#### Structure viewer""")
523
  output_viewer = gr.HTML()
524
- '''
525
- gr.Markdown("""### Don't know where to get started? Click on an example below to try it out!""")
526
- gr.Examples(
527
- [["","125",0.0,0.0,0.2,"","","","20","normal",'','','',None,'','',None],
528
- ["","100",0.0,0.0,0.0,"","W0.2","2","20","normal",'','','',None,'','',None],
529
- # ["","100",0.0,0.0,0.0,
530
- # "XXHHHHHHHHHXXXXXXXHHHHHHHHHXXXXXXXHHHHHHHHXXXXSSSSSSSSSSSXXXXXXXXSSSSSSSSSSSSXXXXXXXSSSSSSSSSXXXXXXX",
531
- # "","","25","normal",'','','',None,'','',None],
532
- # ["XXXXXXXXXXXXXXXXXXXXXXXXXIPDXXXXXXXXXXXXXXXXXXXXXXPEPSEQXXXXXXXXXXXXXXXXXXXXXXXXXXIPDXXXXXXXXXXXXXXXXXXX",
533
- # "",0.0,0.0,0.0,"","","","25","normal",'','','',None,'','',None],
534
- # ["","",0.0,0.0,0.0,"","","","25","normal",'','',
535
- # '9,D10-11,8,D20-20,4,D25-35,65,D101-101,2,D104-105,8,D114-116,15,D132-138,6,D145-145,2,D148-148,12,D161-161,3',
536
- # './tmp/PSSM_lysozyme.csv',
537
- # 'D25-25,D27-31,D33-35,D132-137',
538
- # 'D26-26','./tmp/150l.pdb']
539
- ],
540
- inputs=[sequence,
541
- seq_len,
542
- helix_bias,
543
- strand_bias,
544
- loop_bias,
545
- secondary_structure,
546
- aa_bias,
547
- aa_bias_potential,
548
- #target_charge,
549
- #target_ph,
550
- #charge_potential,
551
- num_steps,
552
- noise,
553
- hydrophobic_target_score,
554
- hydrophobic_potential,
555
- contigs,
556
- pssm,
557
- seq_mask,
558
- str_mask,
559
- rewrite_pdb],
560
- outputs=[output_seq,
561
- output_pdb,
562
- output_viewer,
563
- plddt_plot],
564
- fn=protein_diffusion_model,
565
- )
566
- '''
567
  preview_btn.click(get_motif_preview,[pdb_id_code, contigs],[preview_viewer, rewrite_pdb])
568
 
569
  pssm_gen_btn.click(get_pssm,[fasta_msa,input_pssm],[pssm_view, pssm])
570
 
571
- btn.click(protein_diffusion_model,
572
- [sequence,
573
- seq_len,
574
- helix_bias,
575
- strand_bias,
576
- loop_bias,
577
- secondary_structure,
578
- aa_bias,
579
- aa_bias_potential,
580
- #target_charge,
581
- #target_ph,
582
- #charge_potential,
583
- num_steps,
584
- noise,
585
- hydrophobic_target_score,
586
  hydrophobic_potential,
587
  contigs,
588
  pssm,
589
- seq_mask,
590
  str_mask,
591
- rewrite_pdb],
592
  [output_seq,
593
  output_pdb,
594
  output_viewer,
595
  plddt_plot])
596
 
597
- demo.queue()
598
- demo.launch(debug=True)
599
-
600
-
601
-
 
32
  with open('./tmp/args.json','r') as f:
33
  args = json.load(f)
34
 
 
35
  args['checkpoint'] = None
36
  args['dump_trb'] = False
37
  args['dump_args'] = True
 
42
  args['helix_bias'] = 0.0
43
 
44
 
45
+ def _get_file_path(file_obj):
46
+ """Extract file path from a Gradio file upload (handles both old and new Gradio formats)."""
47
+ if isinstance(file_obj, str):
48
+ return file_obj
49
+ return file_obj.name
50
 
51
+
52
+ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias,
53
+ secondary_structure, aa_bias, aa_bias_potential,
54
  num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
55
  contigs, pssm, seq_mask, str_mask, rewrite_pdb):
56
+
57
  dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
58
  og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'
59
 
60
  model_args = copy.deepcopy(args)
61
 
 
62
  S = HuggingFace_sampler(args=model_args)
63
 
 
64
  S.out_prefix = './tmp/'+secrets.token_hex(nbytes=10).upper()
65
 
 
66
  S.args['checkpoint'] = None
67
  S.args['dump_trb'] = False
68
  S.args['dump_args'] = True
 
75
  S.args['potential_scale'] = None
76
  S.args['aa_composition'] = None
77
 
 
 
78
  alt_aa_dict = {'B':['D','N'],'J':['I','L'],'U':['C'],'Z':['E','Q'],'O':['K']}
79
  if sequence not in ['',None]:
80
  L = len(sequence)
 
91
  else:
92
  S.args['contigs'] = [f'{seq_len}']
93
  L = int(seq_len)
94
+
 
95
  if rewrite_pdb not in ['',None]:
96
+ S.args['pdb'] = _get_file_path(rewrite_pdb)
97
 
98
  if seq_mask not in ['',None]:
99
  S.args['inpaint_seq'] = [seq_mask]
 
111
  else:
112
  dseq = L - len(secondary_structure)
113
  secondary_structure += secondary_structure[-1]*dseq
 
114
 
 
115
  potential_list = []
116
  potential_bias_list = []
117
 
 
121
  if aa_bias_potential in ['',None]:
122
  aa_bias_potential = 3
123
  potential_bias_list.append(str(aa_bias_potential))
124
+
 
 
 
 
 
 
 
 
 
 
 
125
  if hydrophobic_target_score not in ['',None]:
126
  potential_list.append('hydrophobic')
127
  S.args['hydrophobic_score'] = float(hydrophobic_target_score)
128
  if hydrophobic_potential in ['',None]:
129
  hydrophobic_potential = 3
130
  potential_bias_list.append(str(hydrophobic_potential))
131
+
132
  if pssm not in ['',None]:
133
  potential_list.append('PSSM')
134
  potential_bias_list.append('5')
135
+ S.args['PSSM'] = _get_file_path(pssm)
 
136
 
137
  if len(potential_list) > 0:
138
  S.args['potentials'] = ','.join(potential_list)
139
  S.args['potential_scale'] = ','.join(potential_bias_list)
140
 
 
 
141
  S.args['secondary_structure'] = secondary_structure
142
  S.args['helix_bias'] = helix_bias
143
  S.args['strand_bias'] = strand_bias
144
  S.args['loop_bias'] = loop_bias
145
+
 
146
  if num_steps in ['',None]:
147
  S.args['T'] = 20
148
  else:
149
  S.args['T'] = int(num_steps)
150
 
 
151
  if 'normal' in noise:
152
  S.args['sample_distribution'] = noise
153
  S.args['sample_distribution_gmm_means'] = [0]
 
161
  S.args['sample_distribution_gmm_means'] = [-1,0,1]
162
  S.args['sample_distribution_gmm_variances'] = [1,1,1]
163
 
 
 
164
  if secondary_structure not in ['',None] or helix_bias+strand_bias+loop_bias > 0:
165
  S.args['checkpoint'] = dssp_checkpoint
166
  S.args['d_t1d'] = 29
 
169
  S.args['checkpoint'] = og_checkpoint
170
  S.args['d_t1d'] = 24
171
  print('using og checkpoint')
 
172
 
173
  for k,v in S.args.items():
174
  print(f"{k} --> {v}")
175
+
 
176
  S.model_init()
177
  S.diffuser_init()
178
  S.setup()
179
 
 
180
  plddt_data = []
181
  for j in range(S.max_t):
182
  print(f'on step {j}')
183
  output_seq, output_pdb, plddt = S.take_step_get_outputs(j)
184
  plddt_data.append(plddt)
185
  yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
186
+
187
  output_seq, output_pdb, plddt = S.get_outputs()
188
  yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
189
 
 
199
  return fig
200
 
201
  def display_pdb(path_to_pdb):
 
 
 
202
  pdb = open(path_to_pdb, "r").read()
203
+
204
  view = py3Dmol.view(width=500, height=500)
205
  view.addModel(pdb, "pdb")
206
+ view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})
207
  view.zoomTo()
208
  output = view._make_html().replace("'", '"')
209
+ x = f"""<!DOCTYPE html><html></center> {output} </center></html>"""
210
+
 
211
  return f"""<iframe height="500px" width="100%" name="result" allow="midi; geolocation; microphone; camera;
212
  display-capture; encrypted-media;" sandbox="allow-modals allow-forms
213
  allow-scripts allow-same-origin allow-popups
214
  allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
215
  allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
216
 
 
 
 
 
 
 
 
 
 
217
 
218
  # MOTIF SCAFFOLDING
219
  def get_motif_preview(pdb_id, contigs):
 
 
 
220
  input_pdb = fetch_pdb(pdb_id=pdb_id.lower())
221
 
 
222
  parse = parse_pdb(input_pdb)
 
 
 
223
  output_name = input_pdb
224
 
225
  pdb = open(output_name, "r").read()
 
231
  else:
232
  contigs = [contigs]
233
 
 
 
234
  pdb_map = get_mappings(ContigMap(parse,contigs))
 
 
235
  roi = [x[1]-1 for x in pdb_map['con_ref_pdb_idx']]
236
 
237
  colormap = {0:'#D3D3D3', 1:'#F74CFF'}
 
239
  view.setStyle({"cartoon": {"colorscheme": {"prop": "resi", "map": colors}}})
240
  view.zoomTo()
241
  output = view._make_html().replace("'", '"')
242
+ x = f"""<!DOCTYPE html><html></center> {output} </center></html>"""
243
+
 
244
  return f"""<iframe height="500px" width="100%" name="result" allow="midi; geolocation; microphone; camera;
245
  display-capture; encrypted-media;" sandbox="allow-modals allow-forms
246
  allow-scripts allow-same-origin allow-popups
 
256
 
257
  # MSA AND PSSM GUIDANCE
258
  def save_pssm(file_upload):
259
+ filename = _get_file_path(file_upload)
260
+ orig_name = file_upload.orig_name if hasattr(file_upload, 'orig_name') else filename
261
  if filename.split('.')[-1] in ['fasta', 'a3m']:
262
  return msa_to_pssm(file_upload)
263
  return filename
264
 
265
  def msa_to_pssm(msa_file):
 
266
  aa_to_index = {'A': 0, 'R': 1, 'N': 2, 'D': 3, 'C': 4, 'Q': 5, 'E': 6, 'G': 7, 'H': 8, 'I': 9, 'L': 10,
267
  'K': 11, 'M': 12, 'F': 13, 'P': 14, 'S': 15, 'T': 16, 'W': 17, 'Y': 18, 'V': 19, 'X': 20, '-': 21}
268
+ records = list(SeqIO.parse(_get_file_path(msa_file), "fasta"))
 
269
 
270
+ assert len(records) >= 1, "MSA must contain more than one protein sequence."
271
 
272
  first_seq = str(records[0].seq)
273
  aligned_seqs = [first_seq]
 
 
274
  aligner = Align.PairwiseAligner()
275
  aligner.open_gap_score = -0.7
276
  aligner.extend_gap_score = -0.3
 
289
  al1_fin += al1[i]
290
  al2_fin += al2[i]
291
  aligned_seqs.append(str(al2_fin))
 
292
  aligned_seq_length = len(first_seq)
 
293
  matrix = np.zeros((22, aligned_seq_length))
 
294
  for seq in aligned_seqs:
 
295
  for i in range(aligned_seq_length):
296
  if i == len(seq):
297
  break
 
301
  else:
302
  aa_index = aa_to_index[amino_acid.upper()]
303
  matrix[aa_index, i] += 1
 
304
  matrix /= len(aligned_seqs)
305
  print(len(aligned_seqs))
306
  matrix[20:,]=0
307
 
308
+ msa_path = _get_file_path(msa_file)
309
+ outdir = ".".join(msa_path.split('.')[:-1]) + ".csv"
310
  np.savetxt(outdir, matrix[:21,:].T, delimiter=",")
311
  return outdir
312
 
313
  def get_pssm(fasta_msa, input_pssm):
 
314
  if input_pssm not in ['',None]:
315
+ outdir = _get_file_path(input_pssm)
316
  else:
317
  outdir = save_pssm(fasta_msa)
318
 
 
323
  return fig, outdir
324
 
325
 
 
326
  def toggle_seq_input(choice):
327
  if choice == "protein length":
328
+ return gr.Slider(visible=True, value=None), gr.Textbox(visible=False, value=None)
329
  elif choice == "custom sequence":
330
+ return gr.Slider(visible=False, value=None), gr.Textbox(visible=True, value=None)
331
 
332
  def toggle_secondary_structure(choice):
333
  if choice == "sliders":
334
+ return gr.Slider(visible=True, value=None),gr.Slider(visible=True, value=None),gr.Slider(visible=True, value=None),gr.Textbox(visible=False, value=None)
335
  elif choice == "explicit":
336
+ return gr.Slider(visible=False, value=None),gr.Slider(visible=False, value=None),gr.Slider(visible=False, value=None),gr.Textbox(visible=True, value=None)
337
 
338
 
339
+ with gr.Blocks() as demo:
340
+
341
+ gr.Markdown("# Protein Generation via Diffusion in Sequence Space")
 
342
 
343
  with gr.Row():
344
  with gr.Column(min_width=500):
345
+ gr.Markdown("""
346
  ## How does it work?\n
347
+ --- [PREPRINT](https://biorxiv.org/content/10.1101/2023.05.08.539766v1) ---
348
  Protein sequence and structure co-generation is a long outstanding problem in the field of protein design. By implementing [ddpm](https://arxiv.org/abs/2006.11239) style diffusion over protein seqeuence space we generate protein sequence and structure pairs. Starting with [RoseTTAFold](https://www.science.org/doi/10.1126/science.abj8754), a protein structure prediction network, we finetuned it to predict sequence and structure given a partially noised sequence. By applying losses to both the predicted sequence and structure the model is forced to generate meaningful pairs. Diffusing in sequence space makes it easy to implement potentials to guide the diffusive process toward particular amino acid composition, net charge, and more! Furthermore, you can sample proteins from a family of sequences or even train a small sequence to function classifier to guide generation toward desired sequences.
349
  ![fig1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/diffusion_landscape.png)
350
+
351
  ## How to use it?\n
352
  A user can either design a custom input sequence to diffuse from or specify a length below. To scaffold a sequence use the following format where X represent residues to diffuse: XXXXXXXXSCIENCESCIENCEXXXXXXXXXXXXXXXXXXX. You can even design a protein with your name XXXXXXXXXXXXNAMEHEREXXXXXXXXXXXXX!
353
+
354
  ### Acknowledgements\n
355
  Thank you to Simon Dürr and the Hugging Face team for setting us up with a community GPU grant!
356
  """)
357
+
358
  gr.Markdown("""
359
  ## Model in Action
360
  ![gif1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/seqdiff_anim_720p.gif)
361
  """)
362
 
363
+ with gr.Row(equal_height=False):
 
364
  with gr.Column():
365
  with gr.Tabs():
366
+ with gr.Tab("Inputs"):
367
+ gr.Markdown("## INPUTS")
368
  gr.Markdown("""#### Start Sequence
369
  Specify the protein length for complete unconditional generation, or scaffold a motif (or your name) using the custom sequence input""")
370
  seq_opt = gr.Radio(["protein length","custom sequence"], label="How would you like to specify the starting sequence?", value='protein length')
371
 
372
  sequence = gr.Textbox(label="custom sequence", lines=1, placeholder='AMINO ACIDS: A,C,D,E,F,G,H,I,K,L,M,N,P,Q,R,S,T,V,W,Y\n MASK TOKEN: X', visible=False)
373
  seq_len = gr.Slider(minimum=5.0, maximum=250.0, label="protein length", value=100, visible=True)
374
+
375
  seq_opt.change(fn=toggle_seq_input,
376
  inputs=[seq_opt],
377
+ outputs=[seq_len, sequence])
 
378
 
379
+ gr.Markdown("### Optional Parameters")
380
  with gr.Accordion(label='Secondary Structure',open=True):
381
+ gr.Markdown("Try changing the sliders or inputing explicit secondary structure conditioning for each residue")
382
  sec_str_opt = gr.Radio(["sliders","explicit"], label="How would you like to specify secondary structure?", value='sliders')
383
 
384
  secondary_structure = gr.Textbox(label="secondary structure", lines=1, placeholder='HELIX = H STRAND = S LOOP = L MASK = X(must be the same length as input sequence)', visible=False)
385
+
386
  with gr.Column():
387
  helix_bias = gr.Slider(minimum=0.0, maximum=0.05, label="helix bias", visible=True)
388
  strand_bias = gr.Slider(minimum=0.0, maximum=0.05, label="strand bias", visible=True)
389
  loop_bias = gr.Slider(minimum=0.0, maximum=0.20, label="loop bias", visible=True)
390
+
391
  sec_str_opt.change(fn=toggle_secondary_structure,
392
  inputs=[sec_str_opt],
393
+ outputs=[helix_bias,strand_bias,loop_bias,secondary_structure])
394
+
 
395
  with gr.Accordion(label='Amino Acid Compositional Bias',open=False):
396
+ gr.Markdown("Bias sequence composition for particular amino acids by specifying the one letter code followed by the fraction to bias. This can be input as a list for example: W0.2,E0.1")
397
  with gr.Row():
398
  aa_bias = gr.Textbox(label="aa bias", lines=1, placeholder='specify one letter AA and fraction to bias, for example W0.1 or M0.1,K0.1' )
399
  aa_bias_potential = gr.Textbox(label="aa bias scale", lines=1, placeholder='AA Bias potential scale (recomended range 1.0-5.0)')
 
 
 
 
 
 
 
 
 
400
 
401
  with gr.Accordion(label='Hydrophobic Bias',open=False):
402
+ gr.Markdown("Bias for or against hydrophobic composition, to get more soluble proteins, bias away with a negative target score (ex. -5)")
403
  with gr.Row():
404
  hydrophobic_target_score = gr.Textbox(label="hydrophobic score", lines=1, placeholder='hydrophobic score to target (negative score is good for solublility)')
405
  hydrophobic_potential = gr.Textbox(label="hydrophobic potential scale", lines=1, placeholder='hydrophobic potential scale (recomended range 1.0-2.0)')
406
+
407
  with gr.Accordion(label='Diffusion Params',open=False):
408
+ gr.Markdown("Increasing T to more steps can be helpful for harder design challenges, sampling from different distributions can change the sequence and structural composition")
409
  with gr.Row():
410
  num_steps = gr.Textbox(label="T", lines=1, placeholder='number of diffusion steps (25 or less will speed things up)')
411
+ noise = gr.Dropdown(['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'], label='noise type', value='normal')
412
+
413
+ with gr.Tab("Motif Selection"):
414
 
415
+ gr.Markdown("### Motif Selection Preview")
416
  gr.Markdown('Contigs explained: to grab residues (seq and str) on a pdb chain you will provide the chain letter followed by a range of residues as indexed in the pdb file for example (A3-10) is the syntax to select residues 3-10 on chain A (the chain always needs to be specified). To add diffused residues to either side of this motif you can specify a range or discrete value without a chain letter infront. To add 15 residues before the motif and 20-30 residues (randomly sampled) after use the following syntax: 15,A3-10,20-30 commas are used to separate regions selected from the pdb and designed (diffused) resiudes which will be added. ')
417
  pdb_id_code = gr.Textbox(label="PDB ID", lines=1, placeholder='INPUT PDB ID TO FETCH (ex. 1DPX)', visible=True)
418
  contigs = gr.Textbox(label="contigs", lines=1, placeholder='specify contigs to grab particular residues from pdb ()', visible=True)
 
424
  rewrite_pdb = gr.File(label='PDB file')
425
  preview_btn = gr.Button("Preview Motif")
426
 
427
+ with gr.Tab("MSA to PSSM"):
428
+ gr.Markdown("### MSA to PSSM Generation")
429
  gr.Markdown('input either an MSA or PSSM to guide the model toward generating samples within your family of interest')
430
  with gr.Row():
431
  fasta_msa = gr.File(label='MSA')
 
437
 
438
  btn = gr.Button("GENERATE")
439
 
 
440
  with gr.Column():
441
+ gr.Markdown("## OUTPUTS")
442
+ gr.Markdown("#### Confidence score for generated structure at each timestep")
443
  plddt_plot = gr.Plot(label='plddt at step t')
444
+ gr.Markdown("#### Output protein sequence")
445
  output_seq = gr.Textbox(label="sequence")
446
+ gr.Markdown("#### Download PDB file")
447
  output_pdb = gr.File(label="PDB file")
448
+ gr.Markdown("#### Structure viewer")
449
  output_viewer = gr.HTML()
450
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
451
  preview_btn.click(get_motif_preview,[pdb_id_code, contigs],[preview_viewer, rewrite_pdb])
452
 
453
  pssm_gen_btn.click(get_pssm,[fasta_msa,input_pssm],[pssm_view, pssm])
454
 
455
+ btn.click(protein_diffusion_model,
456
+ [sequence,
457
+ seq_len,
458
+ helix_bias,
459
+ strand_bias,
460
+ loop_bias,
461
+ secondary_structure,
462
+ aa_bias,
463
+ aa_bias_potential,
464
+ num_steps,
465
+ noise,
466
+ hydrophobic_target_score,
 
 
 
467
  hydrophobic_potential,
468
  contigs,
469
  pssm,
470
+ seq_mask,
471
  str_mask,
472
+ rewrite_pdb],
473
  [output_seq,
474
  output_pdb,
475
  output_viewer,
476
  plddt_plot])
477
 
478
+ demo.launch(
479
+ debug=True,
480
+ theme=gr.themes.Soft(),
481
+ )
 
requirements.txt CHANGED
@@ -1,10 +1,10 @@
1
- torch==2.0.0
2
  e3nn==0.3.3
3
- dgl==1.0.2
4
  pynvml==11.0.0
5
  decorator==5.1.0
6
  icecream
7
  biopython
8
  py3Dmol
9
- pydantic<2
10
- numpy==1.24.1
 
1
+ torch==2.5.0
2
  e3nn==0.3.3
3
+ dgl==2.2.1
4
  pynvml==11.0.0
5
  decorator==5.1.0
6
  icecream
7
  biopython
8
  py3Dmol
9
+ pydantic>=2.0
10
+ numpy<2