1
0
mirror of https://github.com/kijai/ComfyUI-WanVideoWrapper.git synced 2026-01-26 23:41:35 +03:00

Fix uni3c reloading

This commit is contained in:
kijai
2025-12-14 16:48:55 +02:00
parent 0fa5383106
commit 3611341339
3 changed files with 6 additions and 6 deletions

View File

@@ -863,7 +863,7 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
except Exception:
block_idx = None
if "loras" in name:
if "loras" in name or "uni3c" in name:
continue
# GGUF: skip GGUFParameter params

View File

@@ -890,7 +890,7 @@ class WanVideoSampler:
#uni3c
uni3c_data = uni3c_data_input = None
if uni3c_embeds is not None:
transformer.controlnet = uni3c_embeds["controlnet"]
transformer.uni3c_controlnet = uni3c_embeds["controlnet"]
render_latent = uni3c_embeds["render_latent"].to(device)
if render_latent.shape != noise.shape:
render_latent = torch.nn.functional.interpolate(render_latent, size=(noise.shape[1], noise.shape[2], noise.shape[3]), mode='trilinear', align_corners=False)

View File

@@ -2939,15 +2939,15 @@ class WanModel(torch.nn.Module):
if uni3c_data is not None:
if (uni3c_data["start"] <= current_step_percentage <= uni3c_data["end"]) or \
(uni3c_data["end"] > 0 and current_step == 0 and current_step_percentage >= uni3c_data["start"]):
self.controlnet.to(self.main_device)
self.uni3c_controlnet.to(self.main_device)
with torch.autocast(device_type=mm.get_autocast_device(device), dtype=self.base_dtype, enabled=True):
uni3c_controlnet_states = self.controlnet(
render_latent=render_latent.to(self.main_device, self.controlnet.dtype),
uni3c_controlnet_states = self.uni3c_controlnet(
render_latent=render_latent.to(self.main_device, self.uni3c_controlnet.dtype),
render_mask=uni3c_data["render_mask"],
camera_embedding=uni3c_data["camera_embedding"],
temb=e.to(self.main_device),
device=self.offload_device)
self.controlnet.to(self.offload_device)
self.uni3c_controlnet.to(self.offload_device)
# Asynchronous block offloading with CUDA streams and events
if torch.cuda.is_available():