Update app.py

app.py

@@ -1,11 +1,8 @@
 import os,sys
-from openai import OpenAI
-import gradio as gr
-import json  # json 모듈 추가
 
 # install required packages
-os.system('pip install -q plotly')
-os.system('pip install -q matplotlib')
+os.system('pip install plotly')  # plotly 설치
+os.system('pip install matplotlib')  # matplotlib 설치
 os.system('pip install dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html')
 os.environ["DGLBACKEND"] = "pytorch"
 print('Modules installed')
@@ -32,11 +29,20 @@ if not os.path.exists('./tmp/args.json'):
     with open('./tmp/args.json', 'w') as f:
         json.dump(default_args, f)
 
-# args 로드
-with open('./tmp/args.json', 'r') as f:
-    args = json.load(f)
+# 체크포인트 파일 다운로드
+if not os.path.exists('./SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'):
+    print('Downloading model weights 1')
+    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt')
+    print('Successfully Downloaded')
+
+if not os.path.exists('./SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'):
+    print('Downloading model weights 2')
+    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt')
+    print('Successfully Downloaded')
 
-# 필수 라이브러리 임포트
+from openai import OpenAI
+import gradio as gr
+import json  # json 모듈 추가
 from datasets import load_dataset
 import plotly.graph_objects as go
 import numpy as np
@@ -52,10 +58,21 @@ from model.util import writepdb
 from utils.inpainting_util import *
 import os
 
-# 현재 스크립트의 디렉토리를 기준으로 체크포인트 파일 경로 설정
-current_dir = os.path.dirname(os.path.abspath(__file__))
-dssp_checkpoint = os.path.join(current_dir, 'SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt')
-og_checkpoint = os.path.join(current_dir, 'SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt')
+# args 로드
+with open('./tmp/args.json', 'r') as f:
+    args = json.load(f)
+
+plt.rcParams.update({'font.size': 13})
+
+# manually set checkpoint to load
+args['checkpoint'] = None
+args['dump_trb'] = False
+args['dump_args'] = True
+args['save_best_plddt'] = True
+args['T'] = 25
+args['strand_bias'] = 0.0
+args['loop_bias'] = 0.0
+args['helix_bias'] = 0.0
 
 # Hugging Face 토큰 설정
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
@@ -351,12 +368,15 @@ def generate_explanation(result, params):
     return explanation
 
 # 체크포인트 파일 경로를 절대 경로로 수정
-
-
 def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias, 
                     secondary_structure, aa_bias, aa_bias_potential, 
                     num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                     contigs, pssm, seq_mask, str_mask, rewrite_pdb):
+
+    
+    dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
+    og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'
+
     
     # 체크포인트 파일 존재 확인
     if not os.path.exists(dssp_checkpoint):