fix: keep LDF transformer dtypes aligned
This commit is contained in:
+6
-5
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user