Update app.py

app.py

@@ -545,18 +545,19 @@ def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bia
  
     return output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)
 
-
 def get_plddt_plot(plddt_data, max_t):
-    x = [i+1 for i in range(len(plddt_data))]
     fig, ax = plt.subplots(figsize=(15,6))
-    ax.plot(x,plddt_data,color='#661dbf', linewidth=3,marker='o')
-    ax.set_xticks([i+1 for i in range(max_t)])
-    ax.set_yticks([(i+1)/10 for i in range(10)])
-    ax.set_ylim([0,1])
+    x = list(range(1, len(plddt_data) + 1))
+    ax.plot(x, plddt_data, color='#661dbf', linewidth=3, marker='o')
+    ax.set_xticks(range(1, max_t + 1))
+    ax.set_yticks([i/10 for i in range(11)])  # 0부터 1까지
+    ax.set_ylim([0, 1])
     ax.set_ylabel('model confidence (plddt)')
     ax.set_xlabel('diffusion steps (t)')
+    plt.close()  # 메모리 관리를 위해 닫기
     return fig
 
+
 def display_pdb(path_to_pdb):
     '''
         #function to display pdb in py3dmol
@@ -936,15 +937,25 @@ def process_chat_and_generate(message, history):
             rewrite_pdb=None
         )
         
+        
         # 3. 모든 단계의 결과 수집
         final_result = None
+        plddt_data = []  # plddt 데이터를 저장할 리스트 추가
+        
         for result in generator:
             final_result = result
-            # 중간 결과 업데이트 (선택적)
+            if result[3]:  # plddt_plot이 존재하는 경우
+                # matplotlib 그래프에서 데이터 추출
+                ax = result[3].gca()
+                if ax.lines:
+                    line = ax.lines[0]
+                    plddt_data = line.get_ydata().tolist()
+            
+            # 중간 결과 업데이트
             yield (
                 history + [
                     {"role": "user", "content": message},
-                    {"role": "assistant", "content": f"단백질 생성 중... {len(result[3].get_lines()) if result[3] else 0}단계 완료"}
+                    {"role": "assistant", "content": f"단백질 생성 중... {len(plddt_data)}단계 완료"}
                 ],
                 create_radar_chart(calculate_hero_stats(
                     params['helix_bias'],
@@ -952,7 +963,7 @@ def process_chat_and_generate(message, history):
                     params['loop_bias'],
                     params['hydrophobic_target_score']
                 )),
-                f"단백질 생성 진행 중... {len(result[3].get_lines()) if result[3] else 0}단계 완료",
+                f"단백질 생성 진행 중... {len(plddt_data)}단계 완료",
                 result[0],  # output_seq
                 result[1],  # output_pdb
                 result[2],  # structure_view
@@ -964,14 +975,6 @@ def process_chat_and_generate(message, history):
             
         output_seq, output_pdb, structure_view, plddt_plot = final_result
         
-        # 4. 히어로 능력치 계산
-        stats = calculate_hero_stats(
-            helix_bias=params['helix_bias'],
-            strand_bias=params['strand_bias'],
-            loop_bias=params['loop_bias'],
-            hydrophobic_score=params['hydrophobic_target_score']
-        )
-        
         # 5. 최종 설명 생성
         explanation = f"""
         요청하신 기능에 맞는 단백질을 생성했습니다:
@@ -986,7 +989,8 @@ def process_chat_and_generate(message, history):
         - 루프 구조 비율: {params['loop_bias']*100:.1f}%
         - 소수성 점수: {params['hydrophobic_target_score']}
         
-        생성이 완료되었습니다. 총 {len(plddt_plot.get_lines()) if plddt_plot else 0}단계의 최적화를 수행했습니다.
+        생성이 완료되었습니다. 총 {len(plddt_data)}단계의 최적화를 수행했습니다.
+        최종 안정성 점수: {np.mean(plddt_data):.2f}
         """
         
         # 6. 최종 결과 반환
@@ -1005,6 +1009,7 @@ def process_chat_and_generate(message, history):
         
     except Exception as e:
         print(f"Error in process_chat_and_generate: {str(e)}")
+        traceback.print_exc()  # 상세한 에러 정보 출력
         return (
             history + [
                 {"role": "user", "content": message},