Update app.py

app.py

@@ -285,16 +285,15 @@ def display_pdb(path_to_pdb):
 '''
 
 
-
-# MOTIF SCAFFOLDING
 def get_motif_preview(pdb_id, contigs):
-    '''
-        #function to display selected motif in py3dmol
-    '''
-    input_pdb = fetch_pdb(pdb_id=pdb_id.lower())
+    try:
+        input_pdb = fetch_pdb(pdb_id=pdb_id.lower() if pdb_id else None)
+        if input_pdb is None:
+            return gr.HTML("PDB ID를 입력해주세요"), None
+        
+        parse = parse_pdb(input_pdb)
+        
 
-    # rewrite pdb
-    parse = parse_pdb(input_pdb)
     #output_name = './rewrite_'+input_pdb.split('/')[-1]
     #writepdb(output_name, torch.tensor(parse_og['xyz']),torch.tensor(parse_og['seq']))
     #parse = parse_pdb(output_name)
@@ -330,6 +329,9 @@ def get_motif_preview(pdb_id, contigs):
                             allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
                             allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>""", output_name
 
+  except Exception as e:
+        return gr.HTML(f"오류가 발생했습니다: {str(e)}"), None
+
 def fetch_pdb(pdb_id=None):
     if pdb_id is None or pdb_id == "":
         return None
@@ -402,20 +404,20 @@ def msa_to_pssm(msa_file):
     return outdir
 
 def get_pssm(fasta_msa, input_pssm):
+    try:
+        if input_pssm is not None:
+            outdir = input_pssm.name
+        elif fasta_msa is not None:
+            outdir = save_pssm(fasta_msa)
+        else:
+            return gr.Plot(label="파일을 업로드해주세요"), None
 
-    if input_pssm not in ['',None]:
-        outdir = input_pssm.name
-    else:
-        outdir = save_pssm(fasta_msa)
-
-    pssm = np.loadtxt(outdir, delimiter=",", dtype=float)
-    fig, ax = plt.subplots(figsize=(15,6))
-    plt.imshow(torch.permute(torch.tensor(pssm),(1,0)))
-
-    return fig, outdir
-
-
-
+        pssm = np.loadtxt(outdir, delimiter=",", dtype=float)
+        fig, ax = plt.subplots(figsize=(15,6))
+        plt.imshow(torch.permute(torch.tensor(pssm),(1,0)))
+        return fig, outdir
+    except Exception as e:
+        return gr.Plot(label=f"오류가 발생했습니다: {str(e)}"), None
 
 # 히어로 능력치 계산 함수 추가
 def calculate_hero_stats(helix_bias, strand_bias, loop_bias, hydrophobic_score):
@@ -682,6 +684,10 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                                 visible=True
                             )
 
+
+                            design_btn = gr.Button("🧬 DNA 설계 생성!", variant="primary", scale=2)
+
+                    
                     # 아미노산 구성 설정 추가
                     with gr.Accordion(label='🧬 DNA 구성 설정', open=False):
                         gr.Markdown("""
@@ -771,6 +777,10 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                     rewrite_pdb = gr.File(label='히어로 DNA 파일')
                     preview_btn = gr.Button("🔍 미리보기", variant="secondary")
 
+
+                        enhance_btn = gr.Button("⚡ 강화된 히어로 생성!", variant="primary", scale=2)
+
+                    
                 with gr.TabItem("👑 히어로 가문"):
                     gr.Markdown("""
                     ### 🏰 위대한 히어로 가문의 유산
@@ -789,6 +799,11 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
                     pssm_view = gr.Plot(label='가문 특성 분석 결과')
                     pssm_gen_btn = gr.Button("✨ 가문 특성 분석", variant="secondary")
 
+                    inherit_btn = gr.Button("👑 가문의 힘 계승!", variant="primary", scale=2)
+
+
+
+        
         with gr.Column():
             gr.Markdown("## 🦸‍♂️ 히어로 프로필")
             
@@ -823,11 +838,30 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
         queue=False
     )
 
-    # 나머지 이벤트 연결은 그대로 유지
+
+
+# 이벤트 연결
     preview_btn.click(get_motif_preview,[pdb_id_code, contigs],[preview_viewer, rewrite_pdb])
     pssm_gen_btn.click(get_pssm,[fasta_msa,input_pssm],[pssm_view, pssm])
 
+# 각 탭의 생성 버튼 연결
+    design_btn.click(
+        combined_generation,
+        inputs=[...],  # DNA 설계에 필요한 입력값들
+        outputs=[...]  # 필요한 출력값들
+    )
 
+    enhance_btn.click(
+        combined_generation,
+        inputs=[...],  # 유전자 강화에 필요한 입력값들
+        outputs=[...]  # 필요한 출력값들
+    )
+
+    inherit_btn.click(
+        combined_generation,
+        inputs=[...],  # 가문 계승에 필요한 입력값들
+        outputs=[...]  # 필요한 출력값들
+    )
 
     # generate_hero와 protein_diffusion_model을 combined_generation으로 통합
     create_btn.click(