added selection method

app.py

@@ -37,6 +37,7 @@ pipeline = None
 speaker_encoder = None
 facenet_model = None
 facenet_classify_model = None
+mtcnn_model = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
@@ -347,13 +348,43 @@ def select_best_images(images: list, n: int) -> list:
     return [images[i] for i in top_indices]
 
 
+def select_best_images_combined(images: list, n: int) -> list:
+    global mtcnn_model, facenet_classify_model
+
+    n = min(n, len(images))
+    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)
+
+    scores = []
+    for idx, img in enumerate(images):
+        _, probs = mtcnn_model.detect(img)
+        det_conf = float(probs[0]) if probs is not None and probs[0] is not None else 0.0
+
+        tensor = _facenet_transform()(img.convert("RGB")).unsqueeze(0)
+        with torch.no_grad():
+            logits = facenet_classify_model(tensor)
+        classify_conf = float(F.softmax(logits, dim=1).max(dim=1).values[0])
+
+        combined = det_conf * classify_conf
+        scores.append(combined)
+        print(f"  [combined] idx={idx} det={det_conf:.3f} classify={classify_conf:.3f} combined={combined:.3f}")
+
+    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]
+
+
+SELECTION_METHODS = ["Pairwise similarity", "Detection + Classify confidence"]
+
+
 # ---------------------------------------------------------------------------
 # Generation
 # ---------------------------------------------------------------------------
 INTERNAL_SAMPLES = 10
 
 @spaces.GPU(duration=120)
-def generate(audio_path, num_display, guidance_scale, num_inference_steps, base_seed):
+def generate(audio_path, num_display, guidance_scale, num_inference_steps, base_seed, selection_method="Pairwise similarity"):
     global pipeline, speaker_encoder, facenet_model, device
 
     if audio_path is None:
@@ -402,6 +433,8 @@ def generate(audio_path, num_display, guidance_scale, num_inference_steps, base_
 
             images.append(img)
 
+    if selection_method == "Detection + Classify confidence":
+        return select_best_images_combined(images, int(num_display)), ""
     return select_best_images(images, int(num_display)), ""
 
 # ---------------------------------------------------------------------------
@@ -409,7 +442,7 @@ def generate(audio_path, num_display, guidance_scale, num_inference_steps, base_
 # ---------------------------------------------------------------------------
 
 def load_models():
-    global pipeline, speaker_encoder, facenet_model, facenet_classify_model, device
+    global pipeline, speaker_encoder, facenet_model, facenet_classify_model, mtcnn_model, device
     dtype = torch.float16 if device == "cuda" else torch.float32
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -458,17 +491,19 @@ def load_models():
     pipeline = pipeline.to(device)
     print("  Pipeline ready")
 
-    # FaceNet for best-sample selection
-    print("Loading FaceNet for best-sample selection...")
+    # FaceNet + MTCNN for best-sample selection
+    print("Loading FaceNet + MTCNN for best-sample selection...")
     try:
-        from facenet_pytorch import InceptionResnetV1
+        from facenet_pytorch import InceptionResnetV1, MTCNN
         facenet_model = InceptionResnetV1(pretrained='vggface2', classify=False).eval()
         facenet_classify_model = InceptionResnetV1(pretrained='vggface2', classify=True).eval()
-        print("  FaceNet ready")
+        mtcnn_model = MTCNN(keep_all=False, device='cpu')
+        print("  FaceNet + MTCNN ready")
     except Exception as e:
-        print(f"  FaceNet unavailable ({e}); select-best will fall back to first image")
+        print(f"  FaceNet/MTCNN unavailable ({e}); select-best will fall back to first image")
         facenet_model = None
         facenet_classify_model = None
+        mtcnn_model = None
 
 
 # ---------------------------------------------------------------------------
@@ -491,20 +526,25 @@ def build_demo():
                 guidance_scale = gr.Slider(1.0, 10.0, value=2.5, step=0.5, label="Guidance scale")
                 num_steps = gr.Slider(10, 50, value=25, step=5, label="Inference steps")
                 base_seed = gr.Slider(0, 9999, value=42, step=1, label="Base seed")
+                selection_method = gr.Radio(
+                    choices=SELECTION_METHODS,
+                    value=SELECTION_METHODS[1],
+                    label="Best-image selection method",
+                )
                 generate_btn = gr.Button("Generate", variant="primary")
 
             with gr.Column():
                 gallery = gr.Gallery(label="Generated Images")
                 status = gr.Markdown(visible=False)
 
-        def _generate(audio, n, gs, steps, seed):
-            imgs, msg = generate(audio, n, gs, steps, seed)
+        def _generate(audio, n, gs, steps, seed, sel_method):
+            imgs, msg = generate(audio, n, gs, steps, seed, sel_method)
             visible = bool(msg)
             return imgs, gr.update(value=msg, visible=visible)
 
         generate_btn.click(
             fn=_generate,
-            inputs=[audio_input, num_display, guidance_scale, num_steps, base_seed],
+            inputs=[audio_input, num_display, guidance_scale, num_steps, base_seed, selection_method],
             outputs=[gallery, status],
         )