diff --git a/diffsynth/vram_management/layers.py b/diffsynth/vram_management/layers.py index 6f6f768..3b6e2bb 100644 --- a/diffsynth/vram_management/layers.py +++ b/diffsynth/vram_management/layers.py @@ -75,6 +75,9 @@ def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, for name, module in model.named_children(): for source_module, target_module in module_map.items(): if isinstance(module, source_module): + if "rope_embedder" in name or "patch_embedding" in name: + continue + num_param = sum(p.numel() for p in module.parameters()) if max_num_param is not None and total_num_param + num_param > max_num_param: module_config_ = overflow_module_config diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index e56689f..bb82ddd 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -115,7 +115,7 @@ def rope_apply(x, grid_sizes, freqs): freqs[0][:f].view(*f_size).expand(f, h, w, -1), freqs[1][:h].view(*h_size).expand(f, h, w, -1), freqs[2][:w].view(*w_size).expand(f, h, w, -1) - ], dim=-1).reshape(seq_len, 1, -1) + ], dim=-1).reshape(seq_len, 1, -1).to(x.device) @torch.compiler.disable() def view_as_real_no_compile(x_i):