twnatelo/multi-view-diffusion
04
1import inspect2import math3from inspect import isfunction4from typing import Any, Callable, List, Optional, Union5 6import numpy as np7import torch8import torch.nn as nn9import torch.nn.functional as F10# require xformers!11import xformers12import xformers.ops13from diffusers import AutoencoderKL, DiffusionPipeline14from diffusers.configuration_utils import ConfigMixin, FrozenDict15from diffusers.models.modeling_utils import ModelMixin16from diffusers.schedulers import DDIMScheduler17from diffusers.utils import (deprecate, is_accelerate_available,18 is_accelerate_version, logging)19from diffusers.utils.torch_utils import randn_tensor20from einops import rearrange, repeat21from kiui.cam import orbit_camera22from transformers import (CLIPImageProcessor, CLIPTextModel, CLIPTokenizer,23 CLIPVisionModel)24 25 26def get_camera(27 num_frames,28 elevation=15,29 azimuth_start=0,30 azimuth_span=360,31 blender_coord=True,32 extra_view=False,33):34 angle_gap = azimuth_span / num_frames35 cameras = []36 for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):37 38 pose = orbit_camera(39 -elevation, azimuth, radius=140 ) # kiui's elevation is negated, [4, 4]41 42 # opengl to blender43 if blender_coord:44 pose[2] *= -145 pose[[1, 2]] = pose[[2, 1]]46 47 cameras.append(pose.flatten())48 49 if extra_view:50 cameras.append(np.zeros_like(cameras[0]))51 52 return torch.from_numpy(np.stack(cameras, axis=0)).float() # [num_frames, 16]53 54 55def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):56 """57 Create sinusoidal timestep embeddings.58 :param timesteps: a 1-D Tensor of N indices, one per batch element.59 These may be fractional.60 :param dim: the dimension of the output.61 :param max_period: controls the minimum frequency of the embeddings.62 :return: an [N x dim] Tensor of positional embeddings.63 """64 if not repeat_only:65 half = dim // 266 freqs = torch.exp(67 -math.log(max_period)68 * torch.arange(start=0, end=half, dtype=torch.float32)69 / half70 ).to(device=timesteps.device)71 args = timesteps[:, None] * freqs[None]72 embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)73 if dim % 2:74 embedding = torch.cat(75 [embedding, torch.zeros_like(embedding[:, :1])], dim=-176 )77 else:78 embedding = repeat(timesteps, "b -> b d", d=dim)79 # import pdb; pdb.set_trace()80 return embedding81 82 83def zero_module(module):84 """85 Zero out the parameters of a module and return it.86 """87 for p in module.parameters():88 p.detach().zero_()89 return module90 91 92def conv_nd(dims, *args, **kwargs):93 """94 Create a 1D, 2D, or 3D convolution module.95 """96 if dims == 1:97 return nn.Conv1d(*args, **kwargs)98 elif dims == 2:99 return nn.Conv2d(*args, **kwargs)100 elif dims == 3:101 return nn.Conv3d(*args, **kwargs)102 raise ValueError(f"unsupported dimensions: {dims}")103 104 105def avg_pool_nd(dims, *args, **kwargs):106 """107 Create a 1D, 2D, or 3D average pooling module.108 """109 if dims == 1:110 return nn.AvgPool1d(*args, **kwargs)111 elif dims == 2:112 return nn.AvgPool2d(*args, **kwargs)113 elif dims == 3:114 return nn.AvgPool3d(*args, **kwargs)115 raise ValueError(f"unsupported dimensions: {dims}")116 117 118def default(val, d):119 if val is not None:120 return val121 return d() if isfunction(d) else d122 123 124class GEGLU(nn.Module):125 def __init__(self, dim_in, dim_out):126 super().__init__()127 self.proj = nn.Linear(dim_in, dim_out * 2)128 129 def forward(self, x):130 x, gate = self.proj(x).chunk(2, dim=-1)131 return x * F.gelu(gate)132 133 134class FeedForward(nn.Module):135 def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):136 super().__init__()137 inner_dim = int(dim * mult)138 dim_out = default(dim_out, dim)139 project_in = (140 nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())141 if not glu142 else GEGLU(dim, inner_dim)143 )144 145 self.net = nn.Sequential(146 project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)147 )148 149 def forward(self, x):150 return self.net(x)151 152 153class MemoryEfficientCrossAttention(nn.Module):154 # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223155 def __init__(156 self,157 query_dim,158 context_dim=None,159 heads=8,160 dim_head=64,161 dropout=0.0,162 ip_dim=0,163 ip_weight=1,164 ):165 super().__init__()166 167 inner_dim = dim_head * heads168 context_dim = default(context_dim, query_dim)169 170 self.heads = heads171 self.dim_head = dim_head172 173 self.ip_dim = ip_dim174 self.ip_weight = ip_weight175 176 if self.ip_dim > 0:177 self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)178 self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)179 180 self.to_q = nn.Linear(query_dim, inner_dim, bias=False)181 self.to_k = nn.Linear(context_dim, inner_dim, bias=False)182 self.to_v = nn.Linear(context_dim, inner_dim, bias=False)183 184 self.to_out = nn.Sequential(185 nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)186 )187 self.attention_op: Optional[Any] = None188 189 def forward(self, x, context=None):190 q = self.to_q(x)191 context = default(context, x)192 193 if self.ip_dim > 0:194 # context： [B, 77 + 16(ip), 1024]195 token_len = context.shape[1]196 context_ip = context[:, -self.ip_dim :, :]197 k_ip = self.to_k_ip(context_ip)198 v_ip = self.to_v_ip(context_ip)199 context = context[:, : (token_len - self.ip_dim), :]200 201 k = self.to_k(context)202 v = self.to_v(context)203 204 b, _, _ = q.shape205 q, k, v = map(206 lambda t: t.unsqueeze(3)207 .reshape(b, t.shape[1], self.heads, self.dim_head)208 .permute(0, 2, 1, 3)209 .reshape(b * self.heads, t.shape[1], self.dim_head)210 .contiguous(),211 (q, k, v),212 )213 214 # actually compute the attention, what we cannot get enough of215 out = xformers.ops.memory_efficient_attention(216 q, k, v, attn_bias=None, op=self.attention_op217 )218 219 if self.ip_dim > 0:220 k_ip, v_ip = map(221 lambda t: t.unsqueeze(3)222 .reshape(b, t.shape[1], self.heads, self.dim_head)223 .permute(0, 2, 1, 3)224 .reshape(b * self.heads, t.shape[1], self.dim_head)225 .contiguous(),226 (k_ip, v_ip),227 )228 # actually compute the attention, what we cannot get enough of229 out_ip = xformers.ops.memory_efficient_attention(230 q, k_ip, v_ip, attn_bias=None, op=self.attention_op231 )232 out = out + self.ip_weight * out_ip233 234 out = (235 out.unsqueeze(0)236 .reshape(b, self.heads, out.shape[1], self.dim_head)237 .permute(0, 2, 1, 3)238 .reshape(b, out.shape[1], self.heads * self.dim_head)239 )240 return self.to_out(out)241 242 243class BasicTransformerBlock3D(nn.Module):244 245 def __init__(246 self,247 dim,248 n_heads,249 d_head,250 context_dim,251 dropout=0.0,252 gated_ff=True,253 ip_dim=0,254 ip_weight=1,255 ):256 super().__init__()257 258 self.attn1 = MemoryEfficientCrossAttention(259 query_dim=dim,260 context_dim=None, # self-attention261 heads=n_heads,262 dim_head=d_head,263 dropout=dropout,264 )265 self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)266 self.attn2 = MemoryEfficientCrossAttention(267 query_dim=dim,268 context_dim=context_dim,269 heads=n_heads,270 dim_head=d_head,271 dropout=dropout,272 # ip only applies to cross-attention273 ip_dim=ip_dim,274 ip_weight=ip_weight,275 )276 self.norm1 = nn.LayerNorm(dim)277 self.norm2 = nn.LayerNorm(dim)278 self.norm3 = nn.LayerNorm(dim)279 280 def forward(self, x, context=None, num_frames=1):281 x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()282 x = self.attn1(self.norm1(x), context=None) + x283 x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()284 x = self.attn2(self.norm2(x), context=context) + x285 x = self.ff(self.norm3(x)) + x286 return x287 288 289class SpatialTransformer3D(nn.Module):290 291 def __init__(292 self,293 in_channels,294 n_heads,295 d_head,296 context_dim, # cross attention input dim297 depth=1,298 dropout=0.0,299 ip_dim=0,300 ip_weight=1,301 ):302 super().__init__()303 304 if not isinstance(context_dim, list):305 context_dim = [context_dim]306 307 self.in_channels = in_channels308 309 inner_dim = n_heads * d_head310 self.norm = nn.GroupNorm(311 num_groups=32, num_channels=in_channels, eps=1e-6, affine=True312 )313 self.proj_in = nn.Linear(in_channels, inner_dim)314 315 self.transformer_blocks = nn.ModuleList(316 [317 BasicTransformerBlock3D(318 inner_dim,319 n_heads,320 d_head,321 context_dim=context_dim[d],322 dropout=dropout,323 ip_dim=ip_dim,324 ip_weight=ip_weight,325 )326 for d in range(depth)327 ]328 )329 330 self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))331 332 def forward(self, x, context=None, num_frames=1):333 # note: if no context is given, cross-attention defaults to self-attention334 if not isinstance(context, list):335 context = [context]336 b, c, h, w = x.shape337 x_in = x338 x = self.norm(x)339 x = rearrange(x, "b c h w -> b (h w) c").contiguous()340 x = self.proj_in(x)341 for i, block in enumerate(self.transformer_blocks):342 x = block(x, context=context[i], num_frames=num_frames)343 x = self.proj_out(x)344 x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()345 346 return x + x_in347 348 349class PerceiverAttention(nn.Module):350 def __init__(self, *, dim, dim_head=64, heads=8):351 super().__init__()352 self.scale = dim_head**-0.5353 self.dim_head = dim_head354 self.heads = heads355 inner_dim = dim_head * heads356 357 self.norm1 = nn.LayerNorm(dim)358 self.norm2 = nn.LayerNorm(dim)359 360 self.to_q = nn.Linear(dim, inner_dim, bias=False)361 self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)362 self.to_out = nn.Linear(inner_dim, dim, bias=False)363 364 def forward(self, x, latents):365 """366 Args:367 x (torch.Tensor): image features368 shape (b, n1, D)369 latent (torch.Tensor): latent features370 shape (b, n2, D)371 """372 x = self.norm1(x)373 latents = self.norm2(latents)374 375 b, h, _ = latents.shape376 377 q = self.to_q(latents)378 kv_input = torch.cat((x, latents), dim=-2)379 k, v = self.to_kv(kv_input).chunk(2, dim=-1)380 381 q, k, v = map(382 lambda t: t.reshape(b, t.shape[1], self.heads, -1)383 .transpose(1, 2)384 .reshape(b, self.heads, t.shape[1], -1)385 .contiguous(),386 (q, k, v),387 )388 389 # attention390 scale = 1 / math.sqrt(math.sqrt(self.dim_head))391 weight = (q * scale) @ (k * scale).transpose(392 -2, -1393 ) # More stable with f16 than dividing afterwards394 weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)395 out = weight @ v396 397 out = out.permute(0, 2, 1, 3).reshape(b, h, -1)398 399 return self.to_out(out)400 401 402class Resampler(nn.Module):403 def __init__(404 self,405 dim=1024,406 depth=8,407 dim_head=64,408 heads=16,409 num_queries=8,410 embedding_dim=768,411 output_dim=1024,412 ff_mult=4,413 ):414 super().__init__()415 self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)416 self.proj_in = nn.Linear(embedding_dim, dim)417 self.proj_out = nn.Linear(dim, output_dim)418 self.norm_out = nn.LayerNorm(output_dim)419 420 self.layers = nn.ModuleList([])421 for _ in range(depth):422 self.layers.append(423 nn.ModuleList(424 [425 PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),426 nn.Sequential(427 nn.LayerNorm(dim),428 nn.Linear(dim, dim * ff_mult, bias=False),429 nn.GELU(),430 nn.Linear(dim * ff_mult, dim, bias=False),431 ),432 ]433 )434 )435 436 def forward(self, x):437 latents = self.latents.repeat(x.size(0), 1, 1)438 x = self.proj_in(x)439 for attn, ff in self.layers:440 latents = attn(x, latents) + latents441 latents = ff(latents) + latents442 443 latents = self.proj_out(latents)444 return self.norm_out(latents)445 446 447class CondSequential(nn.Sequential):448 """449 A sequential module that passes timestep embeddings to the children that450 support it as an extra input.451 """452 453 def forward(self, x, emb, context=None, num_frames=1):454 for layer in self:455 if isinstance(layer, ResBlock):456 x = layer(x, emb)457 elif isinstance(layer, SpatialTransformer3D):458 x = layer(x, context, num_frames=num_frames)459 else:460 x = layer(x)461 return x462 463 464class Upsample(nn.Module):465 """466 An upsampling layer with an optional convolution.467 :param channels: channels in the inputs and outputs.468 :param use_conv: a bool determining if a convolution is applied.469 :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then470 upsampling occurs in the inner-two dimensions.471 """472 473 def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):474 super().__init__()475 self.channels = channels476 self.out_channels = out_channels or channels477 self.use_conv = use_conv478 self.dims = dims479 if use_conv:480 self.conv = conv_nd(481 dims, self.channels, self.out_channels, 3, padding=padding482 )483 484 def forward(self, x):485 assert x.shape[1] == self.channels486 if self.dims == 3:487 x = F.interpolate(488 x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"489 )490 else:491 x = F.interpolate(x, scale_factor=2, mode="nearest")492 if self.use_conv:493 x = self.conv(x)494 return x495 496 497class Downsample(nn.Module):498 """499 A downsampling layer with an optional convolution.500 :param channels: channels in the inputs and outputs.501 :param use_conv: a bool determining if a convolution is applied.502 :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then503 downsampling occurs in the inner-two dimensions.504 """505 506 def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):507 super().__init__()508 self.channels = channels509 self.out_channels = out_channels or channels510 self.use_conv = use_conv511 self.dims = dims512 stride = 2 if dims != 3 else (1, 2, 2)513 if use_conv:514 self.op = conv_nd(515 dims,516 self.channels,517 self.out_channels,518 3,519 stride=stride,520 padding=padding,521 )522 else:523 assert self.channels == self.out_channels524 self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)525 526 def forward(self, x):527 assert x.shape[1] == self.channels528 return self.op(x)529 530 531class ResBlock(nn.Module):532 """533 A residual block that can optionally change the number of channels.534 :param channels: the number of input channels.535 :param emb_channels: the number of timestep embedding channels.536 :param dropout: the rate of dropout.537 :param out_channels: if specified, the number of out channels.538 :param use_conv: if True and out_channels is specified, use a spatial539 convolution instead of a smaller 1x1 convolution to change the540 channels in the skip connection.541 :param dims: determines if the signal is 1D, 2D, or 3D.542 :param up: if True, use this block for upsampling.543 :param down: if True, use this block for downsampling.544 """545 546 def __init__(547 self,548 channels,549 emb_channels,550 dropout,551 out_channels=None,552 use_conv=False,553 use_scale_shift_norm=False,554 dims=2,555 up=False,556 down=False,557 ):558 super().__init__()559 self.channels = channels560 self.emb_channels = emb_channels561 self.dropout = dropout562 self.out_channels = out_channels or channels563 self.use_conv = use_conv564 self.use_scale_shift_norm = use_scale_shift_norm565 566 self.in_layers = nn.Sequential(567 nn.GroupNorm(32, channels),568 nn.SiLU(),569 conv_nd(dims, channels, self.out_channels, 3, padding=1),570 )571 572 self.updown = up or down573 574 if up:575 self.h_upd = Upsample(channels, False, dims)576 self.x_upd = Upsample(channels, False, dims)577 elif down:578 self.h_upd = Downsample(channels, False, dims)579 self.x_upd = Downsample(channels, False, dims)580 else:581 self.h_upd = self.x_upd = nn.Identity()582 583 self.emb_layers = nn.Sequential(584 nn.SiLU(),585 nn.Linear(586 emb_channels,587 2 * self.out_channels if use_scale_shift_norm else self.out_channels,588 ),589 )590 self.out_layers = nn.Sequential(591 nn.GroupNorm(32, self.out_channels),592 nn.SiLU(),593 nn.Dropout(p=dropout),594 zero_module(595 conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)596 ),597 )598 599 if self.out_channels == channels:600 self.skip_connection = nn.Identity()601 elif use_conv:602 self.skip_connection = conv_nd(603 dims, channels, self.out_channels, 3, padding=1604 )605 else:606 self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)607 608 def forward(self, x, emb):609 if self.updown:610 in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]611 h = in_rest(x)612 h = self.h_upd(h)613 x = self.x_upd(x)614 h = in_conv(h)615 else:616 h = self.in_layers(x)617 emb_out = self.emb_layers(emb).type(h.dtype)618 while len(emb_out.shape) < len(h.shape):619 emb_out = emb_out[..., None]620 if self.use_scale_shift_norm:621 out_norm, out_rest = self.out_layers[0], self.out_layers[1:]622 scale, shift = torch.chunk(emb_out, 2, dim=1)623 h = out_norm(h) * (1 + scale) + shift624 h = out_rest(h)625 else:626 h = h + emb_out627 h = self.out_layers(h)628 return self.skip_connection(x) + h629 630 631class MultiViewUNetModel(ModelMixin, ConfigMixin):632 """633 The full multi-view UNet model with attention, timestep embedding and camera embedding.634 :param in_channels: channels in the input Tensor.635 :param model_channels: base channel count for the model.636 :param out_channels: channels in the output Tensor.637 :param num_res_blocks: number of residual blocks per downsample.638 :param attention_resolutions: a collection of downsample rates at which639 attention will take place. May be a set, list, or tuple.640 For example, if this contains 4, then at 4x downsampling, attention641 will be used.642 :param dropout: the dropout probability.643 :param channel_mult: channel multiplier for each level of the UNet.644 :param conv_resample: if True, use learned convolutions for upsampling and645 downsampling.646 :param dims: determines if the signal is 1D, 2D, or 3D.647 :param num_classes: if specified (as an int), then this model will be648 class-conditional with `num_classes` classes.649 :param num_heads: the number of attention heads in each attention layer.650 :param num_heads_channels: if specified, ignore num_heads and instead use651 a fixed channel width per attention head.652 :param num_heads_upsample: works with num_heads to set a different number653 of heads for upsampling. Deprecated.654 :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.655 :param resblock_updown: use residual blocks for up/downsampling.656 :param use_new_attention_order: use a different attention pattern for potentially657 increased efficiency.658 :param camera_dim: dimensionality of camera input.659 """660 661 def __init__(662 self,663 image_size,664 in_channels,665 model_channels,666 out_channels,667 num_res_blocks,668 attention_resolutions,669 dropout=0,670 channel_mult=(1, 2, 4, 8),671 conv_resample=True,672 dims=2,673 num_classes=None,674 num_heads=-1,675 num_head_channels=-1,676 num_heads_upsample=-1,677 use_scale_shift_norm=False,678 resblock_updown=False,679 transformer_depth=1,680 context_dim=None,681 n_embed=None,682 num_attention_blocks=None,683 adm_in_channels=None,684 camera_dim=None,685 ip_dim=0, # imagedream uses ip_dim > 0686 ip_weight=1.0,687 **kwargs,688 ):689 super().__init__()690 assert context_dim is not None691 692 if num_heads_upsample == -1:693 num_heads_upsample = num_heads694 695 if num_heads == -1:696 assert (697 num_head_channels != -1698 ), "Either num_heads or num_head_channels has to be set"699 700 if num_head_channels == -1:701 assert (702 num_heads != -1703 ), "Either num_heads or num_head_channels has to be set"704 705 self.image_size = image_size706 self.in_channels = in_channels707 self.model_channels = model_channels708 self.out_channels = out_channels709 if isinstance(num_res_blocks, int):710 self.num_res_blocks = len(channel_mult) * [num_res_blocks]711 else:712 if len(num_res_blocks) != len(channel_mult):713 raise ValueError(714 "provide num_res_blocks either as an int (globally constant) or "715 "as a list/tuple (per-level) with the same length as channel_mult"716 )717 self.num_res_blocks = num_res_blocks718 719 if num_attention_blocks is not None:720 assert len(num_attention_blocks) == len(self.num_res_blocks)721 assert all(722 map(723 lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],724 range(len(num_attention_blocks)),725 )726 )727 print(728 f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "729 f"This option has LESS priority than attention_resolutions {attention_resolutions}, "730 f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "731 f"attention will still not be set."732 )733 734 self.attention_resolutions = attention_resolutions735 self.dropout = dropout736 self.channel_mult = channel_mult737 self.conv_resample = conv_resample738 self.num_classes = num_classes739 self.num_heads = num_heads740 self.num_head_channels = num_head_channels741 self.num_heads_upsample = num_heads_upsample742 self.predict_codebook_ids = n_embed is not None743 744 self.ip_dim = ip_dim745 self.ip_weight = ip_weight746 747 if self.ip_dim > 0:748 self.image_embed = Resampler(749 dim=context_dim,750 depth=4,751 dim_head=64,752 heads=12,753 num_queries=ip_dim, # num token754 embedding_dim=1280,755 output_dim=context_dim,756 ff_mult=4,757 )758 759 time_embed_dim = model_channels * 4760 self.time_embed = nn.Sequential(761 nn.Linear(model_channels, time_embed_dim),762 nn.SiLU(),763 nn.Linear(time_embed_dim, time_embed_dim),764 )765 766 if camera_dim is not None:767 time_embed_dim = model_channels * 4768 self.camera_embed = nn.Sequential(769 nn.Linear(camera_dim, time_embed_dim),770 nn.SiLU(),771 nn.Linear(time_embed_dim, time_embed_dim),772 )773 774 if self.num_classes is not None:775 if isinstance(self.num_classes, int):776 self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)777 elif self.num_classes == "continuous":778 # print("setting up linear c_adm embedding layer")779 self.label_emb = nn.Linear(1, time_embed_dim)780 elif self.num_classes == "sequential":781 assert adm_in_channels is not None782 self.label_emb = nn.Sequential(783 nn.Sequential(784 nn.Linear(adm_in_channels, time_embed_dim),785 nn.SiLU(),786 nn.Linear(time_embed_dim, time_embed_dim),787 )788 )789 else:790 raise ValueError()791 792 self.input_blocks = nn.ModuleList(793 [CondSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))]794 )795 self._feature_size = model_channels796 input_block_chans = [model_channels]797 ch = model_channels798 ds = 1799 for level, mult in enumerate(channel_mult):800 for nr in range(self.num_res_blocks[level]):801 layers: List[Any] = [802 ResBlock(803 ch,804 time_embed_dim,805 dropout,806 out_channels=mult * model_channels,807 dims=dims,808 use_scale_shift_norm=use_scale_shift_norm,809 )810 ]811 ch = mult * model_channels812 if ds in attention_resolutions:813 if num_head_channels == -1:814 dim_head = ch // num_heads815 else:816 num_heads = ch // num_head_channels817 dim_head = num_head_channels818 819 if num_attention_blocks is None or nr < num_attention_blocks[level]:820 layers.append(821 SpatialTransformer3D(822 ch,823 num_heads,824 dim_head,825 context_dim=context_dim,826 depth=transformer_depth,827 ip_dim=self.ip_dim,828 ip_weight=self.ip_weight,829 )830 )831 self.input_blocks.append(CondSequential(*layers))832 self._feature_size += ch833 input_block_chans.append(ch)834 if level != len(channel_mult) - 1:835 out_ch = ch836 self.input_blocks.append(837 CondSequential(838 ResBlock(839 ch,840 time_embed_dim,841 dropout,842 out_channels=out_ch,843 dims=dims,844 use_scale_shift_norm=use_scale_shift_norm,845 down=True,846 )847 if resblock_updown848 else Downsample(849 ch, conv_resample, dims=dims, out_channels=out_ch850 )851 )852 )853 ch = out_ch854 input_block_chans.append(ch)855 ds *= 2856 self._feature_size += ch857 858 if num_head_channels == -1:859 dim_head = ch // num_heads860 else:861 num_heads = ch // num_head_channels862 dim_head = num_head_channels863 864 self.middle_block = CondSequential(865 ResBlock(866 ch,867 time_embed_dim,868 dropout,869 dims=dims,870 use_scale_shift_norm=use_scale_shift_norm,871 ),872 SpatialTransformer3D(873 ch,874 num_heads,875 dim_head,876 context_dim=context_dim,877 depth=transformer_depth,878 ip_dim=self.ip_dim,879 ip_weight=self.ip_weight,880 ),881 ResBlock(882 ch,883 time_embed_dim,884 dropout,885 dims=dims,886 use_scale_shift_norm=use_scale_shift_norm,887 ),888 )889 self._feature_size += ch890 891 self.output_blocks = nn.ModuleList([])892 for level, mult in list(enumerate(channel_mult))[::-1]:893 for i in range(self.num_res_blocks[level] + 1):894 ich = input_block_chans.pop()895 layers = [896 ResBlock(897 ch + ich,898 time_embed_dim,899 dropout,900 out_channels=model_channels * mult,901 dims=dims,902 use_scale_shift_norm=use_scale_shift_norm,903 )904 ]905 ch = model_channels * mult906 if ds in attention_resolutions:907 if num_head_channels == -1:908 dim_head = ch // num_heads909 else:910 num_heads = ch // num_head_channels911 dim_head = num_head_channels912 913 if num_attention_blocks is None or i < num_attention_blocks[level]:914 layers.append(915 SpatialTransformer3D(916 ch,917 num_heads,918 dim_head,919 context_dim=context_dim,920 depth=transformer_depth,921 ip_dim=self.ip_dim,922 ip_weight=self.ip_weight,923 )924 )925 if level and i == self.num_res_blocks[level]:926 out_ch = ch927 layers.append(928 ResBlock(929 ch,930 time_embed_dim,931 dropout,932 out_channels=out_ch,933 dims=dims,934 use_scale_shift_norm=use_scale_shift_norm,935 up=True,936 )937 if resblock_updown938 else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)939 )940 ds //= 2941 self.output_blocks.append(CondSequential(*layers))942 self._feature_size += ch943 944 self.out = nn.Sequential(945 nn.GroupNorm(32, ch),946 nn.SiLU(),947 zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),948 )949 if self.predict_codebook_ids:950 self.id_predictor = nn.Sequential(951 nn.GroupNorm(32, ch),952 conv_nd(dims, model_channels, n_embed, 1),953 # nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits954 )955 956 def forward(957 self,958 x,959 timesteps=None,960 context=None,961 y=None,962 camera=None,963 num_frames=1,964 ip=None,965 ip_img=None,966 **kwargs,967 ):968 """969 Apply the model to an input batch.970 :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).971 :param timesteps: a 1-D batch of timesteps.972 :param context: conditioning plugged in via crossattn973 :param y: an [N] Tensor of labels, if class-conditional.974 :param num_frames: a integer indicating number of frames for tensor reshaping.975 :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).976 """977 assert (978 x.shape[0] % num_frames == 0979 ), "input batch size must be dividable by num_frames!"980 assert (y is not None) == (981 self.num_classes is not None982 ), "must specify y if and only if the model is class-conditional"983 984 hs = []985 986 t_emb = timestep_embedding(987 timesteps, self.model_channels, repeat_only=False988 ).to(x.dtype)989 990 emb = self.time_embed(t_emb)991 992 if self.num_classes is not None:993 assert y is not None994 assert y.shape[0] == x.shape[0]995 emb = emb + self.label_emb(y)996 997 # Add camera embeddings998 if camera is not None:999 emb = emb + self.camera_embed(camera)1000 1001 # imagedream variant1002 if self.ip_dim > 0:1003 x[(num_frames - 1) :: num_frames, :, :, :] = ip_img # place at [4, 9]1004 ip_emb = self.image_embed(ip)1005 context = torch.cat((context, ip_emb), 1)1006 1007 h = x1008 for module in self.input_blocks:1009 h = module(h, emb, context, num_frames=num_frames)1010 hs.append(h)1011 h = self.middle_block(h, emb, context, num_frames=num_frames)1012 for module in self.output_blocks:1013 h = torch.cat([h, hs.pop()], dim=1)1014 h = module(h, emb, context, num_frames=num_frames)1015 h = h.type(x.dtype)1016 if self.predict_codebook_ids:1017 return self.id_predictor(h)1018 else:1019 return self.out(h)1020 1021 1022logger = logging.get_logger(__name__) # pylint: disable=invalid-name1023 1024 1025class MVDreamPipeline(DiffusionPipeline):1026 1027 _optional_components = ["feature_extractor", "image_encoder"]1028 1029 def __init__(1030 self,1031 vae: AutoencoderKL,1032 unet: MultiViewUNetModel,1033 tokenizer: CLIPTokenizer,1034 text_encoder: CLIPTextModel,1035 scheduler: DDIMScheduler,1036 # imagedream variant1037 feature_extractor: CLIPImageProcessor,1038 image_encoder: CLIPVisionModel,1039 requires_safety_checker: bool = False,1040 ):1041 super().__init__()1042 1043 if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: # type: ignore1044 deprecation_message = (1045 f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"1046 f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " # type: ignore1047 "to update the config accordingly as leaving `steps_offset` might led to incorrect results"1048 " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"1049 " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"1050 " file"1051 )1052 deprecate(1053 "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False1054 )1055 new_config = dict(scheduler.config)1056 new_config["steps_offset"] = 11057 scheduler._internal_dict = FrozenDict(new_config)1058 1059 if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: # type: ignore1060 deprecation_message = (1061 f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."1062 " `clip_sample` should be set to False in the configuration file. Please make sure to update the"1063 " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"1064 " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"1065 " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"1066 )1067 deprecate(1068 "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False1069 )1070 new_config = dict(scheduler.config)1071 new_config["clip_sample"] = False1072 scheduler._internal_dict = FrozenDict(new_config)1073 1074 self.register_modules(1075 vae=vae,1076 unet=unet,1077 scheduler=scheduler,1078 tokenizer=tokenizer,1079 text_encoder=text_encoder,1080 feature_extractor=feature_extractor,1081 image_encoder=image_encoder,1082 )1083 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)1084 self.register_to_config(requires_safety_checker=requires_safety_checker)1085 1086 def enable_vae_slicing(self):1087 r"""1088 Enable sliced VAE decoding.1089 1090 When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several1091 steps. This is useful to save some memory and allow larger batch sizes.1092 """1093 self.vae.enable_slicing()1094 1095 def disable_vae_slicing(self):1096 r"""1097 Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to1098 computing decoding in one step.1099 """1100 self.vae.disable_slicing()1101 1102 def enable_vae_tiling(self):1103 r"""1104 Enable tiled VAE decoding.1105 1106 When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in1107 several steps. This is useful to save a large amount of memory and to allow the processing of larger images.1108 """1109 self.vae.enable_tiling()1110 1111 def disable_vae_tiling(self):1112 r"""1113 Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to1114 computing decoding in one step.1115 """1116 self.vae.disable_tiling()1117 1118 def enable_sequential_cpu_offload(self, gpu_id=0):1119 r"""1120 Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,1121 text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a1122 `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.1123 Note that offloading happens on a submodule basis. Memory savings are higher than with1124 `enable_model_cpu_offload`, but performance is lower.1125 """1126 if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):1127 from accelerate import cpu_offload1128 else:1129 raise ImportError(1130 "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"1131 )1132 1133 device = torch.device(f"cuda:{gpu_id}")1134 1135 if self.device.type != "cpu":1136 self.to("cpu", silence_dtype_warnings=True)1137 torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)1138 1139 for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:1140 cpu_offload(cpu_offloaded_model, device)1141 1142 def enable_model_cpu_offload(self, gpu_id=0):1143 r"""1144 Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared1145 to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`1146 method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with1147 `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.1148 """1149 if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):1150 from accelerate import cpu_offload_with_hook1151 else:1152 raise ImportError(1153 "`enable_model_offload` requires `accelerate v0.17.0` or higher."1154 )1155 1156 device = torch.device(f"cuda:{gpu_id}")1157 1158 if self.device.type != "cpu":1159 self.to("cpu", silence_dtype_warnings=True)1160 torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)1161 1162 hook = None1163 for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:1164 _, hook = cpu_offload_with_hook(1165 cpu_offloaded_model, device, prev_module_hook=hook1166 )1167 1168 # We'll offload the last model manually.1169 self.final_offload_hook = hook1170 1171 @property1172 def _execution_device(self):1173 r"""1174 Returns the device on which the pipeline's models will be executed. After calling1175 `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module1176 hooks.1177 """1178 if not hasattr(self.unet, "_hf_hook"):1179 return self.device1180 for module in self.unet.modules():1181 if (1182 hasattr(module, "_hf_hook")1183 and hasattr(module._hf_hook, "execution_device")1184 and module._hf_hook.execution_device is not None1185 ):1186 return torch.device(module._hf_hook.execution_device)1187 return self.device1188 1189 def _encode_prompt(1190 self,1191 prompt,1192 device,1193 num_images_per_prompt,1194 do_classifier_free_guidance: bool,1195 negative_prompt=None,1196 ):1197 r"""1198 Encodes the prompt into text encoder hidden states.1199 1200 Args: