unet_glide.py 19.5 KB
Newer Older
anton-l's avatar
anton-l committed
1
import torch
anton-l's avatar
anton-l committed
2
3
4
import torch.nn as nn
import torch.nn.functional as F

Patrick von Platen's avatar
Patrick von Platen committed
5
6
from ..configuration_utils import ConfigMixin
from ..modeling_utils import ModelMixin
Patrick von Platen's avatar
Patrick von Platen committed
7
from .attention import AttentionBlock
8
from .embeddings import get_timestep_embedding
9
from .resnet import Downsample2D, ResnetBlock2D, Upsample2D
Patrick von Platen's avatar
Patrick von Platen committed
10
from .unet_new import UNetMidBlock2D
anton-l's avatar
anton-l committed
11

anton-l's avatar
anton-l committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

def convert_module_to_f16(l):
    """
    Convert primitive modules to float16.
    """
    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
        l.weight.data = l.weight.data.half()
        if l.bias is not None:
            l.bias.data = l.bias.data.half()


def convert_module_to_f32(l):
    """
    Convert primitive modules to float32, undoing convert_module_to_f16().
    """
    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
        l.weight.data = l.weight.data.float()
        if l.bias is not None:
            l.bias.data = l.bias.data.float()


def conv_nd(dims, *args, **kwargs):
    """
    Create a 1D, 2D, or 3D convolution module.
    """
    if dims == 1:
        return nn.Conv1d(*args, **kwargs)
    elif dims == 2:
        return nn.Conv2d(*args, **kwargs)
    elif dims == 3:
        return nn.Conv3d(*args, **kwargs)
    raise ValueError(f"unsupported dimensions: {dims}")


def linear(*args, **kwargs):
    """
    Create a linear module.
    """
    return nn.Linear(*args, **kwargs)


class GroupNorm32(nn.GroupNorm):
    def __init__(self, num_groups, num_channels, swish, eps=1e-5):
        super().__init__(num_groups=num_groups, num_channels=num_channels, eps=eps)
        self.swish = swish

    def forward(self, x):
        y = super().forward(x.float()).to(x.dtype)
        if self.swish == 1.0:
            y = F.silu(y)
        elif self.swish:
            y = y * F.sigmoid(y * float(self.swish))
        return y


def normalization(channels, swish=0.0):
    """
    Make a standard normalization layer, with an optional swish activation.

Patrick von Platen's avatar
Patrick von Platen committed
71
    :param channels: number of input channels. :return: an nn.Module for normalization.
anton-l's avatar
anton-l committed
72
73
74
75
76
77
78
79
80
81
82
83
84
    """
    return GroupNorm32(num_channels=channels, num_groups=32, swish=swish)


def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


85
class TimestepEmbedSequential(nn.Sequential):
anton-l's avatar
anton-l committed
86
    """
Patrick von Platen's avatar
Patrick von Platen committed
87
    A sequential module that passes timestep embeddings to the children that support it as an extra input.
anton-l's avatar
anton-l committed
88
89
90
91
    """

    def forward(self, x, emb, encoder_out=None):
        for layer in self:
92
            if isinstance(layer, ResnetBlock2D) or isinstance(layer, TimestepEmbedSequential):
anton-l's avatar
anton-l committed
93
94
95
96
97
98
99
100
                x = layer(x, emb)
            elif isinstance(layer, AttentionBlock):
                x = layer(x, encoder_out)
            else:
                x = layer(x)
        return x


