MiniMax-H3 ref2va: aligned video guides (IC-LoRA style v2v) for ai-toolkit

A developer contributed a patch to ai-toolkit's MiniMax-H3 diffusion model extension that fixes reference-to-video-alignment (ref2va) handling of control images and videos in IC-LoRA-style video-to-video workflows. The change refactors keyframe conversion into a helper and preserves the nested per-batch-item structure of control_tensor_list, which previously collapsed multiple references into a single <Picture> block and caused an "image features and image tokens do not match" error in the Qwen3-VL processor.

| | diff --git a/extensions built in/diffusion models/minimax h3/minimax h3.py b/extensions built in/diffusion models/minimax h3/minimax h3.py | | | index 4c21927..2509789 100644 | | | --- a/extensions built in/diffusion models/minimax h3/minimax h3.py | | | +++ b/extensions built in/diffusion models/minimax h3/minimax h3.py | | | @@ -537,44 +537,56 @@ class MinimaxH3Model BaseModel : | | | control tensors arrive in 0, 1 ; the Qwen3-VL processor wants PIL | | | keyframes per prompt = None len prompt | | | if control images is not None: | | | - if isinstance control images, torch.Tensor : | | | - images = control images i for i in range control images.shape 0 | | | - elif isinstance control images, list : | | | - images = | | | - c 0 if isinstance c, torch.Tensor and c.ndim == 4 else c | | | - for c in control images | | | - | | | - else: | | | - images = control images | | | - pil images = | | | - for img in images: | | | + | | | + def to keyframe img : | | | if isinstance img, torch.Tensor : | | | if img.ndim == 4: | | | img = img 0 | | | arr = img.float .clamp 0, 1 255 .round .to torch.uint8 | | | - pil images.append | | | - self. present image control | | | - Image.fromarray arr.permute 1, 2, 0 .cpu .numpy | | | - | | | + return self. present image control | | | + Image.fromarray arr.permute 1, 2, 0 .cpu .numpy | | | | | | - elif isinstance img, str : | | | + if isinstance img, str : | | | a control VIDEO path: 2 fps timestamped presentation over | | | the SAME frames the latent rows use dataset treatment | | | when caching training embeds, sample-length at sampling | | | ds cfg = getattr self, " ref video dataset config", None | | | - pil images.append | | | - load video ref for te | | | - self, img, ds cfg, max frames=self. sample ref max frames | | | - | | | + return load video ref for te | | | + self, img, ds cfg, max frames=self. sample ref max frames | | | | | | - else: | | | - pil images.append img | | | - if len pil images == 1: | | | - keyframes per prompt = pil images len prompt | | | - elif len pil images == len prompt : | | | - keyframes per prompt = img for img in pil images | | | + return img | | | + | | | + batch.control tensor list is item ref : several references per | | | + batch item. It must stay nested per item -- collapsing it emits a | | | + single <Picture block while the image processor flattens the | | | + inner list into N images "image features and image tokens do not | | | + match", tokens counted from image grid thw 0 only . | | | + if | | | + isinstance control images, list | | | + and len control images 0 | | | + and all isinstance c, list, tuple for c in control images | | | + : | | | + per item = to keyframe r for r in refs for refs in control images | | | + if len per item == 1 and len prompt 1: | | | + per item = per item len prompt | | | + keyframes per prompt = per item | | | else: | | | - keyframes per prompt = pil images len prompt | | | + if isinstance control images, torch.Tensor : | | | + images = control images i for i in range control images.shape 0 | | | + elif isinstance control images, list : | | | + images = | | | + c 0 if isinstance c, torch.Tensor and c.ndim == 4 else c | | | + for c in control images | | | + | | | + else: | | | + images = control images | | | + pil images = to keyframe img for img in images | | | + if len pil images == 1: | | | + keyframes per prompt = pil images len prompt | | | + elif len pil images == len prompt : | | | + keyframes per prompt = img for img in pil images | | | + else: | | | + keyframes per prompt = pil images len prompt | | | | | | embeds list, tags list = , | | | for p, keyframes in zip prompt, keyframes per prompt : | | | @@ -721,6 +733,23 @@ class MinimaxH3Model BaseModel : | | | ------------------------------------------------------------------ | | | Training forward | | | ------------------------------------------------------------------ | | | + def aligned ref flags self, ref blocks : | | | + """Which reference blocks share the target's coordinates. | | | + | | | + With align video refs a video reference stops being a loose reference and | | | + becomes a v2v guide: same rotary clock, same spatial grid as the target, so guide | | | + frame i drives output frame i without the model having to search for the | | | + correspondence. Image references are never aligned — an identity photo has no | | | + spatial correspondence with the output and should not claim one. | | | + | | | + A head-swap pack is exactly this pair: aligned driving video, unaligned identity. | | | + """ | | | + if not ref blocks or not self.model config.model kwargs.get | | | + "align video refs", False | | | + : | | | + return | | | + return tuple block 0 1 for block in ref blocks | | | + | | | def build condition | | | self, batch: "DataLoaderBatchDTO", latent shape, device, dtype | | | : | | | @@ -890,6 +919,7 @@ class MinimaxH3Model BaseModel : | | | text tag list = t for , t in trimmed | | | | | | --- packed layout per item: text lengths differ -------------- | | | + aligned refs = self. aligned ref flags ref blocks | | | layouts = | | | for i in range batch size : | | | layouts.append | | | @@ -901,6 +931,7 @@ class MinimaxH3Model BaseModel : | | | num audio latents=a lat, | | | keyframe anchors=keyframe anchors, | | | ref blocks=ref blocks, | | | + aligned refs=aligned refs, | | | | | | | | | | | | @@ -1128,9 +1159,11 @@ class MinimaxH3Ref2VAModel MinimaxH3Model : | | | def text embedding space version self : | | | the presentation of image references changes the embeds - new cache key | | | n = self. image ref video frames | | | - if n: | | | - return f"{self.arch}:img as vid{n}" | | | - return self.arch | | | + base = f"{self.arch}:img as vid{n}" if n else self.arch | | | + if getattr self, "control latent only", False : | | | + embeds drop the control media entirely - different space | | | + return f"{base}:ctrl latent only" | | | + return base | | | | | | def present image control self, image: Image.Image : | | | n = self. image ref video frames | | | @@ -1150,6 +1183,22 @@ class MinimaxH3Ref2VAModel MinimaxH3Model : | | | D-OPSD: a no-grad teacher pass with the target as its own reference | | | becomes the training target for the reference-free student pass | | | self.dopsd = bool self.model config.model kwargs.get "dopsd", False | | | + VLM-gap variant of D-OPSD: teacher and student share the SAME condition | | | + the aligned guide ; the information gap is the VLM's reading of that | | | + guide, which only the teacher's text embeds carry. The student's | | | + sequence loses the vision tokens -- here ~8.2k of ~14.6k -- so it is | | | + cheaper to train and to sample, and it matches ComfyUI's AddGuide path, | | | + where the guide reaches the DiT as latents but never reaches the VLM. | | | + Control media reaches the DiT as latents only -- it is never presented | | | + to the VLM. For a v2v guide that is usually what you want: the guide | | | + already lands on the target's rotary grid via align video refs, the | | | + vision tokens are pure cost ~8.2k of a ~14.6k sequence on a 73-frame | | | + 512x288 clip , and inference paths that inject the guide as a keyframe | | | + ComfyUI's Add Guide never show it to the VLM either -- so training | | | + with the presentation creates a train/inference gap. | | | + self.control latent only = bool | | | + self.model config.model kwargs.get "control latent only", False | | | + | | | if self.dopsd: | | | self.dopsd self ref = True | | | self.require pixel tensor cache = True | | | @@ -1336,10 +1385,16 @@ class MinimaxH3Ref2VAModel MinimaxH3Model : | | | | | | cap.release | | | h0, w0 = frames 0 .shape :2 | | | - match the sample canvas's pixel area, own aspect kept | | | - ph, pw = packing.reference video pixel size | | | - w0, h0, gen config.height, gen config.width | | | - | | | + if bool self.model config.model kwargs.get "align video refs", False : | | | + aligned guide: the sample canvas exactly, same rule the training path uses. | | | + Area-matching the guide's own aspect here would hand build packed sequence a | | | + grid the target does not have and fail the sample, not the step. | | | + ph, pw = int gen config.height , int gen config.width | | | + else: | | | + match the sample canvas's pixel area, own aspect kept | | | + ph, pw = packing.reference video pixel size | | | + w0, h0, gen config.height, gen config.width | | | + | | | pixels = torch.from numpy np.stack frames .float / 255.0 2.0 - 1.0 | | | pixels = pixels.permute 3, 0, 1, 2 None 1, 3, T, H, W | | | pixels = | | | @@ -1422,12 +1477,18 @@ class MinimaxH3Ref2VAModel MinimaxH3Model : | | | "ref2va: every item in a batch must have the same number of " | | | "reference videos" | | | | | | + with align video refs a control VIDEO is a v2v guide, not a loose reference: | | | + it has to sit on the target's exact latent grid build packed sequence enforces | | | + it , so it is encoded at the target's size instead of area-matched to its own | | | + aspect. Identity images never take this path. | | | + align = bool self.model config.model kwargs.get "align video refs", False | | | for ref idx in range vid count : | | | lats = | | | auds = | | | for per item in paths per item: | | | entry = load ref video latent | | | - self, per item ref idx , batch.dataset config, target h, target w | | | + self, per item ref idx , batch.dataset config, target h, target w, | | | + align=align, | | | | | | lats.append entry "latent" .to device, torch.float32 | | | auds.append entry.get "audio rows" | | | diff --git a/extensions built in/diffusion models/minimax h3/src/packing.py b/extensions built in/diffusion models/minimax h3/src/packing.py | | | index 6e77429..39c1c72 100644 | | | --- a/extensions built in/diffusion models/minimax h3/src/packing.py | | | +++ b/extensions built in/diffusion models/minimax h3/src/packing.py | | | @@ -297,6 +297,7 @@ def build packed sequence | | | patch size= 1, 2, 2 , | | | keyframe anchors: Tuple str, ... = , | | | ref blocks: Tuple Tuple int, int, int , ... = , | | | + aligned refs: Tuple bool, ... = , | | | - PackedLayout: | | | """Build the text \| conditions \| target audio \| target video layout. | | | | | | @@ -307,7 +308,22 @@ def build packed sequence | | | for images. References keep their OWN aspect on their own | | | aspect-normalized grid; an image block advances the shared media clock by | | | 1.0, a video block by its temporal span, and the target streams start | | | - after the cumulative advance .""" | | | + after the cumulative advance . | | | + | | | + aligned refs marks blocks that should be aligned to the target rather | | | + than placed beside it: an aligned block reuses the target's rotary clock and | | | + the target's spatial grid, so its latent frame i and pixel h, w | | | + land on exactly the coordinates of the target's, and it does not advance the | | | + media clock. That is the v2v/IC-LoRA arrangement — attention between a guide | | | + row and the target row at the same position costs nothing positionally, which | | | + is what makes frame-accurate control pose, depth, a driving performance | | | + learnable. Ordinary references stay unaligned: an identity reference has no | | | + spatial correspondence with the output and should not claim one. | | | + | | | + A head-swap pack therefore uses both — an aligned guide video plus an | | | + unaligned identity reference. Aligned blocks must match the target's latent | | | + height and width; their frame count may be shorter they align from the | | | + target's first frame .""" | | | if keyframe anchors and ref blocks: | | | raise ValueError "keyframe anchors and ref blocks are mutually exclusive" | | | , ph, pw = patch size | | | @@ -317,6 +333,17 @@ def build packed sequence | | | reference's soundtrack packs as clean audio rows immediately BEFORE its | | | own video rows | | | ref blocks = tuple tuple b + 0, 4 - len b for b in ref blocks | | | + aligned flags = tuple aligned refs + False, len ref blocks - len aligned refs | | | + if len aligned flags = len ref blocks : | | | + raise ValueError | | | + f"aligned refs has {len aligned refs } entries for {len ref blocks } ref blocks" | | | + | | | + for , b h, b w, , aligned in zip ref blocks, aligned flags : | | | + if aligned and b h, b w = latent height, latent width : | | | + raise ValueError | | | + "an aligned reference must share the target's latent resolution " | | | + f" {latent height}x{latent width} , got {b h}x{b w}" | | | + | | | ref vid rows = t h // ph w // pw for t, h, w, in ref blocks | | | ref aud rows = a AUDIO CHANNELS for , , , a in ref blocks | | | num cond = | | | @@ -337,7 +364,12 @@ def build packed sequence | | | text rows sit on the time axis at their row index; the media clock | | | continues from there past the reference blocks, so prompt length and | | | reference count/length shifts the whole media clock | | | - media advance = sum block advance t, a for t, , , a in ref blocks | | | + Aligned blocks sit ON the target's clock, so they must not push it forward. | | | + media advance = sum | | | + block advance t, a | | | + for t, , , a , aligned in zip ref blocks, aligned flags | | | + if not aligned | | | + | | | media origin = float num text + media advance | | | position ids = torch.zeros seq len, 3, dtype=torch.float64 | | | position ids :num text, 0 = torch.arange num text, dtype=torch.float64 | | | @@ -377,25 +409,34 @@ def build packed sequence | | | cond video idx.append torch.arange cond start, audio start | | | ref cursor = audio start | | | for i, ref t, ref h, ref w, ref a in enumerate ref blocks : | | | - each reference on its own aspect-normalized grid area-matched to | | | - the target, so the grids span comparable ranges | | | - ref sqrt area = math.sqrt ref h ref w | | | - w grid = spatial position grid ref w, pw, ref sqrt area | | | - ref grid = torch.stack | | | - | | | - g.reshape -1 | | | - for g in torch.meshgrid | | | - spatial position grid ref h, ph, ref sqrt area , | | | - w grid, | | | - indexing="ij", | | | - | | | - , | | | - dim=-1, | | | - | | | + aligned = aligned flags i | | | + if aligned: | | | + the target's own grid and clock: row i, h, w of the guide lands on | | | + the coordinate of row i, h, w of the output | | | + w grid = width grid | | | + ref grid = frame grid | | | + block clock = media origin | | | + else: | | | + each reference on its own aspect-normalized grid area-matched to | | | + the target, so the grids span comparable ranges | | | + ref sqrt area = math.sqrt ref h ref w | | | + w grid = spatial position grid ref w, pw, ref sqrt area | | | + ref grid = torch.stack | | | + | | | + g.reshape -1 | | | + for g in torch.meshgrid | | | + spatial position grid ref h, ph, ref sqrt area , | | | + w grid, | | | + indexing="ij", | | | + | | | + , | | | + dim=-1, | | | + | | | + block clock = ref clock | | | if ref a: | | | soundtrack rows first: channel-major, shared 40/s clock from the | | | block origin, width pinned to the ref grid's extremes | | | - a time = ref clock + torch.arange ref a, dtype=torch.float64 | | | + a time = block clock + torch.arange ref a, dtype=torch.float64 | | | rows = slice ref cursor, ref cursor + ref aud rows i | | | position ids rows, 0 = a time.repeat AUDIO CHANNELS | | | position ids rows, 2 = torch.cat | | | @@ -408,12 +449,13 @@ def build packed sequence | | | ref cursor += ref aud rows i | | | rows per ref frame = ref grid.shape 0 | | | block = torch.empty ref t, rows per ref frame, 3, dtype=torch.float64 | | | - block :, :, 0 = temporal position grid ref t, ref clock :, None | | | + block :, :, 0 = temporal position grid ref t, block clock :, None | | | block :, :, 1: = ref grid None | | | position ids ref cursor : ref cursor + ref vid rows i = block.reshape -1, 3 | | | cond video idx.append torch.arange ref cursor, ref cursor + ref vid rows i | | | ref cursor += ref vid rows i | | | - ref clock += block advance ref t, ref a | | | + if not aligned: | | | + ref clock += block advance ref t, ref a | | | | | | audio rows: channel-major, one rotary unit per latent 40/s = 24fps 5/3 , | | | no height coordinate, width pinned to the grid extremes per channel | | | diff --git a/extensions built in/diffusion models/minimax h3/src/pipeline.py b/extensions built in/diffusion models/minimax h3/src/pipeline.py | | | index 6e08cd3..7a65747 100644 | | | --- a/extensions built in/diffusion models/minimax h3/src/pipeline.py | | | +++ b/extensions built in/diffusion models/minimax h3/src/pipeline.py | | | @@ -136,6 +136,9 @@ class MiniMaxH3Pipeline: | | | num audio latents=a lat, | | | keyframe anchors=anchors, | | | ref blocks=ref blocks, | | | + same rule as training, from the same helper: a layout mismatch between the | | | + two is silent — sampling just quietly stops matching what was trained | | | + aligned refs=model. aligned ref flags ref blocks , | | | | | | num cond = layout.num condition video rows | | | | | | diff --git a/extensions built in/diffusion models/minimax h3/src/ref video cache.py b/extensions built in/diffusion models/minimax h3/src/ref video cache.py | | | index c54c263..a320cbc 100644 | | | --- a/extensions built in/diffusion models/minimax h3/src/ref video cache.py | | | +++ b/extensions built in/diffusion models/minimax h3/src/ref video cache.py | | | @@ -164,7 +164,8 @@ def cache path path: str, hash dict: dict - str: | | | | | | @torch.no grad | | | def load ref video latent | | | - model, path: str, dataset config, target height: int, target width: int | | | + model, path: str, dataset config, target height: int, target width: int, | | | + align: bool = False, | | | - dict: | | | """Returns {"latent": C, T, h, w cpu tensor, "num frames": int}, | | | encoding + disk-caching on first use. model is the MinimaxH3 model | | | @@ -173,8 +174,14 @@ def load ref video latent | | | if mem cache is None: | | | mem cache = {} | | | model. ref video cache = mem cache | | | - if path in mem cache: | | | - return mem cache path | | | + The entry depends on the TARGET the ref was sized against, not just the file: with | | | + multi-resolution buckets the same guide is requested at several target sizes in one | | | + run, and a path-only key handed back whatever size happened to be encoded first. | | | + Unaligned that passed silently the guide just sat on a wrong grid, and cost whatever | | | + its own size cost ; aligned it raises "must share the target's latent resolution". | | | + mem key = path, int target height , int target width , bool align | | | + if mem key in mem cache: | | | + return mem cache mem key | | | | | | cap = cv2.VideoCapture path | | | if not cap.isOpened : | | | @@ -198,10 +205,14 @@ def load ref video latent | | | | | | hash dict = { | | | "signature": get quick signature string path , | | | - "ref sizing": "match target area", | | | + an aligned guide is resized to the target's EXACT grid, so the cache key | | | + needs the dimensions, not just the area -- two buckets with the same area | | | + and different aspects are different tensors here | | | + "ref sizing": "match target exact" if align else "match target area", | | | "target area": int target height target width | | | if target height and target width | | | else 0, | | | + "target hw": int target height , int target width if align else 0, | | | "num frames": num frames, | | | "fps": dataset config.fps, | | | "trim tail": trim tail, | | | @@ -217,13 +228,22 @@ def load ref video latent | | | "num frames": int sd "num frames" .item , | | | "audio rows": sd.get "audio latent" , | | | } | | | - mem cache path = entry | | | + mem cache mem key = entry | | | return entry | | | | | | - match the target's pixel area the dataset bucket the target trains at | | | - with the ref's own aspect: same aspect - identical size; aspect- | | | - preserving resize, no crop | | | - out h, out w = reference video pixel size src w, src h, target height, target width | | | + if align: | | | + An aligned guide shares the target's rotary clock AND its spatial grid, so it | | | + must land on the target's exact latent resolution -- build packed sequence | | | + rejects anything else. Area-matching the guide's own aspect would leave 25% of | | | + a real head-swap set a few percent off and crash the pack, so the guide is | | | + resized to the target outright; the small anisotropy is the price of frame- | | | + accurate correspondence, and it only applies to guides, never to identity refs. | | | + out h, out w = int target height , int target width | | | + else: | | | + match the target's pixel area the dataset bucket the target trains at | | | + with the ref's own aspect: same aspect - identical size; aspect- | | | + preserving resize, no crop | | | + out h, out w = reference video pixel size src w, src h, target height, target width | | | | | | indices = ref frame indices | | | total, src fps, num frames, dataset config.fps, trim tail | | | @@ -274,5 +294,5 @@ def load ref video latent | | | os.makedirs os.path.dirname cache file , exist ok=True | | | save file state dict, cache file | | | entry = {"latent": latent, "num frames": num frames, "audio rows": audio rows} | | | - mem cache path = entry | | | + mem cache mem key = entry | | | return entry | | | diff --git a/extensions built in/diffusion models/minimax h3/src/text encoder.py b/extensions built in/diffusion models/minimax h3/src/text encoder.py | | | index f58863d..a9680e4 100644 | | | --- a/extensions built in/diffusion models/minimax h3/src/text encoder.py | | | +++ b/extensions built in/diffusion models/minimax h3/src/text encoder.py | | | @@ -120,6 +120,8 @@ def encode minimax h3 prompt | | | processor, Qwen3VLProcessor needed only when keyframes are present | | | prompt: str, | | | keyframes: Optional List = None, PIL images already on the target canvas | | | + reference videos: Optional List = None, videos sampled to 2 fps, each TCHW/THWC | | | + reference video timestamps: Optional List List float = None, | | | device: Optional torch.device = None, | | | dtype: Optional torch.dtype = None, | | | max length: Optional | | | @@ -220,6 +222,41 @@ def encode minimax h3 prompt | | | | | | pic idx += 1 | | | | | | + if reference videos: | | | + vision = processor.video processor | | | + videos=reference videos, do sample frames=False, return tensors="pt" | | | + | | | + pixel values videos = vision "pixel values videos" | | | + video grid thw = vision "video grid thw" | | | + merge = processor.video processor.merge size 2 | | | + vision start = tokenizer.convert tokens to ids "<\|vision start\| " | | | + vision end = tokenizer.convert tokens to ids "<\|vision end\| " | | | + video pad = tokenizer.convert tokens to ids "<\|video pad\| " | | | + for i in range len reference videos : | | | + label ids = tokenizer f"<Video {i + 1} : ", add special tokens=False | | | + "input ids" | | | + | | | + token ids += label ids | | | + token tags += TEXT TAG len label ids | | | + grid t, grid h, grid w = int x for x in video grid thw i | | | + tokens per block = grid h grid w // merge | | | + timestamps = | | | + reference video timestamps i | | | + if reference video timestamps is not None | | | + else j / 2.0 for j in range grid t 2 | | | + | | | + if len timestamps % 2: | | | + timestamps = list timestamps + timestamps -1 | | | + for block in range grid t : | | | + ts idx = min block 2, len timestamps - 2 | | | + block ts = timestamps ts idx + timestamps ts idx + 1 / 2.0 | | | + time ids = tokenizer | | | + f"<{block ts:.1f} seconds ", add special tokens=False | | | + "input ids" | | | + vision ids = vision start + video pad tokens per block + vision end | | | + token ids += time ids + vision ids | | | + token tags += TEXT TAG len time ids + VIDEO TAG len vision ids | | | + | | | prompt ids = tokenizer prompt, add special tokens=False "input ids" | | | if max length is not None and max length 0: | | | the cap applies to the caption only; a keyframe's vision block is | | | diff --git a/tests/test minimax h3 aligned refs.py b/tests/test minimax h3 aligned refs.py | | | new file mode 100644 | | | index 0000000..8c565d3 | | | --- /dev/null | | | +++ b/tests/test minimax h3 aligned refs.py | | | @@ -0,0 +1,129 @@ | | | +"""An aligned reference must land on the target's coordinates exactly. | | | + | | | +That is the whole point of the flag: a guide row and the target row it drives share a rotary | | | +position, so attention between them costs nothing positionally. If the clock or the spatial grid | | | +drifts by even one unit the guide stops being frame-accurate, and the failure is invisible — | | | +training still runs, the model just never learns tight control. | | | +""" | | | + | | | +import pytest | | | +import torch | | | + | | | +from extensions built in.diffusion models.minimax h3.src.packing import build packed sequence | | | + | | | +TEXT = torch.ones 4, dtype=torch.long | | | +FRAMES, H, W = 7, 8, 8 | | | + | | | + | | | +def layout ref blocks= , aligned refs= : | | | + return build packed sequence | | | + text token tags=TEXT, | | | + num latent frames=FRAMES, | | | + latent height=H, | | | + latent width=W, | | | + num audio latents=0, | | | + ref blocks=ref blocks, | | | + aligned refs=aligned refs, | | | + | | | + | | | + | | | +def rows layout, start, count : | | | + return layout.position ids start : start + count | | | + | | | + | | | +def test aligned reference matches the target coordinates row for row - None: | | | + layout = layout ref blocks= FRAMES, H, W , , aligned refs= True, | | | + | | | + rows per ref = FRAMES H // 2 W // 2 | | | + guide = rows layout, layout.video indices 0 .item , rows per ref | | | + target = layout.position ids -rows per ref: | | | + | | | + assert torch.equal guide, target | | | + | | | + | | | +def test aligned reference does not shift the target clock - None: | | | + """The target must sit where it would with no reference at all.""" | | | + bare = layout | | | + aligned = layout ref blocks= FRAMES, H, W , , aligned refs= True, | | | + | | | + rows = FRAMES H // 2 W // 2 | | | + assert torch.equal bare.position ids -rows: , aligned.position ids -rows: | | | + | | | + | | | +def test unaligned reference still sits beside the target - None: | | | + """The default path is untouched: an ordinary reference keeps its own clock.""" | | | + layout = layout ref blocks= 1, H, W , , aligned refs= False, | | | + | | | + rows per ref = H // 2 W // 2 | | | + ref = rows layout, layout.video indices 0 .item , rows per ref | | | + target = layout.position ids -FRAMES rows per ref : | | | + | | | + assert ref :, 0 .max < target :, 0 .min | | | + | | | + | | | +def test omitting the flag keeps upstream behaviour - None: | | | + explicit = layout ref blocks= 1, H, W , , aligned refs= False, | | | + implicit = layout ref blocks= 1, H, W , | | | + | | | + assert torch.equal explicit.position ids, implicit.position ids | | | + | | | + | | | +def test head swap pack aligns only the guide - None: | | | + """Guide aligned to the target, identity reference parked beside it.""" | | | + layout = layout ref blocks= FRAMES, H, W , 1, H, W , aligned refs= True, False | | | + | | | + rows per frame = H // 2 W // 2 | | | + guide rows = FRAMES rows per frame | | | + start = layout.video indices 0 .item | | | + | | | + guide = rows layout, start, guide rows | | | + identity = rows layout, start + guide rows, rows per frame | | | + target = layout.position ids -guide rows: | | | + | | | + assert torch.equal guide, target | | | + assert identity :, 0 .max < target :, 0 .min | | | + | | | + | | | +def test shorter aligned guide starts at the targets first frame - None: | | | + short = 3 | | | + layout = layout ref blocks= short, H, W , , aligned refs= True, | | | + | | | + rows per frame = H // 2 W // 2 | | | + guide = rows layout, layout.video indices 0 .item , short rows per frame | | | + target = layout.position ids -FRAMES rows per frame : | | | + | | | + assert torch.equal guide, target : short rows per frame | | | + | | | + | | | +def test aligned reference must match the target resolution - None: | | | + """A different grid cannot be aligned, and silently misaligning is the worst outcome.""" | | | + with pytest.raises ValueError, match="aligned reference" : | | | + layout ref blocks= FRAMES, H // 2, W , , aligned refs= True, | | | + | | | + | | | +class FakeModel: | | | + """Just enough of the model to exercise the flag's plumbing.""" | | | + | | | + def init self, enabled : | | | + self.model config = type "C", , {"model kwargs": {"align video refs": enabled}} | | | + | | | + aligned ref flags = None bound below | | | + | | | + | | | +def flags enabled, ref blocks : | | | + from extensions built in.diffusion models.minimax h3.minimax h3 import MinimaxH3Model | | | + | | | + return MinimaxH3Model. aligned ref flags FakeModel enabled , ref blocks | | | + | | | + | | | +def test flag off aligns nothing - None: | | | + assert flags False, 7, 8, 8, 0 , 1, 8, 8, 0 == | | | + | | | + | | | +def test flag on aligns video blocks only - None: | | | + """Video reference becomes the guide; the image reference stays a reference.""" | | | + assert flags True, 7, 8, 8, 0 , 1, 8, 8, 0 == True, False | | | + | | | + | | | +def test no references needs no flags - None: | | | + assert flags True, == | | | diff --git a/toolkit/data loader.py b/toolkit/data loader.py | | | index 38b7ad4..a94e369 100644 | | | --- a/toolkit/data loader.py | | | +++ b/toolkit/data loader.py | | | @@ -540,6 +540,7 @@ class AiToolkitDataset LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin | | | encode control in text embeddings=self.sd.encode control in text embeddings if self.sd else False, | | | encode first frame in text embeddings=getattr self.sd, 'encode first frame in text embeddings', False if self.sd else False, | | | dopsd self ref=getattr self.sd, 'dopsd self ref', False if self.sd else False, | | | + control latent only=getattr self.sd, 'control latent only', False if self.sd else False, | | | text embedding space version=self.sd.text embedding space version if self.sd else "sd1", | | | te padding side=self.sd.te padding side if self.sd else "right", | | | latent space version=latent space version, | | | diff --git a/toolkit/data transfer object/data loader.py b/toolkit/data transfer object/data loader.py | | | index 8c7dfc6..2c9b615 100644 | | | --- a/toolkit/data transfer object/data loader.py | | | +++ b/toolkit/data transfer object/data loader.py | | | @@ -85,6 +85,8 @@ class FileItemDTO | | | | | | D-OPSD: also cache teacher embeds with the item's own media as reference 1 | | | self.dopsd self ref = kwargs.get "dopsd self ref", False | | | + control media rides as latents only, never into the text encoder | | | + self.control latent only = kwargs.get "control latent only", False | | | self.te padding side = kwargs.get "te padding side", "right" | | | self.latent space version = kwargs.get "latent space version", "sd1" | | | self.text embedding space version = kwargs.get "text embedding space version", "sd1" | | | diff --git a/toolkit/dataloader mixins.py b/toolkit/dataloader mixins.py | | | index 60e0a9e..ae6a1e5 100644 | | | --- a/toolkit/dataloader mixins.py | | | +++ b/toolkit/dataloader mixins.py | | | @@ -2414,8 +2414,15 @@ class TextEmbeddingCachingMixin: | | | ctrl img = ctrl img list 0 | | | else: | | | ctrl img = ctrl img list | | | + if getattr file item, 'control latent only', False : | | | + caption-only embeds. The control | | | + media still reaches the DiT as aligned latents; what the | | | + student loses is the VLM's reading of it, which is exactly | | | + the information the teacher distills down. Also matches | | | + ComfyUI's AddGuide path, where the guide never hits the VLM. | | | + ctrl img = None | | | for path, caption in encode targets: | | | - if path in dropout target paths: | | | + if path in dropout target paths or ctrl img is None: | | | dropout embeds are plain text. Only fall back to the | | | control images if the model cannot encode without them | | | try: |