diff --git a/comfy/ldm/seedvr/vae.py b/comfy/ldm/seedvr/vae.py index c9f430184..7a8070b65 100644 --- a/comfy/ldm/seedvr/vae.py +++ b/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) diff --git a/tests-unit/comfy_test/test_seedvr2_dtype.py b/tests-unit/comfy_test/test_seedvr2_dtype.py index 8e08b6dde..d743cc848 100644 --- a/tests-unit/comfy_test/test_seedvr2_dtype.py +++ b/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 diff --git a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py index a2866b609..d64f51918 100644 --- a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py +++ b/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):