Patrick von Platen's avatar
Patrick von Platen committed
101
class GlideUNetModel(ModelMixin, ConfigMixin):
anton-l's avatar
anton-l committed
102
103
104
    """
    The full UNet model with attention and timestep embedding.

Patrick von Platen's avatar
Patrick von Platen committed
105
106
    :param in_channels: channels in the input Tensor. :param model_channels: base channel count for the model. :param
    out_channels: channels in the output Tensor. :param num_res_blocks: number of residual blocks per downsample.
anton-l's avatar
anton-l committed
107
    :param attention_resolutions: a collection of downsample rates at which
Patrick von Platen's avatar
Patrick von Platen committed
108
109
110
111
        attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x
        downsampling, attention will be used.
    :param dropout: the dropout probability. :param channel_mult: channel multiplier for each level of the UNet. :param
    conv_resample: if True, use learned convolutions for upsampling and
anton-l's avatar
anton-l committed
112
        downsampling.
Patrick von Platen's avatar
Patrick von Platen committed
113
114
    :param dims: determines if the signal is 1D, 2D, or 3D. :param num_classes: if specified (as an int), then this
    model will be
anton-l's avatar
anton-l committed
115
        class-conditional with `num_classes` classes.
Patrick von Platen's avatar
Patrick von Platen committed
116
117
    :param use_checkpoint: use gradient checkpointing to reduce memory usage. :param num_heads: the number of attention
    heads in each attention layer. :param num_heads_channels: if specified, ignore num_heads and instead use
anton-l's avatar
anton-l committed
118
119
120
                               a fixed channel width per attention head.
    :param num_heads_upsample: works with num_heads to set a different number
                               of heads for upsampling. Deprecated.
Patrick von Platen's avatar
Patrick von Platen committed
121
122
    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. :param resblock_updown: use residual blocks
    for up/downsampling.
anton-l's avatar
anton-l committed
123
124
125
126
    """

    def __init__(
        self,
127
        in_channels=3,
anton-l's avatar
anton-l committed
128
        resolution=64,
129
130
131
132
        model_channels=192,
        out_channels=6,
        num_res_blocks=3,
        attention_resolutions=(2, 4, 8),
anton-l's avatar
anton-l committed
133
134
135
136
137
138
139
140
141
142
143
        dropout=0,
        channel_mult=(1, 2, 4, 8),
        conv_resample=True,
        dims=2,
        use_checkpoint=False,
        use_fp16=False,
        num_heads=1,
        num_head_channels=-1,
        num_heads_upsample=-1,
        use_scale_shift_norm=False,
        resblock_updown=False,
anton-l's avatar
anton-l committed
144
        transformer_dim=None,
anton-l's avatar
anton-l committed
145
146
147
148
149
150
151
    ):
        super().__init__()

        if num_heads_upsample == -1:
            num_heads_upsample = num_heads

        self.in_channels = in_channels
anton-l's avatar
anton-l committed
152
        self.resolution = resolution
anton-l's avatar
anton-l committed
153
154
155
156
157
158
159
160
        self.model_channels = model_channels
        self.out_channels = out_channels
        self.num_res_blocks = num_res_blocks
        self.attention_resolutions = attention_resolutions
        self.dropout = dropout
        self.channel_mult = channel_mult
        self.conv_resample = conv_resample
        self.use_checkpoint = use_checkpoint
anton-l's avatar
Style  
anton-l committed
161
        # self.dtype = torch.float16 if use_fp16 else torch.float32
