mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-12 12:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a --force-channels-last to inference models in channel last mode.

Browse Source
This commit is contained in:
comfyanonymous
2024-06-15 01:08:12 -04:00
parent 0e06b370db
commit 0ec513d877
3 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/model_base.py
View File

@@ -66,6 +66,9 @@ class BaseModel(torch.nn.Module):
else:
operations = comfy.ops.disable_weight_init
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
if comfy.model_management.force_channels_last():
self.diffusion_model.to(memory_format=torch.channels_last)
logging.debug("using channels last mode for diffusion model")
self.model_type = model_type
self.model_sampling = model_sampling(model_config, model_type)