fix: keep LDF transformer dtypes aligned

ldf_backend.py

@@ -187,11 +187,12 @@ class LDFVFIModel:
         target = torch.device(device)
         if target.type == "cuda" and not _cuda_bf16_supported(target):
             raise RuntimeError("LDF-VFI requires a CUDA GPU with BF16 support (Ampere or newer)")
-        # from_pretrained(torch_dtype=...) keeps numerically sensitive modules
+        # LDF's custom Wan fork expects its complete condition embedder to use
-        # (time embedding, norms, scale/shift) in FP32. Passing dtype here would
+        # one dtype: time_embedder feeds time_proj directly without an explicit
-        # flatten that mixed-precision policy and diffusers warns that results
+        # cast. Match the official generator's model.to(..., dtype=BF16) call;
-        # can become inconsistent.
+        # preserving diffusers' generic FP32-module policy leaves that pair as
-        self.transformer.to(device=target)
+        # Float/BFloat16 and fails in the first sampling step.
+        self.transformer.to(device=target, dtype=self.dtype)
         self._move_auxiliary_models(target)
         self.device = str(target)
         return self

tests/test_ldf_backend.py

@@ -1,5 +1,6 @@
 import torch
 import pytest
+from types import SimpleNamespace
 
 from ldf_backend import LDFVFIModel
 
@@ -21,6 +22,15 @@ class _RecordingConditionalVAE:
         return torch.zeros(1, 3, 40, 1, 1)
 
 
+class _RecordingTransformer:
+    def __init__(self):
+        self.to_kwargs = None
+
+    def to(self, **kwargs):
+        self.to_kwargs = kwargs
+        return self
+
+
 class _ShapeOnlyLDF(LDFVFIModel):
     """Exercise sequence chunking without loading the multi-GB checkpoint."""
 
@@ -80,6 +90,21 @@ def test_decode_tiles_40_frame_720p_condition_for_conditional_vae():
     assert result.shape == (40, 3, 1, 1)
 
 
+def test_device_move_casts_complete_ldf_transformer_to_bfloat16():
+    model = LDFVFIModel.__new__(LDFVFIModel)
+    transformer = _RecordingTransformer()
+    model.model = SimpleNamespace(transformer=transformer)
+    model.dtype = torch.bfloat16
+    model._move_auxiliary_models = lambda device: None
+
+    model.to("cpu")
+
+    assert transformer.to_kwargs == {
+        "device": torch.device("cpu"),
+        "dtype": torch.bfloat16,
+    }
+
+
 def test_decode_rejects_condition_length_that_cannot_tile():
     model = LDFVFIModel.__new__(LDFVFIModel)
     model.vae = _RecordingConditionalVAE()