anton-l's avatar
anton-l committed
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
        self.num_heads = num_heads
        self.num_head_channels = num_head_channels
        self.num_heads_upsample = num_heads_upsample

        time_embed_dim = model_channels * 4
        self.time_embed = nn.Sequential(
            linear(model_channels, time_embed_dim),
            nn.SiLU(),
            linear(time_embed_dim, time_embed_dim),
        )

        ch = input_ch = int(channel_mult[0] * model_channels)
        self.input_blocks = nn.ModuleList([TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))])
        self._feature_size = ch
        input_block_chans = [ch]
        ds = 1
        for level, mult in enumerate(channel_mult):
            for _ in range(num_res_blocks):
                layers = [
Patrick von Platen's avatar
Patrick von Platen committed
181
                    ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
182
183
184
185
186
187
                        in_channels=ch,
                        out_channels=mult * model_channels,
                        dropout=dropout,
                        temb_channels=time_embed_dim,
                        eps=1e-5,
                        non_linearity="silu",
Patrick von Platen's avatar
Patrick von Platen committed
188
                        time_embedding_norm="scale_shift" if use_scale_shift_norm else "default",
Patrick von Platen's avatar
Patrick von Platen committed
189
                        overwrite_for_glide=True,
anton-l's avatar
anton-l committed
190
191
192
193
194
195
196
197
198
                    )
                ]
                ch = int(mult * model_channels)
                if ds in attention_resolutions:
                    layers.append(
                        AttentionBlock(
                            ch,
                            num_heads=num_heads,
                            num_head_channels=num_head_channels,
anton-l's avatar
anton-l committed
199
                            encoder_channels=transformer_dim,
anton-l's avatar
anton-l committed
200
201
202
203
204
205
206
207
208
                        )
                    )
                self.input_blocks.append(TimestepEmbedSequential(*layers))
                self._feature_size += ch
                input_block_chans.append(ch)
            if level != len(channel_mult) - 1:
                out_ch = ch
                self.input_blocks.append(
                    TimestepEmbedSequential(
Patrick von Platen's avatar
Patrick von Platen committed
209
                        ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
210
                            in_channels=ch,
anton-l's avatar
anton-l committed
211
                            out_channels=out_ch,
Patrick von Platen's avatar
Patrick von Platen committed
212
213
214
215
                            dropout=dropout,
                            temb_channels=time_embed_dim,
                            eps=1e-5,
                            non_linearity="silu",
Patrick von Platen's avatar
Patrick von Platen committed
216
                            time_embedding_norm="scale_shift" if use_scale_shift_norm else "default",
Patrick von Platen's avatar
Patrick von Platen committed
217
                            overwrite_for_glide=True,
Patrick von Platen's avatar
Patrick von Platen committed
218
                            down=True,
anton-l's avatar
anton-l committed
219
220
                        )
                        if resblock_updown
221
                        else Downsample2D(ch, use_conv=conv_resample, out_channels=out_ch, padding=1, name="op")
anton-l's avatar
anton-l committed
222
223
224
225
226
227
228
                    )
                )
                ch = out_ch
                input_block_chans.append(ch)
                ds *= 2
                self._feature_size += ch

Patrick von Platen's avatar
Patrick von Platen committed
229
230
231
232
233
234
235
236
237
238
239
240
241
242
        self.mid = UNetMidBlock2D(
            in_channels=ch,
            dropout=dropout,
            temb_channels=time_embed_dim,
            resnet_eps=1e-5,
            resnet_act_fn="silu",
            resnet_time_scale_shift="scale_shift" if use_scale_shift_norm else "default",
            attn_num_heads=num_heads,
            attn_num_head_channels=num_head_channels,
            attn_encoder_channels=transformer_dim,
        )

        # TODO(Patrick) - delete after weight conversion
        # init to be able to overwrite `self.mid`
anton-l's avatar
anton-l committed
243
        self.middle_block = TimestepEmbedSequential(
Patrick von Platen's avatar
Patrick von Platen committed
244
            ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
245
246
247
248
249
                in_channels=ch,
                dropout=dropout,
                temb_channels=time_embed_dim,
                eps=1e-5,
                non_linearity="silu",
Patrick von Platen's avatar
Patrick von Platen committed
250
                time_embedding_norm="scale_shift" if use_scale_shift_norm else "default",
Patrick von Platen's avatar
Patrick von Platen committed
251
                overwrite_for_glide=True,
anton-l's avatar
anton-l committed
252
253
254
255
256
            ),
            AttentionBlock(
                ch,
                num_heads=num_heads,
                num_head_channels=num_head_channels,
anton-l's avatar
anton-l committed
257
                encoder_channels=transformer_dim,
anton-l's avatar
anton-l committed
258
            ),
Patrick von Platen's avatar
Patrick von Platen committed
259
            ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
260
261
262
263
264
                in_channels=ch,
                dropout=dropout,
                temb_channels=time_embed_dim,
                eps=1e-5,
                non_linearity="silu",
Patrick von Platen's avatar
Patrick von Platen committed
265
                time_embedding_norm="scale_shift" if use_scale_shift_norm else "default",
Patrick von Platen's avatar
Patrick von Platen committed
266
                overwrite_for_glide=True,
Patrick von Platen's avatar
Patrick von Platen committed
267
            ),
