Rewrite causual forcing using custom sampler with KSampler node.

comfy/k_diffusion/sampling.py

@@ -1810,3 +1810,84 @@ def sample_sa_solver(model, x, sigmas, extra_args=None, callback=None, disable=F
 def sample_sa_solver_pece(model, x, sigmas, extra_args=None, callback=None, disable=False, tau_func=None, s_noise=1.0, noise_sampler=None, predictor_order=3, corrector_order=4, simple_order_2=False):
     """Stochastic Adams Solver with PECE (Predict–Evaluate–Correct–Evaluate) mode (NeurIPS 2023)."""
     return sample_sa_solver(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, tau_func=tau_func, s_noise=s_noise, noise_sampler=noise_sampler, predictor_order=predictor_order, corrector_order=corrector_order, use_pece=True, simple_order_2=simple_order_2)
+
+
+@torch.no_grad()
+def sample_ar_video(model, x, sigmas, extra_args=None, callback=None, disable=None):
+    """
+    Autoregressive video sampler: block-by-block denoising with KV cache
+    and flow-match re-noising for Causal Forcing / Self-Forcing models.
+    """
+    extra_args = {} if extra_args is None else extra_args
+    model_options = extra_args.get("model_options", {})
+    transformer_options = model_options.get("transformer_options", {})
+    ar_config = transformer_options.get("ar_config", {})
+
+    num_frame_per_block = ar_config.get("num_frame_per_block", 1)
+    seed = extra_args.get("seed", 0)
+
+    bs, c, lat_t, lat_h, lat_w = x.shape
+    frame_seq_len = (lat_h // 2) * (lat_w // 2)
+    num_blocks = lat_t // num_frame_per_block
+
+    inner_model = model.inner_model.inner_model
+    causal_model = inner_model.diffusion_model
+    device = x.device
+    model_dtype = inner_model.get_dtype()
+
+    kv_caches = causal_model.init_kv_caches(bs, lat_t * frame_seq_len, device, model_dtype)
+    crossattn_caches = causal_model.init_crossattn_caches(bs, device, model_dtype)
+
+    output = torch.zeros_like(x)
+    s_in = x.new_ones([x.shape[0]])
+    current_start_frame = 0
+    num_sigma_steps = len(sigmas) - 1
+    total_real_steps = num_blocks * num_sigma_steps
+    step_count = 0
+
+    for block_idx in trange(num_blocks, disable=disable):
+        bf = num_frame_per_block
+        fs, fe = current_start_frame, current_start_frame + bf
+        noisy_input = x[:, :, fs:fe]
+
+        ar_state = {
+            "start_frame": current_start_frame,
+            "kv_caches": kv_caches,
+            "crossattn_caches": crossattn_caches,
+        }
+        transformer_options["ar_state"] = ar_state
+
+        for i in range(num_sigma_steps):
+            denoised = model(noisy_input, sigmas[i] * s_in, **extra_args)
+
+            if callback is not None:
+                # Scale step_count to [0, num_sigma_steps) so the progress bar fills gradually
+                scaled_i = step_count * num_sigma_steps // total_real_steps
+                callback({"x": noisy_input, "i": scaled_i, "sigma": sigmas[i],
+                          "sigma_hat": sigmas[i], "denoised": denoised})
+
+            if sigmas[i + 1] == 0:
+                noisy_input = denoised
+            else:
+                sigma_next = sigmas[i + 1]
+                torch.manual_seed(seed + block_idx * 1000 + i)
+                fresh_noise = torch.randn_like(denoised)
+                noisy_input = (1.0 - sigma_next) * denoised + sigma_next * fresh_noise
+
+                for cache in kv_caches:
+                    cache["end"].fill_(cache["end"].item() - bf * frame_seq_len)
+
+            step_count += 1
+
+        output[:, :, fs:fe] = noisy_input
+
+        # Cache update: run model at t=0 with clean output to fill KV cache
+        for cache in kv_caches:
+            cache["end"].fill_(cache["end"].item() - bf * frame_seq_len)
+        zero_sigma = sigmas.new_zeros([1])
+        _ = model(noisy_input, zero_sigma * s_in, **extra_args)
+
+        current_start_frame += bf
+
+    transformer_options.pop("ar_state", None)
+    return output

comfy/ldm/wan/ar_model.py

@@ -281,7 +281,7 @@ class CausalWanModel(torch.nn.Module):
 
         # Per-frame time embedding → [B, block_frames, 6, dim]
         e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()))
