Update app.py

app.py

@@ -720,6 +720,130 @@ def build_advanced_character_pack(
         return [], None, None, None, traceback.format_exc()
 
 
+@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+def build_two_character_pack(
+    character_0_ref,
+    character_1_ref,
+    driving_video,
+    prompt,
+    sam3_text,
+    progress=gr.Progress(track_tqdm=True),
+):
+    try:
+        progress(0.0, desc="Checking SCAIL-Pose repo")
+        _require_repo_layout()
+        progress(0.04, desc="Preparing SAM3 weights")
+        sam3_model = _ensure_sam3_weights()
+
+        job_dir = _new_job_dir("two_character_animation")
+        pack_root = job_dir / "scail2_input_pack"
+        pack_root.mkdir(parents=True, exist_ok=True)
+
+        progress(0.08, desc="Preparing two-character inputs")
+        work_driving = job_dir / "driving.mp4"
+        _copy_video(driving_video, work_driving, "driving video")
+        shutil.copy2(work_driving, pack_root / "rendered_v2.mp4")
+
+        prepared_refs = [
+            CharacterRef(
+                character="character_0",
+                view="front",
+                path=_copy_reference_to_png(_as_path(character_0_ref, "character 0 reference"), pack_root / "characters" / "character_0" / "front.png"),
+            ),
+            CharacterRef(
+                character="character_1",
+                view="front",
+                path=_copy_reference_to_png(_as_path(character_1_ref, "character 1 reference"), pack_root / "characters" / "character_1" / "front.png"),
+            ),
+        ]
+        character_names = ["character_0", "character_1"]
+        primary = prepared_refs[0]
+        _build_advanced_metadata(pack_root, "animation", prompt, primary, character_names, sam3_text)
+
+        progress(0.12, desc="Loading SAM3 predictors")
+        video_predictor, image_predictor = _load_sam3_predictors(sam3_model)
+
+        from TrackSam3.track import get_mask_from_image, get_mask_from_image_via_video, get_mask_from_video
+        from NLFPoseExtract.v2_helper import save_colored_mask_image, write_colored_mask_video
+
+        text = _text_args(sam3_text)
+        colors = SCAIL_COLORS[:2]
+
+        progress(0.20, desc="Tracking two subjects left-to-right")
+        drv_masks, drv_colors = get_mask_from_video(
+            str(work_driving),
+            video_predictor,
+            max_targets=2,
+            sort_by="x",
+            fixed_colors=colors,
+            text=text,
+        )
+        if len(drv_masks) < 2:
+            raise RuntimeError(
+                f"SAM3 detected {len(drv_masks)} subject(s), but this mode expects two. "
+                "Try a clearer driving video or use the Advanced Character Pack."
+            )
+
+        import cv2
+
+        cap = cv2.VideoCapture(str(work_driving))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        cap.release()
+        fps_int = max(1, int(round(fps or 24)))
+        write_colored_mask_video(
+            drv_masks[:2],
+            drv_colors[:2],
+            str(pack_root / "rendered_mask_v2.mp4"),
+            fps_int,
+            bg_color=(255, 255, 255),
+        )
+
+        progress(0.55, desc="Generating reference masks")
+        logs = [
+            "Mode: two-character animation",
+            "Mapping: left-most driving subject -> character_0; next subject to the right -> character_1",
+            f"Driving tracks detected: {len(drv_masks)}",
+        ]
+        for ref in prepared_refs:
+            character_idx = character_names.index(ref.character)
+            color = drv_colors[character_idx]
+            mask_path = ref.path.with_name(f"{ref.view}_mask.png")
+            progress(None, desc=f"Masking {ref.character}")
+            try:
+                ref_masks, _ = get_mask_from_image_via_video(
+                    str(ref.path),
+                    video_predictor,
+                    max_targets=1,
+                    sort_by="x",
+                    fixed_colors=[color],
+                    text=text,
+                )
+            except Exception as exc:
+                logging.warning("Video-based image masking failed for %s, trying image predictor: %s", ref.path, exc)
+                ref_masks, _ = get_mask_from_image(
+                    str(ref.path),
+                    image_predictor,
+                    max_targets=1,
+                    sort_by="x",
+                    fixed_colors=[color],
+                    text=text,
+                )
+            if len(ref_masks) == 0:
+                raise RuntimeError(f"SAM3 did not detect a subject in {ref.character}.")
+            save_colored_mask_image([ref_masks[0]], [color], str(mask_path), bg_color=(255, 255, 255))
+            logs.append(f"- {ref.character}: {mask_path.relative_to(pack_root).as_posix()}")
+
+        progress(0.92, desc="Packaging two-character pack")
+        zip_path = _write_canonical_pack_zip(pack_root, "two_character_animation")
+        gallery = _advanced_gallery(pack_root)
+        progress(1.0, desc="Done")
+        status = f"Done. Two-character pack created at {zip_path}\n\n" + "\n".join(logs)
+        return gallery, str(pack_root / "rendered_v2.mp4"), str(pack_root / "rendered_mask_v2.mp4"), str(zip_path), status
+    except Exception:
+        logging.exception("Two-character pack generation failed")
+        return [], None, None, None, traceback.format_exc()
+
+
 def build_ui():
     with gr.Blocks(title="SCAIL-Pose Pack Builder") as demo:
         gr.Markdown(
@@ -788,6 +912,48 @@ def build_ui():
                 outputs=[simple_ref_mask, simple_rendered, simple_mask_video, simple_zip, simple_status],
             )
 
+        with gr.Tab("Two Characters"):
+            gr.Markdown(
+                "Fast path for the common two-character case. Upload two reference images and one driving video. "
+                "`Character 0` should match the left-most subject in the driving video; `Character 1` should match "
+                "the next subject to the right."
+            )
+            with gr.Row():
+                two_char0 = gr.Image(type="filepath", label="Character 0 reference")
+                two_char1 = gr.Image(type="filepath", label="Character 1 reference")
+            two_driving = gr.Video(label="Driving video")
+            with gr.Accordion("Optional settings", open=False):
+                two_prompt = gr.Textbox(label="Prompt for SCAIL-2", lines=3)
+                two_sam3_text = gr.Textbox(
+                    value="human character",
+                    label="SAM3 text prompt",
+                    info="Use a broad prompt such as `human character`, or comma/newline-separated prompts for distinct subjects.",
+                )
+            two_run = gr.Button("Generate two-character pack", variant="primary")
+            two_gallery = gr.Gallery(
+                label="References and generated masks",
+                columns=4,
+                height=320,
+                selected_index=0,
+                preview=True,
+            )
+            with gr.Row():
+                two_rendered = gr.Video(label="Rendered / driving video")
+                two_mask_video = gr.Video(label="Driving mask video")
+            two_zip = gr.File(label="Download two-character SCAIL-2 pack")
+            two_status = gr.Textbox(label="Run logs", lines=14)
+            two_run.click(
+                build_two_character_pack,
+                inputs=[
+                    two_char0,
+                    two_char1,
+                    two_driving,
+                    two_prompt,
+                    two_sam3_text,
+                ],
+                outputs=[two_gallery, two_rendered, two_mask_video, two_zip, two_status],
+            )
+
         with gr.Tab("Advanced Character Pack"):
             gr.Markdown(
                 "Create a canonical SCAIL-2 animation pack for multi-reference or multi-character cases. "
@@ -869,8 +1035,8 @@ def build_ui():
 
         with gr.Tab("Pack Format"):
             gr.Markdown(
-                "The Simple Pack tab exports flat packs. The Advanced Character Pack tab exports a canonical "
-                "character pack with explicit character folders.\n\n"
+                "The Simple Pack tab exports flat packs. The Two Characters and Advanced Character Pack tabs "
+                "export canonical character packs with explicit character folders.\n\n"
                 "Simple animation pack:\n"
                 "```text\n"
                 "ref.png\n"
@@ -913,4 +1079,4 @@ def build_ui():
 
 
 if __name__ == "__main__":
-    build_ui().queue(max_size=4).launch(show_error=True)
+    build_ui().queue(max_size=4).launch(show_error=True)