anton-l's avatar
anton-l committed
268
        )
269
270
271
        self.mid.resnets[0] = self.middle_block[0]
        self.mid.attentions[0] = self.middle_block[1]
        self.mid.resnets[1] = self.middle_block[2]
Patrick von Platen's avatar
Patrick von Platen committed
272

anton-l's avatar
anton-l committed
273
274
275
276
277
278
279
        self._feature_size += ch

        self.output_blocks = nn.ModuleList([])
        for level, mult in list(enumerate(channel_mult))[::-1]:
            for i in range(num_res_blocks + 1):
                ich = input_block_chans.pop()
                layers = [
Patrick von Platen's avatar
Patrick von Platen committed
280
                    ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
281
282
283
284
285
286
                        in_channels=ch + ich,
                        out_channels=model_channels * mult,
                        dropout=dropout,
                        temb_channels=time_embed_dim,
                        eps=1e-5,
                        non_linearity="silu",
Patrick von Platen's avatar
Patrick von Platen committed
287
                        time_embedding_norm="scale_shift" if use_scale_shift_norm else "default",
Patrick von Platen's avatar
Patrick von Platen committed
288
289
                        overwrite_for_glide=True,
                    ),
anton-l's avatar
anton-l committed
290
291
292
293
294
295
296
297
                ]
                ch = int(model_channels * mult)
                if ds in attention_resolutions:
                    layers.append(
                        AttentionBlock(
                            ch,
                            num_heads=num_heads_upsample,
                            num_head_channels=num_head_channels,
anton-l's avatar
anton-l committed
298
                            encoder_channels=transformer_dim,
anton-l's avatar
anton-l committed
299
300
301
302
303
                        )
                    )
                if level and i == num_res_blocks:
                    out_ch = ch
                    layers.append(
Patrick von Platen's avatar
Patrick von Platen committed
304
                        ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
305
                            in_channels=ch,
anton-l's avatar
anton-l committed
306
                            out_channels=out_ch,
Patrick von Platen's avatar
Patrick von Platen committed
307
308
309
310
                            dropout=dropout,
                            temb_channels=time_embed_dim,
                            eps=1e-5,
                            non_linearity="silu",
Patrick von Platen's avatar
Patrick von Platen committed
311
                            time_embedding_norm="scale_shift" if use_scale_shift_norm else "default",
Patrick von Platen's avatar
Patrick von Platen committed
312
                            overwrite_for_glide=True,
anton-l's avatar
anton-l committed
313
314
315
                            up=True,
                        )
                        if resblock_updown
316
                        else Upsample2D(ch, use_conv=conv_resample, out_channels=out_ch)
anton-l's avatar
anton-l committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
                    )
                    ds //= 2
                self.output_blocks.append(TimestepEmbedSequential(*layers))
                self._feature_size += ch

        self.out = nn.Sequential(
            normalization(ch, swish=1.0),
            nn.Identity(),
            zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),
        )
        self.use_fp16 = use_fp16

    def convert_to_fp16(self):
        """
        Convert the torso of the model to float16.
        """
        self.input_blocks.apply(convert_module_to_f16)
        self.middle_block.apply(convert_module_to_f16)
        self.output_blocks.apply(convert_module_to_f16)

    def convert_to_fp32(self):
        """
        Convert the torso of the model to float32.
        """
        self.input_blocks.apply(convert_module_to_f32)
        self.middle_block.apply(convert_module_to_f32)
        self.output_blocks.apply(convert_module_to_f32)

345
    def forward(self, x, timesteps):
anton-l's avatar
anton-l committed
346
347
348
        """
        Apply the model to an input batch.

Patrick von Platen's avatar
Patrick von Platen committed
349
350
        :param x: an [N x C x ...] Tensor of inputs. :param timesteps: a 1-D batch of timesteps. :param y: an [N]
        Tensor of labels, if class-conditional. :return: an [N x C x ...] Tensor of outputs.
anton-l's avatar
anton-l committed
351
        """
anton-l's avatar
anton-l committed
352
353

        hs = []
