Fix save node to force-load model weights before saving

Model weights stay on meta device (no data) until ComfyUI's model
management loads them. Use load_models_gpu with force_full_load
to materialize weights before extracting state dict.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-12 15:20:46 +01:00
parent 3acf57b7a1
commit d2c96726e5
+7
View File
@@ -3,6 +3,7 @@ import json
import logging import logging
import torch import torch
import folder_paths import folder_paths
import comfy.model_management
from safetensors.torch import save_file from safetensors.torch import save_file
from comfy.utils import ProgressBar from comfy.utils import ProgressBar
@@ -76,6 +77,12 @@ class WanVideoSaveMergedModel:
# Extract state dict from the diffusion model (keys are already bare, # Extract state dict from the diffusion model (keys are already bare,
# e.g. "blocks.0.self_attn.k.weight" — matching original checkpoint format) # e.g. "blocks.0.self_attn.k.weight" — matching original checkpoint format)
diffusion_model = model.model.diffusion_model diffusion_model = model.model.diffusion_model
# Force ComfyUI to load model weights into real memory.
# Without this, weights stay on meta device (shape-only, no data).
log.info("Loading model weights into memory for saving...")
comfy.model_management.load_models_gpu([model], force_full_load=True)
state_dict = diffusion_model.state_dict() state_dict = diffusion_model.state_dict()
target_dtype = dtype_map.get(save_dtype) target_dtype = dtype_map.get(save_dtype)