From 21f04fe632a5192d8f29f1bc0c852b24eb9dce2f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 14 Jun 2023 19:46:08 -0400 Subject: [PATCH] Disable default weight values in unet conv2d for faster loading. --- comfy/ldm/modules/diffusionmodules/util.py | 6 +++--- comfy/ops.py | 4 ++++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/modules/diffusionmodules/util.py b/comfy/ldm/modules/diffusionmodules/util.py index 82ea3f0a..d6a4778e 100644 --- a/comfy/ldm/modules/diffusionmodules/util.py +++ b/comfy/ldm/modules/diffusionmodules/util.py @@ -16,7 +16,7 @@ import numpy as np from einops import repeat from comfy.ldm.util import instantiate_from_config - +import comfy.ops def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): if schedule == "linear": @@ -233,7 +233,7 @@ def conv_nd(dims, *args, **kwargs): if dims == 1: return nn.Conv1d(*args, **kwargs) elif dims == 2: - return nn.Conv2d(*args, **kwargs) + return comfy.ops.Conv2d(*args, **kwargs) elif dims == 3: return nn.Conv3d(*args, **kwargs) raise ValueError(f"unsupported dimensions: {dims}") @@ -243,7 +243,7 @@ def linear(*args, **kwargs): """ Create a linear module. """ - return nn.Linear(*args, **kwargs) + return comfy.ops.Linear(*args, **kwargs) def avg_pool_nd(dims, *args, **kwargs): diff --git a/comfy/ops.py b/comfy/ops.py index 0654dbcd..c39b994a 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -15,3 +15,7 @@ class Linear(torch.nn.Module): def forward(self, input): return torch.nn.functional.linear(input, self.weight, self.bias) + +class Conv2d(torch.nn.Conv2d): + def reset_parameters(self): + return None