You've already forked ComfyUI-WanVideoWrapper
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:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user