[pull] master from Comfy-Org:master

comfy/ldm/wan/model.py

@@ -1665,7 +1665,7 @@ def forward_orig(self, x, t, context, clip_fea=None, freqs=None, transformer_opt
 
         # embeddings
         x = self.patch_embedding(x.float()).to(x.dtype)
-        if ref_mask_latents is not None:  # SCAIL-2 additive mask stream
+        if ref_mask_latents is not None:  # SCAIL-2 additive mask stream (one identity mask frame per reference, then video)
             x = x + self.patch_embedding_mask(ref_mask_latents.float()).to(x.dtype)
         grid_sizes = x.shape[2:]
         transformer_options["grid_sizes"] = grid_sizes
@@ -1728,22 +1728,25 @@ def block_wrap(args):
 
     # ref_mask_flag is a scalar bool (CONDConstant, SCAIL-2 only). False => replacement mode,
     # which places ref/pose via H/W rope shifts instead of the animation-mode temporal offset.
+    # reference_latent may stack several frames: the last is the primary reference adjacent to the video, the earlier frames are additional references.
     def rope_encode(self, t, h, w, t_start=0, steps_t=None, steps_h=None, steps_w=None, device=None, dtype=None, pose_latents=None, reference_latent=None, ref_mask_flag=None, transformer_options={}):
+        ref_t_patches = 0
+        if reference_latent is not None:
+            ref_t_patches = (reference_latent.shape[2] + (self.patch_size[0] // 2)) // self.patch_size[0]
+
         if ref_mask_flag is not None and not bool(ref_mask_flag):
             REF_ROPE_H = 120.0
             POSE_ROPE_W = 120.0
 
-            ref_t_patches = 0
-            if reference_latent is not None:
-                ref_t_patches = (reference_latent.shape[2] + (self.patch_size[0] // 2)) // self.patch_size[0]
             main_t_patches = t - ref_t_patches
+            video_t_start = max(ref_t_patches - 1, 0)
 
             parts = []
             if ref_t_patches > 0:
                 ref_tf = {"rope_options": {"shift_y": REF_ROPE_H, "shift_x": 0.0, "scale_y": 1.0, "scale_x": 1.0}}
                 parts.append(super().rope_encode(ref_t_patches, h, w, t_start=0, device=device, dtype=dtype, transformer_options=ref_tf))
             if main_t_patches > 0:
-                parts.append(super().rope_encode(main_t_patches, h, w, t_start=0, device=device, dtype=dtype, transformer_options=transformer_options))
+                parts.append(super().rope_encode(main_t_patches, h, w, t_start=video_t_start, device=device, dtype=dtype, transformer_options=transformer_options))
 
             if pose_latents is not None:
                 F_pose, H_pose, W_pose = pose_latents.shape[-3], pose_latents.shape[-2], pose_latents.shape[-1]
@@ -1752,7 +1755,7 @@ def rope_encode(self, t, h, w, t_start=0, steps_t=None, steps_h=None, steps_w=No
                 h_shift = (h_scale - 1) / 2
                 w_shift = (w_scale - 1) / 2
                 pose_tf = {"rope_options": {"shift_y": h_shift, "shift_x": POSE_ROPE_W + w_shift, "scale_y": h_scale, "scale_x": w_scale}}
-                parts.append(super().rope_encode(F_pose, H_pose, W_pose, t_start=0, device=device, dtype=dtype, transformer_options=pose_tf))
+                parts.append(super().rope_encode(F_pose, H_pose, W_pose, t_start=video_t_start, device=device, dtype=dtype, transformer_options=pose_tf))
 
             return torch.cat(parts, dim=1)
 
@@ -1761,10 +1764,6 @@ def rope_encode(self, t, h, w, t_start=0, steps_t=None, steps_h=None, steps_w=No
         if pose_latents is None:
             return main_freqs
 
-        ref_t_patches = 0
-        if reference_latent is not None:
-            ref_t_patches = (reference_latent.shape[2] + (self.patch_size[0] // 2)) // self.patch_size[0]
-
         F_pose, H_pose, W_pose = pose_latents.shape[-3], pose_latents.shape[-2], pose_latents.shape[-1]
 
         # if pose is at half resolution, scale_y/scale_x=2 stretches the position range to cover the same RoPE extent as the main frames

comfy/model_base.py

@@ -1747,10 +1747,14 @@ def extra_conds(self, **kwargs):
 
         reference_latents = kwargs.get("reference_latents", None)
         if reference_latents is not None:
-            ref_latent = self.process_latent_in(reference_latents[-1])
-            ref_mask = torch.ones_like(ref_latent[:, :4])
-            ref_latent = torch.cat([ref_latent, ref_mask], dim=1)
-            out['reference_latent'] = comfy.conds.CONDRegular(ref_latent)
+            # SCAIL-2 multi-reference: reference_latents[0] is the primary ref, [1:] are additional
+            # references. Stack as [additional..., primary] so the primary stays adjacent to the video.
+            ordered = list(reference_latents[1:]) + list(reference_latents[:1])
+            stacked = []
+            for lat in ordered:
+                lat = self.process_latent_in(lat)
+                stacked.append(torch.cat([lat, torch.ones_like(lat[:, :4])], dim=1))
+            out['reference_latent'] = comfy.conds.CONDRegular(torch.cat(stacked, dim=2))
 
         pose_latents = kwargs.get("pose_video_latent", None)
         if pose_latents is not None:
@@ -1792,6 +1796,7 @@ def extra_conds(self, **kwargs):
         if driving_mask_28ch is not None:
             out['sam_latents'] = comfy.conds.CONDRegular(driving_mask_28ch.movedim(1, 2).contiguous())
 
+        # ref_mask_28ch holds one identity mask per stacked reference frame (additional refs first, then the primary ref), followed by zeros over the video frames.
         ref_mask_28ch = kwargs.get("ref_mask_28ch", None)
         if ref_mask_28ch is not None:
             out['ref_mask_latents'] = comfy.conds.CONDRegular(ref_mask_28ch.movedim(1, 2).contiguous())
@@ -1819,10 +1824,11 @@ def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, dev
             # Return sliced view omitting retain_index_list
             return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=0)
         if cond_key == "ref_mask_latents" and hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor):
-            # The ref mask is just a single frame padded with frames of zeros, so just grab the first frames for all windows
+            # The ref mask is N leading ref frames padded with frames of zeros, so just grab the first frames for all windows
             full_ref_mask = cond_value.cond
             video_frame_count = x_in.shape[2]
-            if full_ref_mask.shape[2] != video_frame_count + 1:
+            ref_frame_count = full_ref_mask.shape[2] - video_frame_count
+            if ref_frame_count < 1:
                 return None
             window_length = len(window.index_list)
 
@@ -1831,7 +1837,7 @@ def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, dev
             if anchor_index is not None and anchor_index >= 0:
                 window_length += 1
 
-            window_ref_mask = full_ref_mask[:, :, :window_length + 1].to(device)
+            window_ref_mask = full_ref_mask[:, :, :window_length + ref_frame_count].to(device)
             return cond_value._copy_with(window_ref_mask)
 
         return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)

comfy/sd.py

@@ -1622,6 +1622,10 @@ class EmptyClass:
                 clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
                 clip_target.clip = comfy.text_encoders.ideogram4.te_qwen3vl(**llama_detect(clip_data))
                 clip_target.tokenizer = comfy.text_encoders.ideogram4.Ideogram4Qwen3VLTokenizer
+            elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2):  # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
+                klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
+                clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
+                clip_target.tokenizer = comfy.text_encoders.flux.KleinTokenizer8B if te_model == TEModel.QWEN3VL_8B else comfy.text_encoders.flux.KleinTokenizer
             else:
                 clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
                 qwen3vl_type = {TEModel.QWEN3VL_4B: "qwen3vl_4b", TEModel.QWEN3VL_8B: "qwen3vl_8b"}[te_model]