354
355
356
        emb = self.time_embed(
            get_timestep_embedding(timesteps, self.model_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
        )
anton-l's avatar
anton-l committed
357
358
359
360
361

        h = x.type(self.dtype)
        for module in self.input_blocks:
            h = module(h, emb)
            hs.append(h)
Patrick von Platen's avatar
Patrick von Platen committed
362
        h = self.mid(h, emb)
anton-l's avatar
anton-l committed
363
364
365
366
367
368
369
        for module in self.output_blocks:
            h = torch.cat([h, hs.pop()], dim=1)
            h = module(h, emb)
        h = h.type(x.dtype)
        return self.out(h)


Patrick von Platen's avatar
Patrick von Platen committed
370
class GlideTextToImageUNetModel(GlideUNetModel):
anton-l's avatar
anton-l committed
371
372
373
374
375
376
    """
    A UNetModel that performs super-resolution.

    Expects an extra kwarg `low_res` to condition on a low-resolution image.
    """

377
    def __init__(
Patrick von Platen's avatar
Patrick von Platen committed
378
379
        self,
        in_channels=3,
anton-l's avatar
anton-l committed
380
        resolution=64,
Patrick von Platen's avatar
Patrick von Platen committed
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
        model_channels=192,
        out_channels=6,
        num_res_blocks=3,
        attention_resolutions=(2, 4, 8),
        dropout=0,
        channel_mult=(1, 2, 4, 8),
        conv_resample=True,
        dims=2,
        use_checkpoint=False,
        use_fp16=False,
        num_heads=1,
        num_head_channels=-1,
        num_heads_upsample=-1,
        use_scale_shift_norm=False,
        resblock_updown=False,
        transformer_dim=512,
397
398
399
    ):
        super().__init__(
            in_channels=in_channels,
anton-l's avatar
anton-l committed
400
            resolution=resolution,
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
            model_channels=model_channels,
            out_channels=out_channels,
            num_res_blocks=num_res_blocks,
            attention_resolutions=attention_resolutions,
            dropout=dropout,
            channel_mult=channel_mult,
            conv_resample=conv_resample,
            dims=dims,
            use_checkpoint=use_checkpoint,
            use_fp16=use_fp16,
            num_heads=num_heads,
            num_head_channels=num_head_channels,
            num_heads_upsample=num_heads_upsample,
            use_scale_shift_norm=use_scale_shift_norm,
            resblock_updown=resblock_updown,
Patrick von Platen's avatar
Patrick von Platen committed
416
            transformer_dim=transformer_dim,
417
        )
418
        self.register_to_config(
419
            in_channels=in_channels,
anton-l's avatar
anton-l committed
420
            resolution=resolution,
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
            model_channels=model_channels,
            out_channels=out_channels,
            num_res_blocks=num_res_blocks,
            attention_resolutions=attention_resolutions,
            dropout=dropout,
            channel_mult=channel_mult,
            conv_resample=conv_resample,
            dims=dims,
            use_checkpoint=use_checkpoint,
            use_fp16=use_fp16,
            num_heads=num_heads,
            num_head_channels=num_head_channels,
            num_heads_upsample=num_heads_upsample,
            use_scale_shift_norm=use_scale_shift_norm,
            resblock_updown=resblock_updown,
Patrick von Platen's avatar
Patrick von Platen committed
436
            transformer_dim=transformer_dim,
437
        )
anton-l's avatar
anton-l committed
438

439
        self.transformer_proj = nn.Linear(transformer_dim, self.model_channels * 4)
anton-l's avatar
anton-l committed
440

441
442
    def forward(self, sample, timestep, transformer_out=None):
        timesteps = timestep
443
        x = sample
anton-l's avatar
anton-l committed
444
        hs = []
445
446
447
        emb = self.time_embed(
            get_timestep_embedding(timesteps, self.model_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
        )
anton-l's avatar
anton-l committed
448
449
450
451
452

        # project the last token
        transformer_proj = self.transformer_proj(transformer_out[:, -1])
        transformer_out = transformer_out.permute(0, 2, 1)  # NLC -> NCL

453
454
        emb = emb + transformer_proj.to(emb)

anton-l's avatar
anton-l committed
455
        h = x
anton-l's avatar
anton-l committed
456
        for module in self.input_blocks:
457
            h = module(h, emb, transformer_out)
anton-l's avatar
anton-l committed
458
            hs.append(h)
Patrick von Platen's avatar
Patrick von Platen committed
459
        h = self.mid(h, emb, transformer_out)
anton-l's avatar
anton-l committed
460
        for module in self.output_blocks:
anton-l's avatar
anton-l committed
461
462
            other = hs.pop()
            h = torch.cat([h, other], dim=1)
463
            h = module(h, emb, transformer_out)
anton-l's avatar
anton-l committed
464
        return self.out(h)
anton-l's avatar
anton-l committed
465
466


Patrick von Platen's avatar
Patrick von Platen committed
467
class GlideSuperResUNetModel(GlideUNetModel):
anton-l's avatar
anton-l committed
468
469
470
471
472
473
    """
    A UNetModel that performs super-resolution.

    Expects an extra kwarg `low_res` to condition on a low-resolution image.
    """

474
    def __init__(
Patrick von Platen's avatar
Patrick von Platen committed
475
476
        self,
        in_channels=3,
anton-l's avatar
anton-l committed
477
        resolution=256,
Patrick von Platen's avatar
Patrick von Platen committed
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
        model_channels=192,
        out_channels=6,
        num_res_blocks=3,
        attention_resolutions=(2, 4, 8),
        dropout=0,
        channel_mult=(1, 2, 4, 8),
        conv_resample=True,
        dims=2,
        use_checkpoint=False,
        use_fp16=False,
        num_heads=1,
        num_head_channels=-1,
        num_heads_upsample=-1,
        use_scale_shift_norm=False,
        resblock_updown=False,
493
494
495
    ):
        super().__init__(
            in_channels=in_channels,
anton-l's avatar
anton-l committed
496
            resolution=resolution,
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
            model_channels=model_channels,
            out_channels=out_channels,
            num_res_blocks=num_res_blocks,
            attention_resolutions=attention_resolutions,
            dropout=dropout,
            channel_mult=channel_mult,
            conv_resample=conv_resample,
            dims=dims,
            use_checkpoint=use_checkpoint,
            use_fp16=use_fp16,
            num_heads=num_heads,
            num_head_channels=num_head_channels,
            num_heads_upsample=num_heads_upsample,
            use_scale_shift_norm=use_scale_shift_norm,
            resblock_updown=resblock_updown,
        )
513
        self.register_to_config(
514
            in_channels=in_channels,
anton-l's avatar
anton-l committed
515
            resolution=resolution,
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
            model_channels=model_channels,
            out_channels=out_channels,
            num_res_blocks=num_res_blocks,
            attention_resolutions=attention_resolutions,
            dropout=dropout,
            channel_mult=channel_mult,
            conv_resample=conv_resample,
            dims=dims,
            use_checkpoint=use_checkpoint,
            use_fp16=use_fp16,
            num_heads=num_heads,
            num_head_channels=num_head_channels,
            num_heads_upsample=num_heads_upsample,
            use_scale_shift_norm=use_scale_shift_norm,
            resblock_updown=resblock_updown,
        )
anton-l's avatar
anton-l committed
532

533
534
    def forward(self, sample, timestep, low_res=None):
        timesteps = timestep
535
        x = sample
anton-l's avatar
anton-l committed
536
537
538
        _, _, new_height, new_width = x.shape
        upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
        x = torch.cat([x, upsampled], dim=1)
539
540

        hs = []
541
542
543
        emb = self.time_embed(
            get_timestep_embedding(timesteps, self.model_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
        )
544
545
546
547
548

        h = x
        for module in self.input_blocks:
            h = module(h, emb)
            hs.append(h)
549
        h = self.mid(h, emb)
550
551
552
553
        for module in self.output_blocks:
            h = torch.cat([h, hs.pop()], dim=1)
            h = module(h, emb)

Patrick von Platen's avatar
Patrick von Platen committed
554
        return self.out(h)