+            sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).to(dtype=x.dtype))
         e = e.reshape(timestep.shape[0], -1, e.shape[-1])
         e0 = self.time_projection(e).unflatten(2, (6, self.dim))
 
@@ -351,8 +351,20 @@ class CausalWanModel(torch.nn.Module):
     def head_dim(self):
         return self.dim // self.num_heads
 
-    # Standard forward for non-causal use (compatibility with ComfyUI infrastructure)
     def forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, **kwargs):
+        ar_state = transformer_options.get("ar_state")
+        if ar_state is not None:
+            bs = x.shape[0]
+            block_frames = x.shape[2]
+            t_per_frame = timestep.unsqueeze(1).expand(bs, block_frames)
+            return self.forward_block(
+                x=x, timestep=t_per_frame, context=context,
+                start_frame=ar_state["start_frame"],
+                kv_caches=ar_state["kv_caches"],
+                crossattn_caches=ar_state["crossattn_caches"],
+                clip_fea=clip_fea,
+            )
+
         bs, c, t, h, w = x.shape
         x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
 
@@ -369,7 +381,7 @@ class CausalWanModel(torch.nn.Module):
         freqs = self.rope_encode(t_len, h, w, device=x.device, dtype=x.dtype)
 
         e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()))
+            sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).to(dtype=x.dtype))
         e = e.reshape(timestep.shape[0], -1, e.shape[-1])
         e0 = self.time_projection(e).unflatten(2, (6, self.dim))
 

comfy/model_base.py

@@ -42,6 +42,7 @@ import comfy.ldm.cosmos.predict2
 import comfy.ldm.lumina.model
 import comfy.ldm.wan.model
 import comfy.ldm.wan.model_animate
+import comfy.ldm.wan.ar_model
 import comfy.ldm.hunyuan3d.model
 import comfy.ldm.hidream.model
 import comfy.ldm.chroma.model
@@ -1365,6 +1366,13 @@ class WAN21(BaseModel):
         return out
 
 
+class WAN21_CausalAR(WAN21):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super(WAN21, self).__init__(model_config, model_type, device=device,
+                                    unet_model=comfy.ldm.wan.ar_model.CausalWanModel)
+        self.image_to_video = False
+
+
 class WAN21_Vace(WAN21):
     def __init__(self, model_config, model_type=ModelType.FLOW, image_to_video=False, device=None):
         super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.VaceWanModel)

comfy/samplers.py

@@ -723,7 +723,8 @@ KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_c
                   "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu",
                   "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_2m_sde_heun", "dpmpp_2m_sde_heun_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
                   "ipndm", "ipndm_v", "deis", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp",
-                  "gradient_estimation", "gradient_estimation_cfg_pp", "er_sde", "seeds_2", "seeds_3", "sa_solver", "sa_solver_pece"]
+                  "gradient_estimation", "gradient_estimation_cfg_pp", "er_sde", "seeds_2", "seeds_3", "sa_solver", "sa_solver_pece",
+                  "ar_video"]
 
 class KSAMPLER(Sampler):
     def __init__(self, sampler_function, extra_options={}, inpaint_options={}):

comfy/supported_models.py

@@ -1167,6 +1167,15 @@ class WAN21_T2V(supported_models_base.BASE):
         t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}umt5xxl.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.wan.WanT5Tokenizer, comfy.text_encoders.wan.te(**t5_detect))
 
+class WAN21_CausalAR_T2V(WAN21_T2V):
+    sampling_settings = {
+        "shift": 5.0,
+    }
+
+    def get_model(self, state_dict, prefix="", device=None):
+        return model_base.WAN21_CausalAR(self, device=device)
+
+
 class WAN21_I2V(WAN21_T2V):
     unet_config = {
         "image_model": "wan2.1",