From 8d6d0ae3d84a9cd6fb0220420ab49dbdf517a082 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 1 Mar 2025 03:23:36 +0200
Subject: [PATCH] Allow loading comfy models

---
 nodes.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/nodes.py b/nodes.py
index 0ded5ca..245ec19 100644
--- a/nodes.py
+++ b/nodes.py
@@ -296,6 +296,16 @@ class WanVideoModelLoader:
 
         model_path = folder_paths.get_full_path_or_raise("diffusion_models", model)
         sd = load_torch_file(model_path, device=transformer_load_device, safe_load=True)
+
+        first_key = next(iter(sd))
+        if first_key.startswith("model.diffusion_model."):
+            # Create new state dict with modified keys
+            new_sd = {}
+            for key, value in sd.items():
+                new_key = key.replace("model.diffusion_model.", "", 1)
+                new_sd[new_key] = value
+            sd = new_sd
+
         
         dim = sd["patch_embedding.weight"].shape[0]
         in_channels = sd["patch_embedding.weight"].shape[1]