:tada: init

app.py

@@ -277,7 +277,7 @@ def inference(
         if has_include_special_tokens:
             mean_cross_attn_probs = mean_cross_attn_probs[:n_cond_tokens, ...]  # (n_tokens, 512, 512)
         else:
-            mean_cross_attn_probs = mean_cross_attn_probs[1:n_cond_tokens - 1, ...]
+            mean_cross_attn_probs = mean_cross_attn_probs[1:n_cond_tokens - 1, ...]  # (n_tokens-2, 512, 512)
 
         cross_attention_probs_list.append(mean_cross_attn_probs)
 
@@ -290,10 +290,7 @@ def inference(
     image_list = []
     # 各行ごとに画像を作成し保存
     for i in tqdm(range(cross_attention_probs.shape[0]), desc="Saving images..."):
-        if has_include_special_tokens:
-            fig, ax = plt.subplots(1, n_cond_tokens, figsize=(16, 4))
-        else:
-            fig, ax = plt.subplots(1, n_cond_tokens - 2, figsize=(16, 4))
+        fig, ax = plt.subplots(1, n_cond_tokens, figsize=(16, 4))
 
         for j in range(cross_attention_probs.shape[1]):
             # 各クラスのアテンションマップを Min-Max 正規化 (0~1)