| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import gradio as gr |
|
|
|
|
| @dataclass(frozen=True) |
| class ChunkPlan: |
| control_start_frame: int |
| requested_frames: int |
| overlap_frames: int |
|
|
|
|
| @dataclass(frozen=True) |
| class FramePlanRules: |
| frame_step: int |
| minimum_requested_frames: int |
|
|
|
|
| class FramePlanningError(RuntimeError): |
| pass |
|
|
|
|
| def require_model_def(model_type: str, get_model_def) -> dict: |
| model_def = get_model_def(str(model_type)) |
| if not isinstance(model_def, dict): |
| raise gr.Error(f"Unsupported model type: {model_type}") |
| return model_def |
|
|
|
|
| def get_frame_plan_rules(model_type: str, get_model_def) -> FramePlanRules: |
| model_def = require_model_def(model_type, get_model_def) |
| return FramePlanRules(frame_step=int(model_def.get("frames_steps", 1)), minimum_requested_frames=int(model_def.get("frames_minimum", 1))) |
|
|
|
|
| def get_vae_temporal_latent_size(model_type: str, get_model_def) -> int: |
| model_def = require_model_def(model_type, get_model_def) |
| return int(model_def.get("latent_size", model_def.get("frames_steps", 1))) |
|
|
|
|
| def get_overlap_slider_max(model_type: str, get_model_def, *, exclusive_upper_bound: int = 100) -> int: |
| step = get_vae_temporal_latent_size(model_type, get_model_def) |
| last_allowed_value = int(exclusive_upper_bound) - 1 |
| return 1 + ((last_allowed_value - 1) // step) * step |
|
|
|
|
| def align_requested_frames(frame_count: int, *, frame_step: int, round_up: bool) -> int: |
| if frame_count <= 1: |
| return 1 |
| if round_up: |
| return int(math.ceil((frame_count - 1) / float(frame_step)) * frame_step + 1) |
| return int(math.floor((frame_count - 1) / float(frame_step)) * frame_step + 1) |
|
|
|
|
| def normalize_chunk_frames(chunk_seconds: float, fps_float: float, *, frame_step: int, minimum_requested_frames: int) -> int: |
| if chunk_seconds < 0.1: |
| raise FramePlanningError("Chunk Size must be at least 0.1 seconds.") |
| if fps_float <= 0.0: |
| raise FramePlanningError("Source FPS must be positive.") |
| target_frames = int(round(chunk_seconds * fps_float)) |
| if target_frames < minimum_requested_frames: |
| target_frames = minimum_requested_frames |
| below = align_requested_frames(target_frames, frame_step=frame_step, round_up=False) |
| if below < minimum_requested_frames: |
| below = minimum_requested_frames |
| above = align_requested_frames(target_frames, frame_step=frame_step, round_up=True) |
| if above < minimum_requested_frames: |
| above = minimum_requested_frames |
| return below if abs(below - target_frames) <= abs(above - target_frames) else above |
|
|
|
|
| def normalize_overlap_frames(overlap_frames: float, *, frame_step: int) -> int: |
| if overlap_frames < 1: |
| raise FramePlanningError("Sliding Window Overlap must be at least 1 frame.") |
| target_frames = int(round(float(overlap_frames))) |
| below = align_requested_frames(target_frames, frame_step=frame_step, round_up=False) |
| if below < 1: |
| below = 1 |
| above = align_requested_frames(target_frames, frame_step=frame_step, round_up=True) |
| if above < 1: |
| above = 1 |
| return below if abs(below - target_frames) <= abs(above - target_frames) else above |
|
|
|
|
| def align_total_unique_frames(total_unique_frames: int, *, frame_step: int, minimum_requested_frames: int, initial_overlap_frames: int) -> int: |
| if total_unique_frames <= 0: |
| return 0 |
| if initial_overlap_frames < 0: |
| raise FramePlanningError("Initial overlap cannot be negative.") |
| if initial_overlap_frames > 0: |
| minimum_unique_frames = minimum_requested_frames - initial_overlap_frames |
| if minimum_unique_frames < 1: |
| minimum_unique_frames = 1 |
| return 0 if total_unique_frames < minimum_unique_frames else total_unique_frames - (total_unique_frames % frame_step) |
| return 0 if total_unique_frames < minimum_requested_frames else ((total_unique_frames - 1) // frame_step) * frame_step + 1 |
|
|
|
|
| def count_planned_unique_frames(plans: list[ChunkPlan]) -> int: |
| return sum(plan.requested_frames - plan.overlap_frames for plan in plans) |
|
|
|
|
| def describe_frame_range(start_frame: int, frame_count: int) -> str: |
| if frame_count <= 0: |
| return "0 frame(s)" |
| return f"{frame_count} frame(s) [{start_frame}..{start_frame + frame_count - 1}]" |
|
|
|
|
| def build_chunk_plan( |
| start_frame: int, |
| end_frame_exclusive: int, |
| total_source_frames: int, |
| chunk_frames: int, |
| *, |
| frame_step: int, |
| minimum_requested_frames: int, |
| overlap_frames: int, |
| initial_overlap_frames: int = 0, |
| ) -> list[ChunkPlan]: |
| if chunk_frames < minimum_requested_frames: |
| raise FramePlanningError("Chunk size is below the model minimum frame count.") |
| if overlap_frames < 0: |
| raise FramePlanningError("Sliding Window Overlap cannot be negative.") |
| if initial_overlap_frames < 0: |
| raise FramePlanningError("Initial overlap cannot be negative.") |
| if overlap_frames >= chunk_frames: |
| raise FramePlanningError("Sliding Window Overlap must stay below the computed chunk size.") |
| if initial_overlap_frames >= chunk_frames: |
| raise FramePlanningError("Initial overlap must stay below the computed chunk size.") |
| plans: list[ChunkPlan] = [] |
| cursor = start_frame |
| total_unique_frames = align_total_unique_frames( |
| end_frame_exclusive - start_frame, |
| frame_step=frame_step, |
| minimum_requested_frames=minimum_requested_frames, |
| initial_overlap_frames=initial_overlap_frames, |
| ) |
| if total_unique_frames <= 0: |
| raise FramePlanningError("The selected range ends too close to the source video end to build a valid chunk for the current model.") |
| written_unique_frames = 0 |
| while written_unique_frames < total_unique_frames: |
| plan_overlap_frames = initial_overlap_frames if len(plans) == 0 else overlap_frames |
| remaining_unique = total_unique_frames - written_unique_frames |
| max_unique_frames = chunk_frames - plan_overlap_frames |
| requested_frames = chunk_frames if remaining_unique > max_unique_frames else remaining_unique + plan_overlap_frames |
| control_start_frame = cursor - plan_overlap_frames |
| max_available_frames = total_source_frames - control_start_frame |
| if max_available_frames < requested_frames: |
| raise FramePlanningError("The selected range ends too close to the source video end to build a valid chunk for the current model.") |
| if requested_frames < minimum_requested_frames: |
| raise FramePlanningError("The selected range ends too close to the source video end to build a valid chunk for the current model.") |
| plans.append(ChunkPlan(control_start_frame=control_start_frame, requested_frames=requested_frames, overlap_frames=plan_overlap_frames)) |
| unique_frames = requested_frames - plan_overlap_frames |
| written_unique_frames += unique_frames |
| cursor += unique_frames |
| return plans |
|
|
|
|
| def count_completed_chunks(plans: list[ChunkPlan], completed_unique_frames: int) -> tuple[int, int]: |
| if completed_unique_frames <= 0: |
| return 0, 0 |
| completed_chunks = 0 |
| consumed_frames = 0 |
| for plan in plans: |
| unique_frames = plan.requested_frames - plan.overlap_frames |
| if consumed_frames + unique_frames <= completed_unique_frames: |
| consumed_frames += unique_frames |
| completed_chunks += 1 |
| continue |
| break |
| return completed_chunks, consumed_frames |
|
|
|
|
| def count_completed_written_chunks(plans: list[ChunkPlan], completed_unique_frames: int) -> tuple[int, int]: |
| if completed_unique_frames <= 0: |
| return 0, 0 |
| completed_chunks = 0 |
| consumed_frames = 0 |
| for index, plan in enumerate(plans): |
| next_overlap_frames = plans[index + 1].overlap_frames if index + 1 < len(plans) else 0 |
| unique_frames = plan.requested_frames - next_overlap_frames |
| if consumed_frames + unique_frames <= completed_unique_frames: |
| consumed_frames += unique_frames |
| completed_chunks += 1 |
| continue |
| break |
| return completed_chunks, consumed_frames |
|
|
