Update app.py

app.py

@@ -925,7 +925,7 @@ def process_chat_and_generate(message, history):
             secondary_structure=None,
             aa_bias=None,
             aa_bias_potential=None,
-            num_steps="25",
+            num_steps="25",  # 25단계 설정
             noise="normal",
             hydrophobic_target_score=str(params['hydrophobic_target_score']),
             hydrophobic_potential="2",
@@ -936,8 +936,33 @@ def process_chat_and_generate(message, history):
             rewrite_pdb=None
         )
         
-        # 3. 생성 결과 가져오기
-        output_seq, output_pdb, structure_view, plddt_plot = next(generator)
+        # 3. 모든 단계의 결과 수집
+        final_result = None
+        for result in generator:
+            final_result = result
+            # 중간 결과 업데이트 (선택적)
+            yield (
+                history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": f"단백질 생성 중... {len(result[3].get_lines()) if result[3] else 0}단계 완료"}
+                ],
+                create_radar_chart(calculate_hero_stats(
+                    params['helix_bias'],
+                    params['strand_bias'],
+                    params['loop_bias'],
+                    params['hydrophobic_target_score']
+                )),
+                f"단백질 생성 진행 중... {len(result[3].get_lines()) if result[3] else 0}단계 완료",
+                result[0],  # output_seq
+                result[1],  # output_pdb
+                result[2],  # structure_view
+                result[3]   # plddt_plot
+            )
+        
+        if final_result is None:
+            raise Exception("생성 결과가 없습니다")
+            
+        output_seq, output_pdb, structure_view, plddt_plot = final_result
         
         # 4. 히어로 능력치 계산
         stats = calculate_hero_stats(
@@ -947,7 +972,7 @@ def process_chat_and_generate(message, history):
             hydrophobic_score=params['hydrophobic_target_score']
         )
         
-        # 5. 설명 생성
+        # 5. 최종 설명 생성
         explanation = f"""
         요청하신 기능에 맞는 단백질을 생성했습니다:
         
@@ -960,23 +985,22 @@ def process_chat_and_generate(message, history):
         - 베타 시트 비율: {params['strand_bias']*100:.1f}%
         - 루프 구조 비율: {params['loop_bias']*100:.1f}%
         - 소수성 점수: {params['hydrophobic_target_score']}
-        """
         
-        # 6. 챗봇 히스토리 업데이트
-        new_history = history + [
-            {"role": "user", "content": message},
-            {"role": "assistant", "content": explanation}
-        ]
+        생성이 완료되었습니다. 총 {len(plddt_plot.get_lines()) if plddt_plot else 0}단계의 최적화를 수행했습니다.
+        """
         
-        # 7. 모든 결과 반환
+        # 6. 최종 결과 반환
         return (
-            new_history,  # 챗봇 히스토리
-            create_radar_chart(stats),  # 능력치 차트
-            explanation,  # 설명
-            output_seq,  # 단백질 서열
-            output_pdb,  # PDB 파일
-            structure_view,  # 3D 구조
-            plddt_plot  # 신뢰도 플롯
+            history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": explanation}
+            ],
+            create_radar_chart(stats),
+            explanation,
+            output_seq,
+            output_pdb,
+            structure_view,
+            plddt_plot
         )
         
     except Exception as e:
@@ -1495,9 +1519,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
         queue=False
     )
 
-
-
-    
     # 챗봇 응답에 따른 결과 업데이트
     msg.submit(
         update_protein_display,
@@ -1514,6 +1535,26 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
               [chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
+    # 이벤트 핸들러 수정
+    msg.submit(
+        fn=process_chat_and_generate,
+        inputs=[msg, chatbot],
+        outputs=[
+            chatbot,
+            hero_stats,
+            hero_description,
+            output_seq,
+            output_pdb,
+            output_viewer,
+            plddt_plot
+        ],
+        show_progress=True
+    ).then(
+        lambda: None,  # 진행 상태 초기화
+        None,
+        None
+    )
+
 # 실행
 demo.queue()
 demo.launch(debug=True)