improved UI

app.py

@@ -333,31 +333,34 @@ def _extract_facenet_logits(img: Image.Image, model) -> torch.Tensor:
     return logits.squeeze(0)
 
 
-def select_best_images(images: list, n: int) -> list:
+def select_best_images(pairs: list, n: int) -> list:
+    """pairs: list of (image, seed). Returns top-n (image, seed) pairs."""
     global facenet_model
 
-    n = min(n, len(images))
+    n = min(n, len(pairs))
+    images = [p[0] for p in pairs]
     if facenet_model is None:
-        return images[:n]
+        return pairs[:n]
 
     embeddings = torch.stack([_extract_facenet_emb(img, facenet_model) for img in images])
     sim_matrix = F.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=2)
     avg_sims = (sim_matrix.sum(dim=1) - 1) / (len(images) - 1)
     top_indices = avg_sims.argsort(descending=True)[:n].tolist()
     print(f"[select_best:pairwise] top {n} indices={top_indices} avg_sims={avg_sims[top_indices].tolist()}")
-    return [images[i] for i in top_indices]
+    return [pairs[i] for i in top_indices]
 
 
-def select_best_images_combined(images: list, n: int) -> list:
+def select_best_images_combined(pairs: list, n: int) -> list:
+    """pairs: list of (image, seed). Returns top-n (image, seed) pairs."""
     global mtcnn_model, facenet_classify_model
 
-    n = min(n, len(images))
+    n = min(n, len(pairs))
     if mtcnn_model is None or facenet_classify_model is None:
         print("[select_best:combined] models unavailable, falling back to pairwise")
-        return select_best_images(images, n)
+        return select_best_images(pairs, n)
 
     scores = []
-    for idx, img in enumerate(images):
+    for idx, (img, _) in enumerate(pairs):
         _, probs = mtcnn_model.detect(img)
         det_conf = float(probs[0]) if probs is not None and probs[0] is not None else 0.0
 
@@ -372,7 +375,7 @@ def select_best_images_combined(images: list, n: int) -> list:
 
     top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:n]
     print(f"[select_best:combined] top {n} indices={top_indices} scores={[scores[i] for i in top_indices]}")
-    return [images[i] for i in top_indices]
+    return [pairs[i] for i in top_indices]
 
 
 SELECTION_METHODS = ["Pairwise similarity", "Detection + Classify confidence"]
@@ -421,7 +424,7 @@ def generate(audio_path, num_display, guidance_scale, num_inference_steps, base_
         id_emb = speech_z.to(dtype)
         id_emb_projected = project_face_embs(pipeline, id_emb)
 
-        images = []
+        pairs = []
         for seed in GENERATION_SEEDS:
             generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -433,11 +436,14 @@ def generate(audio_path, num_display, guidance_scale, num_inference_steps, base_
                 generator=generator,
             ).images[0]
 
-            images.append(img)
+            pairs.append((img, seed))
 
     if selection_method == "Detection + Classify confidence":
-        return select_best_images_combined(images, int(num_display)), ""
-    return select_best_images(images, int(num_display)), ""
+        best = select_best_images_combined(pairs, int(num_display))
+    else:
+        best = select_best_images(pairs, int(num_display))
+
+    return [(img, f"Seed: {seed}") for img, seed in best], ""
 
 # ---------------------------------------------------------------------------
 # Model loading