Skip to content

Commit d5acb41

Browse files
patrickvonplatenpatrickvonplaten
authored
Finalize ldm (huggingface#96)
* upload * make checkpoint work * finalize
1 parent 6cabc59 commit d5acb41

7 files changed

Lines changed: 999 additions & 135 deletions

File tree

‎src/diffusers/__init__.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,15 @@
77
__version__ = "0.0.4"
88

99
from .modeling_utils import ModelMixin
10-
from .models import AutoencoderKL, NCSNpp, UNetLDMModel, UNetModel, UNetUnconditionalModel, VQModel
10+
from .models import (
11+
AutoencoderKL,
12+
NCSNpp,
13+
UNetConditionalModel,
14+
UNetLDMModel,
15+
UNetModel,
16+
UNetUnconditionalModel,
17+
VQModel,
18+
)
1119
from .pipeline_utils import DiffusionPipeline
1220
from .pipelines import (
1321
DDIMPipeline,

‎src/diffusers/models/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
# limitations under the License.
1818

1919
from .unet import UNetModel
20+
from .unet_conditional import UNetConditionalModel
2021
from .unet_glide import GlideSuperResUNetModel, GlideTextToImageUNetModel, GlideUNetModel
2122
from .unet_ldm import UNetLDMModel
2223
from .unet_sde_score_estimation import NCSNpp

‎src/diffusers/models/attention.py‎

Lines changed: 22 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def __init__(
4242
self.value = nn.Linear(channels, channels)
4343

4444
self.rescale_output_factor = rescale_output_factor
45-
self.proj_attn = zero_module(nn.Linear(channels, channels, 1))
45+
self.proj_attn = nn.Linear(channels, channels, 1)
4646

4747
def transpose_for_scores(self, projection: torch.Tensor) -> torch.Tensor:
4848
new_projection_shape = projection.size()[:-1] + (self.num_heads, -1)
@@ -147,6 +147,8 @@ class SpatialTransformer(nn.Module):
147147

148148
def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_dim=None):
149149
super().__init__()
150+
self.n_heads = n_heads
151+
self.d_head = d_head
150152
self.in_channels = in_channels
151153
inner_dim = n_heads * d_head
152154
self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
@@ -160,7 +162,7 @@ def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_d
160162
]
161163
)
162164

163-
self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
165+
self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
164166

165167
def forward(self, x, context=None):
166168
# note: if no context is given, cross-attention defaults to self-attention
@@ -175,6 +177,12 @@ def forward(self, x, context=None):
175177
x = self.proj_out(x)
176178
return x + x_in
177179

180+
def set_weight(self, layer):
181+
self.norm = layer.norm
182+
self.proj_in = layer.proj_in
183+
self.transformer_blocks = layer.transformer_blocks
184+
self.proj_out = layer.proj_out
185+
178186

179187
class BasicTransformerBlock(nn.Module):
180188
def __init__(self, dim, n_heads, d_head, dropout=0.0, context_dim=None, gated_ff=True, checkpoint=True):
@@ -270,14 +278,15 @@ def forward(self, x):
270278
return self.net(x)
271279

272280

273-
# TODO(Patrick) - this can and should be removed
274-
def zero_module(module):
275-
"""
276-
Zero out the parameters of a module and return it.
277-
"""
278-
for p in module.parameters():
279-
p.detach().zero_()
280-
return module
281+
# feedforward
282+
class GEGLU(nn.Module):
283+
def __init__(self, dim_in, dim_out):
284+
super().__init__()
285+
self.proj = nn.Linear(dim_in, dim_out * 2)
286+
287+
def forward(self, x):
288+
x, gate = self.proj(x).chunk(2, dim=-1)
289+
return x * F.gelu(gate)
281290

282291

283292
# TODO(Patrick) - remove once all weights have been converted -> not needed anymore then
@@ -298,17 +307,6 @@ def default(val, d):
298307
return d() if isfunction(d) else d
299308

300309

301-
# feedforward
302-
class GEGLU(nn.Module):
303-
def __init__(self, dim_in, dim_out):
304-
super().__init__()
305-
self.proj = nn.Linear(dim_in, dim_out * 2)
306-
307-
def forward(self, x):
308-
x, gate = self.proj(x).chunk(2, dim=-1)
309-
return x * F.gelu(gate)
310-
311-
312310
# the main attention block that is used for all models
313311
class AttentionBlock(nn.Module):
314312
"""
@@ -348,7 +346,7 @@ def __init__(
348346
if encoder_channels is not None:
349347
self.encoder_kv = nn.Conv1d(encoder_channels, channels * 2, 1)
350348

351-
self.proj = zero_module(nn.Conv1d(channels, channels, 1))
349+
self.proj = nn.Conv1d(channels, channels, 1)
352350

353351
self.overwrite_qkv = overwrite_qkv
354352
self.overwrite_linear = overwrite_linear
@@ -370,7 +368,7 @@ def __init__(
370368

371369
self.GroupNorm_0 = nn.GroupNorm(num_groups=num_groups, num_channels=channels, eps=1e-6)
372370
else:
373-
self.proj_out = zero_module(nn.Conv1d(channels, channels, 1))
371+
self.proj_out = nn.Conv1d(channels, channels, 1)
374372
self.set_weights(self)
375373

376374
self.is_overwritten = False
@@ -385,7 +383,7 @@ def set_weights(self, module):
385383
self.qkv.weight.data = qkv_weight
386384
self.qkv.bias.data = qkv_bias
387385

388-
proj_out = zero_module(nn.Conv1d(self.channels, self.channels, 1))
386+
proj_out = nn.Conv1d(self.channels, self.channels, 1)
389387
proj_out.weight.data = module.proj_out.weight.data[:, :, :, 0]
390388
proj_out.bias.data = module.proj_out.bias.data
391389

0 commit comments

Comments
 (0)