Spaces:
Sleeping
Sleeping
File size: 1,914 Bytes
0da718b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | from diffsynth.pipelines.wan_video import WanVideoPipeline
from diffsynth.diffusion.base_pipeline import PipelineUnit
class WanVideoUnit_CutInjector(PipelineUnit):
"""Turns per-sample cut annotations into a planning-token schedule.
Consumes the `cut_info` field emitted by the dataset and produces the
`cut_schedule` consumed by `model_fn_wan_video_with_cut`.
"""
def __init__(self):
super().__init__(
input_params=("cut_info",),
output_params=("cut_schedule",),
)
def process(self, pipe: WanVideoPipeline, cut_info):
"""
Expected `cut_info` structure:
{
"cuts": [36, 43] or [15] or [[start, end], ...],
"type": "hardcut" | "softcut" | "normal" | ...
}
Only 'hardcut' samples inject a token. Every element of 'cuts' becomes
one hardcut_embedding at frame index f (mapped to the fractional
latent coordinate t = 1 + f/4 for a 4x temporally compressed VAE).
Other types (normal / softcut / camera motion) inject nothing so the
model still sees them as regular video without any control token.
"""
if cut_info is None:
return {}
raw_cuts = cut_info.get("cuts", [])
cut_type = cut_info.get("type", "hardcut")
if cut_type != "hardcut":
return {"cut_schedule": []}
def frame_to_t(f):
return 1.0 + float(f) / 4.0
schedule = []
for f in raw_cuts:
# Robust to both flat [f, ...] and nested [[f], ...] formats.
val = f
if isinstance(f, (list, tuple)):
if len(f) == 0:
continue
val = f[0]
schedule.append({
"t": frame_to_t(val),
"token_name": "hardcut_embedding",
})
return {"cut_schedule": schedule}
|