Update app.py

app.py

@@ -441,8 +441,8 @@ def toggle_secondary_structure(choice):
 
 def generate_hero(name, strength, flexibility, speed, defense, size, abilities):
     try:
-        # protein_diffusion_model 호출
-        results = protein_diffusion_model(
+        # protein_diffusion_model 호출하여 마지막 결과 얻기
+        generator = protein_diffusion_model(
             sequence=None,
             seq_len=size,
             helix_bias=flexibility,
@@ -462,16 +462,27 @@ def generate_hero(name, strength, flexibility, speed, defense, size, abilities):
             rewrite_pdb=None
         )
         
+        # 마지막 결과만 가져오기
+        final_result = None
+        for result in generator:
+            final_result = result
+            
+        if final_result is None:
+            raise Exception("생성 결과가 없습니다")
+        
+        output_seq, output_pdb, structure_view, plddt_plot = final_result
+            
         # 능력치 계산
         stats = calculate_hero_stats(flexibility, strength, speed, defense)
         
         # 결과 반환
         return (
-            results[2],  # 3D 구조
+            structure_view,  # 3D 구조
             create_radar_chart(stats),  # 능력치 차트
             generate_hero_description(name, stats, abilities)  # 히어로 설명
         )
     except Exception as e:
+        print(f"Error in generate_hero: {str(e)}")  # 디버깅을 위한 출력 추가
         return (
             gr.HTML("에러가 발생했습니다"),
             None,