From 897e343268084e13d1f6f5c058a2a9bd11162167 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Sat, 15 Aug 2026 21:58:53 +0200
Subject: [PATCH] fix: keep LDF transformer dtypes aligned

---
 ldf_backend.py            | 11 ++++++-----
 tests/test_ldf_backend.py | 25 +++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/ldf_backend.py b/ldf_backend.py
index 9e33e77..f0c6fef 100644
--- a/ldf_backend.py
+++ b/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
-        # (time embedding, norms, scale/shift) in FP32. Passing dtype here would
-        # flatten that mixed-precision policy and diffusers warns that results
-        # can become inconsistent.
-        self.transformer.to(device=target)
+        # LDF's custom Wan fork expects its complete condition embedder to use
+        # one dtype: time_embedder feeds time_proj directly without an explicit
+        # cast. Match the official generator's model.to(..., dtype=BF16) call;
+        # preserving diffusers' generic FP32-module policy leaves that pair as
+        # 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
diff --git a/tests/test_ldf_backend.py b/tests/test_ldf_backend.py
index 6f28fd3..a3b4844 100644
--- a/tests/test_ldf_backend.py
+++ b/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()