Coverage for simulator/workflows.py: 95%
64 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1from __future__ import annotations
3import math
5from typing import Optional
7from sim_types import WorkflowConfig
8from sim_types import Model
9from sim_types import QualityLevel
11from constants import FPS
12from constants import FRAMES_OPTIONS
13from constants import FRAMES_PER_STEP_IDX
14from constants import NUM_STEPS
15from constants import SECONDS_IN_HOUR, SECONDS_IN_MINUTE
16from constants import TOTAL_INPUT_TOKENS
19# Shared physical constants
20MAX_FT_FRAMES: int = 1 + 80
21SUBSCENE_SECONDS: float = MAX_FT_FRAMES / FPS[Model.FT] # 81 frames @ 23 FPS → ~3.52 s
22SUBSCENES_PER_SCENE: int = 4 # default subscene grouping
23TOKENS_PER_FRAME = 500 # 1 frame generates around 500 tokens
26def _get_num_subscenes(total_video_seconds: int) -> int:
27 """Return the number of subscenes needed to cover the given video duration."""
28 return math.ceil(total_video_seconds / SUBSCENE_SECONDS)
31def _get_num_scenes(total_video_seconds: int) -> int:
32 """Return the number of scenes needed to cover the given video duration."""
33 return math.ceil(_get_num_subscenes(total_video_seconds) / SUBSCENES_PER_SCENE)
36def _get_num_frames(total_video_seconds: int, model: Model) -> int:
37 """Return the number of frames needed for the given video duration and model."""
38 return math.ceil(total_video_seconds * FPS[model])
41def _video_gen_work(
42 total_video_seconds: int,
43 num_scenes: int,
44 num_subscenes: int,
45 model_work_overrides: Optional[dict[Model, int | str | None]] = None,
46) -> dict[Model, int]:
47 """Standard model work for video-generation workflows (Podcast, Movie, etc.)."""
48 ret = {
49 Model.GEMMA: 1,
50 Model.FLUX: 1,
51 Model.HF: num_subscenes,
52 Model.HF_VAE: _get_num_frames(total_video_seconds, Model.HF),
53 Model.FT: num_subscenes,
54 Model.FT_VAE: _get_num_frames(total_video_seconds, Model.FT),
55 Model.UPSCALER: _get_num_frames(total_video_seconds, Model.FT),
56 Model.OTHERS: 1,
57 }
58 if model_work_overrides:
59 for model, value in model_work_overrides.items():
60 if value == "num_scenes":
61 ret[model] = num_scenes
62 elif value == "num_subscenes":
63 ret[model] = num_subscenes
64 elif isinstance(value, str):
65 raise ValueError(f"Invalid model_work override value: {value}")
66 elif value == 0 or value is None:
67 del ret[model]
68 else:
69 ret[model] = value
70 return ret
73class WorkOverrideType:
74 def __init__(self, value: int | str | None = None):
75 self.value = value
78def build_workflow_config(
79 total_video_seconds: int,
80 input_tokens: int,
81 model_work: dict[Model, int] | None = None,
82 *,
83 model_work_overrides: dict[Model, int | str | None] | None = None,
84 num_scenes_override: int | None = None,
85 num_steps_override: dict[Model, int] | None = None,
86 target_resolution: QualityLevel = QualityLevel.HIGH,
87) -> WorkflowConfig:
88 """Build a ``WorkflowConfig`` from base parameters, computing all derived values.
90 Parameters
91 ----------
92 model_work:
93 Explicit model-work dictionary. When ``None`` (default), standard
94 video-generation work is auto-generated from the other parameters.
95 exclude_models:
96 Models to remove from auto-generated ``model_work``.
97 model_work_overrides:
98 Key-value overrides applied on top of auto-generated ``model_work``.
99 If a value is set to "num_scenes", it will be replaced with the number of scenes (i.e. per-scene work).
100 target_resolution:
101 The target output resolution for the workflow (default HIGH).
102 When not HIGH, UPSCALER is automatically removed from model_work.
103 """
104 num_subscenes = _get_num_subscenes(total_video_seconds)
106 num_scenes = _get_num_scenes(total_video_seconds)
107 if num_scenes_override is not None:
108 num_scenes = num_scenes_override
110 num_steps = dict(NUM_STEPS)
111 if num_steps_override:
112 num_steps.update(num_steps_override)
114 if model_work is None:
115 model_work = _video_gen_work(
116 total_video_seconds,
117 num_scenes,
118 num_subscenes,
119 model_work_overrides,
120 )
122 return WorkflowConfig(
123 total_video_seconds=total_video_seconds,
124 total_scenes=num_scenes,
125 total_subscenes=num_subscenes,
126 total_frames={
127 Model.HF: _get_num_frames(total_video_seconds, Model.HF),
128 Model.FT: _get_num_frames(total_video_seconds, Model.FT),
129 },
130 per_subscene_frames={
131 Model.HF: math.ceil(_get_num_frames(total_video_seconds, Model.HF) / num_subscenes),
132 Model.FT: math.ceil(_get_num_frames(total_video_seconds, Model.FT) / num_subscenes),
133 },
134 num_steps=num_steps,
135 hf_frames=FRAMES_OPTIONS[Model.HF],
136 ft_frames=FRAMES_OPTIONS[Model.FT],
137 frames_per_step_idx=FRAMES_PER_STEP_IDX,
138 target_resolution=target_resolution,
139 total_input_tokens=input_tokens,
140 model_work=model_work,
141 )
144WORKFLOW_DURATIONS = { # in seconds
145 "podcast": int(10 * SECONDS_IN_MINUTE),
146 # TODO The input is two hours but the output should be shorter something like 1 or 2 minutes
147 "short": int(2 * SECONDS_IN_HOUR),
148 "movie": int(2 * SECONDS_IN_HOUR),
149 "story": int(10 * SECONDS_IN_MINUTE),
150 "lecture": int(5 * SECONDS_IN_MINUTE),
151 "slide": int(10 * SECONDS_IN_MINUTE),
152 "dubbing": int(10 * SECONDS_IN_MINUTE),
153 "editing": int(10 * SECONDS_IN_MINUTE),
154 "chat": 5,
155}
158# Podcast: 10-minute video from text/PDF input
159PODCAST_WORKFLOW = build_workflow_config(
160 total_video_seconds=WORKFLOW_DURATIONS["podcast"],
161 input_tokens=TOTAL_INPUT_TOKENS,
162)
164# Shorts: short clips from a 2-hour input video
165_SHORTS_SECONDS = WORKFLOW_DURATIONS["short"]
166_SHORTS_SCENES = _SHORTS_SECONDS // 10 # 10-second scene segmentation → 720
167SHORTS_WORKFLOW = build_workflow_config(
168 total_video_seconds=_SHORTS_SECONDS,
169 input_tokens=int(_SHORTS_SECONDS * TOKENS_PER_FRAME), # 1 fps × 500 tokens/frame
170 model_work={
171 Model.GEMMA: _SHORTS_SCENES,
172 Model.OTHERS: 1, # TODO isn't this 1 by default?
173 },
174 num_scenes_override=_SHORTS_SCENES,
175)
177# Movie: 2-hour movie
178MOVIE_WORKFLOW = build_workflow_config(
179 total_video_seconds=WORKFLOW_DURATIONS["movie"],
180 input_tokens=TOTAL_INPUT_TOKENS,
181 model_work_overrides={
182 Model.FLUX: "num_scenes",
183 },
184)
186# Animated Story: Podcast + 5% more HF denoising steps (LoRA overhead)
187OVERHEAD_PCT = 5
188ANIMATED_STORY_WORKFLOW = build_workflow_config(
189 total_video_seconds=WORKFLOW_DURATIONS["story"],
190 input_tokens=TOTAL_INPUT_TOKENS,
191 num_steps_override={
192 Model.HF: int(NUM_STEPS[Model.HF] * 1 + (OVERHEAD_PCT / 100.0))
193 },
194)
196# Lecture: 5-minute video, Flux generates per-scene images
197LECTURE_WORKFLOW = build_workflow_config(
198 total_video_seconds=WORKFLOW_DURATIONS["lecture"],
199 input_tokens=TOTAL_INPUT_TOKENS,
200 model_work_overrides={
201 Model.FLUX: "num_scenes",
202 },
203)
205# Slide Persona: same as Podcast but at low resolution, no upscaler
206SLIDE_PERSONA_WORKFLOW = build_workflow_config(
207 total_video_seconds=WORKFLOW_DURATIONS["slide"],
208 input_tokens=TOTAL_INPUT_TOKENS,
209 target_resolution=QualityLevel.LOW,
210 model_work_overrides={
211 Model.UPSCALER: None,
212 },
213)
215# Dubbing: like Podcast but without Flux, and double the audio work
216DUBBING_WORKFLOW = build_workflow_config(
217 total_video_seconds=WORKFLOW_DURATIONS["dubbing"],
218 input_tokens=TOTAL_INPUT_TOKENS,
219 model_work_overrides={
220 Model.FLUX: None,
221 Model.OTHERS: 2, # Double audio work
222 },
223)
225# Editing: like Podcast but without GEMMA, FLUX, or OTHERS
226EDITING_WORKFLOW = build_workflow_config(
227 total_video_seconds=WORKFLOW_DURATIONS["editing"],
228 input_tokens=TOTAL_INPUT_TOKENS,
229 model_work_overrides={
230 Model.GEMMA: None,
231 Model.FLUX: None,
232 Model.OTHERS: None,
233 }
234)
236# Video Chat: like Podcast but only 5 seconds of output video
237VIDEO_CHAT_WORKFLOW = build_workflow_config(
238 total_video_seconds=WORKFLOW_DURATIONS["chat"],
239 input_tokens=TOTAL_INPUT_TOKENS,
240)
243WORKFLOWS = {
244 "podcast": PODCAST_WORKFLOW,
245 "chat": VIDEO_CHAT_WORKFLOW,
246 "dubbing": DUBBING_WORKFLOW,
247 "editing": EDITING_WORKFLOW,
248 "lecture": LECTURE_WORKFLOW,
249 "movie": MOVIE_WORKFLOW,
250 "short": SHORTS_WORKFLOW,
251 "slide": SLIDE_PERSONA_WORKFLOW,
252 "story": ANIMATED_STORY_WORKFLOW,
253}