feat : decoder 의 마지막 layer 에서만 heatmap 생성

app.py

@@ -345,25 +345,22 @@ def predict_captioning(image):
     )
 
     tmp_dir = tempfile.mkdtemp(prefix="combined_captioning_gradio_")
-    heatmap_images = []
-    n_layers = len(runtime["decoder"].layers)
-
-    # 각 decoder layer별 cross-attention heatmap 이미지를 만들어 Gallery에 표시한다.
-    for layer in range(1, n_layers + 1):
-        cross_atten_path = Path(tmp_dir) / f"cross_attention_layer_{layer}.jpg"
-        runtime["decoder"].show_cross_atten(
-            enc_dec_atten[0],
-            caption_tokens,
-            layer,
-            image_tensor.squeeze(0).detach().cpu(),
+    last_layer = len(runtime["decoder"].layers)
+    cross_atten_path = Path(tmp_dir) / "cross_attention_last_layer.jpg"
+
+    runtime["decoder"].show_cross_atten(
+        enc_dec_atten[0],
+        caption_tokens,
+        last_layer,
+        image_tensor.squeeze(0).detach().cpu(),
+        str(cross_atten_path),
+    )
+    heatmap_images = [
+        (
             str(cross_atten_path),
+            f"Last Layer ({last_layer})",
         )
-        heatmap_images.append(
-            (
-                str(cross_atten_path),
-                f"Layer {layer}",
-            )
-        )
+    ]
 
     return caption, heatmap_images