segment length awareness

app.py

@@ -64,6 +64,8 @@ COPY_SOURCE_AUDIO = os.getenv("SCAIL_COPY_SOURCE_AUDIO", "1") == "1"
 MATCH_SOURCE_FPS_AT_EXPORT = os.getenv("SCAIL_MATCH_SOURCE_FPS_AT_EXPORT", "1") == "1"
 CONFORM_OUTPUT_TO_SOURCE_FPS = os.getenv("SCAIL_CONFORM_OUTPUT_TO_SOURCE_FPS", "1") == "1"
 FPS_MATCH_EPSILON = float(os.getenv("SCAIL_FPS_MATCH_EPSILON", "0.05"))
+AUTO_AVOID_SEGMENT_TRUNCATION = os.getenv("SCAIL_AUTO_AVOID_SEGMENT_TRUNCATION", "1") == "1"
+MAX_AUTO_SEGMENT_LEN = int(os.getenv("SCAIL_MAX_AUTO_SEGMENT_LEN", "161"))
 
 CLIP_CKPT_NAME = "models_clip_open-clip-xlm-roberta-large-vit-huge-14-onlyvisual.pth"
 ORIGINAL_DIT_REL_PATH = "model/1/fsdp2_rank_0000_checkpoint.pt"
@@ -1001,6 +1003,50 @@ def _video_fps(video_path: str | Path) -> float | None:
     return None
 
 
+def _video_frame_count(video_path: str | Path) -> int | None:
+    try:
+        import cv2
+
+        capture = cv2.VideoCapture(str(video_path))
+        frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+        capture.release()
+        if frame_count > 0:
+            return frame_count
+    except Exception as exc:
+        logging.warning("Could not read frame count for %s: %s", video_path, exc)
+    return None
+
+
+def _effective_segment_len(requested_segment_len: int, segment_overlap: int, source_video_path: str | Path) -> int:
+    requested_segment_len = int(requested_segment_len)
+    segment_overlap = int(segment_overlap)
+    if not AUTO_AVOID_SEGMENT_TRUNCATION:
+        return requested_segment_len
+
+    frame_count = _video_frame_count(source_video_path)
+    if frame_count is None or frame_count <= requested_segment_len:
+        return requested_segment_len
+
+    if frame_count <= MAX_AUTO_SEGMENT_LEN:
+        effective = max(frame_count, segment_overlap + 1)
+        logging.info(
+            "Increasing segment_len to avoid SCAIL-2 tail truncation: requested=%s frame_count=%s effective=%s",
+            requested_segment_len,
+            frame_count,
+            effective,
+        )
+        return effective
+
+    logging.warning(
+        "Source has %s frames, but segment_len=%s and max_auto_segment_len=%s. "
+        "SCAIL-2 may drop the final partial segment unless you increase segment_len or trim/resample the input.",
+        frame_count,
+        requested_segment_len,
+        MAX_AUTO_SEGMENT_LEN,
+    )
+    return requested_segment_len
+
+
 def _generation_config_for_source_video(source_video_path: str | Path):
     cfg = copy.deepcopy(_SCAIL_CONFIGS[MODEL_NAME])
     if not MATCH_SOURCE_FPS_AT_EXPORT:
@@ -1189,6 +1235,7 @@ def _run_scail_job(
     pipeline = _get_pipeline()
     cfg = _generation_config_for_source_video(pose_path)
     save_file = OUTPUT_DIR / f"scail2_{uuid.uuid4().hex}.mp4"
+    effective_segment_len = _effective_segment_len(segment_len, segment_overlap, pose_path)
 
     if progress is not None:
         progress(0.12, desc="Preparing inputs")
@@ -1197,7 +1244,7 @@ def _run_scail_job(
         target_w=int(target_w),
         sample_shift=float(sample_shift),
         sample_solver=DEFAULT_SOLVER,
-        segment_len=int(segment_len),
+        segment_len=effective_segment_len,
         segment_overlap=int(segment_overlap),
         sample_steps=int(sample_steps),
         sample_guide_scale=float(guide_scale),