Try to fix some issues with the seedvr VAE. (#14877)

comfy/ldm/seedvr/vae.py

@@ -30,7 +30,7 @@ from enum import Enum
 import logging
 import comfy.model_management
 import comfy.ops
-ops = comfy.ops.disable_weight_init
+ops = comfy.ops.manual_cast
 
 
 def _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap, temporal_scale=1):
@@ -103,11 +103,10 @@ def tiled_vae(
     storage_device = vae_model.device
     result = None
     count = None
-    def run_temporal_chunks(spatial_tile, model=vae_model, device=storage_device):
-        device = torch.device(device)
-        t_chunk = spatial_tile.to(device=device, dtype=next(model.parameters()).dtype, non_blocking=True).contiguous()
+    def run_temporal_chunks(spatial_tile, model=vae_model):
+        t_chunk = spatial_tile.contiguous()
         old_device = getattr(model, "device", None)
-        model.device = device
+        model.device = t_chunk.device
         old_slicing_min_size = getattr(model, slicing_attr, None)
         if old_slicing_min_size is not None and slicing_min_size is not None:
             if slicing_min_size <= 0:
@@ -397,7 +396,7 @@ class Attention(nn.Module):
 
 def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
     input_dtype = x.dtype
-    if isinstance(norm_layer, (ops.LayerNorm, ops.RMSNorm)):
+    if isinstance(norm_layer, (nn.LayerNorm, nn.RMSNorm)):
         if x.ndim == 4:
             x = x.permute(0, 2, 3, 1)
             x = norm_layer(x)
@@ -408,14 +407,14 @@ def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
             x = norm_layer(x)
             x = x.permute(0, 4, 1, 2, 3)
             return x.to(input_dtype)
-    if isinstance(norm_layer, (ops.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
+    if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
         if x.ndim <= 4:
             return norm_layer(x).to(input_dtype)
         if x.ndim == 5:
             b, c, t, h, w = x.shape
             x = x.transpose(1, 2).reshape(b * t, c, h, w)
             memory_occupy = x.numel() * x.element_size() / 1024**3
-            if isinstance(norm_layer, ops.GroupNorm) and memory_occupy > get_norm_limit():
+            if isinstance(norm_layer, nn.GroupNorm) and memory_occupy > get_norm_limit():
                 num_chunks = min(BYTEDANCE_GN_CHUNKS_FP16 if x.element_size() == 2 else BYTEDANCE_GN_CHUNKS_FP32, norm_layer.num_groups)
                 if norm_layer.num_groups % num_chunks != 0:
                     raise ValueError(
@@ -423,9 +422,9 @@ def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
                     )
                 num_groups_per_chunk = norm_layer.num_groups // num_chunks
 
+                weights = comfy.ops.cast_to_input(norm_layer.weight, x).chunk(num_chunks, dim=0)
+                biases = comfy.ops.cast_to_input(norm_layer.bias, x).chunk(num_chunks, dim=0)
                 x = list(x.chunk(num_chunks, dim=1))
-                weights = norm_layer.weight.chunk(num_chunks, dim=0)
-                biases = norm_layer.bias.chunk(num_chunks, dim=0)
                 for i, (w, bias) in enumerate(zip(weights, biases)):
                     x[i] = F.group_norm(x[i], num_groups_per_chunk, w, bias, norm_layer.eps)
                     x[i] = x[i].to(input_dtype)
@@ -1459,7 +1458,6 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
     def _encode_with_raw_latent(self, x):
         if x.ndim == 4:
             x = x.unsqueeze(2)
-        x = x.to(dtype=next(self.parameters()).dtype)
         self.device = x.device
         p = super().encode(x)
         z = p.squeeze(2)

tests-unit/comfy_test/test_seedvr2_dtype.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 
 from comfy.cli_args import args as cli_args
 
@@ -48,3 +49,31 @@ def test_seedvr2_vae_decode_memory_covers_full_frame_lab_transfer():
     assert estimate == 101 * 960 * 1280 * 160
     assert estimate > 15 * 1024 ** 3
     assert estimate > old_estimate * 100
+
+
+def test_seedvr2_vae_encode_preserves_compute_dtype(monkeypatch):
+    wrapper = seedvr_vae.VideoAutoencoderKLWrapper.__new__(seedvr_vae.VideoAutoencoderKLWrapper)
+    nn.Module.__init__(wrapper)
+    wrapper._dummy = nn.Parameter(torch.empty(1, dtype=torch.float16))
+    input_dtype = None
+
+    def encode(self, x):
+        nonlocal input_dtype
+        input_dtype = x.dtype
+        return x
+
+    monkeypatch.setattr(seedvr_vae.VideoAutoencoderKL, "encode", encode)
+
+    x = torch.zeros((1, 3, 1, 8, 8), dtype=torch.float32)
+    wrapper._encode_with_raw_latent(x)
+
+    assert input_dtype == torch.float32
+
+
+def test_seedvr2_vae_ops_cast_weights_to_compute_dtype():
+    attention = seedvr_vae.Attention(query_dim=4, heads=1, dim_head=4).to(torch.float16)
+    hidden_states = torch.zeros((1, 2, 4), dtype=torch.float32)
+
+    output = attention(hidden_states)
+
+    assert output.dtype == torch.float32

tests-unit/comfy_test/test_seedvr2_vae_tiled.py

@@ -122,6 +122,31 @@ def test_tiled_vae_encode_uses_tensor_return_without_indexing():
     assert tuple(out.shape) == (2, _LATENT_CHANNELS, 1, 8, 8)
 
 
+def test_tiled_vae_preserves_compute_dtype_with_different_parameter_dtype():
+    class DummyVAE(nn.Module):
+        spatial_downsample_factor = 8
+        temporal_downsample_factor = 4
+        slicing_sample_min_size = 8
+
+        def __init__(self):
+            super().__init__()
+            self.device = torch.device("cpu")
+            self._dummy = nn.Parameter(torch.zeros(1, dtype=torch.float16))
+            self.input_dtype = None
+
+        def encode(self, t_chunk):
+            self.input_dtype = t_chunk.dtype
+            b, _, _, h, w = t_chunk.shape
+            return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=t_chunk.dtype)
+
+    vae = DummyVAE()
+    x = torch.zeros((1, 3, 1, 64, 64), dtype=torch.float32)
+
+    tiled_vae(x, vae, tile_size=(64, 64), tile_overlap=(16, 16), encode=True)
+
+    assert vae.input_dtype == torch.float32
+
+
 def test_tiled_vae_preserves_input_dtype_on_single_tile():
     class FloatOutputVAEModel(torch.nn.Module):
         def __init__(self):