Disable gradients on diffusion models. (#15218)

This commit is contained in:
comfyanonymous 2026-08-01 20:15:51 -07:00 committed by GitHub
parent e0951a9c7e
commit 532a16f3b9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions

View File

@ -169,6 +169,7 @@ class BaseModel(torch.nn.Module):
else:
operations = model_config.custom_operations
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
self.diffusion_model.requires_grad_(False)
self.diffusion_model.eval()
if comfy.model_management.force_channels_last():
self.diffusion_model.to(memory_format=torch.channels_last)