nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implement conditional offloading for diffusion model

Browse Source 
Add condition to skip offloading if load_device is main_device.
This commit is contained in:
Jean J. de Jong
2026-01-21 17:33:58 +01:00
committed by GitHub
parent 9f9a8e71c2
commit b99eac73da
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes_model_loading.py
+8 -4
View File
@@ -1803,10 +1803,14 @@ class WanVideoModelLoader:
)
if merge_loras and lora is not None:
log.info(f"Moving diffusion model from {patcher.model.diffusion_model.device} to {offload_device}")
patcher.model.diffusion_model.to(offload_device)
gc.collect()
mm.soft_empty_cache()
# Skip offloading if load_device is main_device (for unified memory systems like AMD Strix Halo)
if load_device != "main_device":
log.info(f"Moving diffusion model from {patcher.model.diffusion_model.device} to {offload_device}")
patcher.model.diffusion_model.to(offload_device)
gc.collect()
mm.soft_empty_cache()
else:
log.info(f"Skipping offload (load_device=main_device, keeping model on {patcher.model.diffusion_model.device})")
patcher.model["base_dtype"] = base_dtype
patcher.model["weight_dtype"] = weight_dtype