(self, x, timesteps, context=None, y=None, features_adapter=None, is_imgbatch=False, T=None, **kwargs)
| 69 | return x |
| 70 | |
| 71 | def selfattn_forward_unet(self, x, timesteps, context=None, y=None, features_adapter=None, is_imgbatch=False, T=None, **kwargs): |
| 72 | b,_,t,_,_ = x.shape |
| 73 | |
| 74 | t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) |
| 75 | emb = self.time_embed(t_emb) |
| 76 | if self.micro_condition and y is not None: |
| 77 | micro_emb = timestep_embedding(y, self.model_channels, repeat_only=False) |
| 78 | emb = emb + self.micro_embed(micro_emb) |
| 79 | |
| 80 | |
| 81 | |
| 82 | # pose_emb = pose_emb.reshape(-1, pose_emb.shape[-1]) |
| 83 | ## repeat t times for context [(b t) 77 768] & time embedding |
| 84 | if not is_imgbatch: |
| 85 | context = context.repeat_interleave(repeats=t, dim=0) |
| 86 | |
| 87 | if 'pose_emb' in kwargs: |
| 88 | pose_emb = kwargs.pop('pose_emb') |
| 89 | context = { 'context': context, 'pose_emb': pose_emb } |
| 90 | |
| 91 | emb = emb.repeat_interleave(repeats=t, dim=0) |
| 92 | |
| 93 | ## always in shape (b t) c h w, except for temporal layer |
| 94 | x = rearrange(x, 'b c t h w -> (b t) c h w') |
| 95 | if features_adapter is not None: |
| 96 | features_adapter = [rearrange(feature, 'b c t h w -> (b t) c h w') for feature in features_adapter] |
| 97 | |
| 98 | h = x.type(self.dtype) |
| 99 | adapter_idx = 0 |
| 100 | hs = [] |
| 101 | for id, module in enumerate(self.input_blocks): |
| 102 | h = module(h, emb, context=context, batch_size=b,is_imgbatch=is_imgbatch) |
| 103 | if id ==0 and self.addition_attention: |
| 104 | h = self.init_attn(h, emb, context=context, batch_size=b,is_imgbatch=is_imgbatch) |
| 105 | ## plug-in adapter features |
| 106 | if ((id+1)%3 == 0) and features_adapter is not None: |
| 107 | # if adapter_idx == 0 or adapter_idx == 1 or adapter_idx == 2: |
| 108 | h = h + features_adapter[adapter_idx] |
| 109 | adapter_idx += 1 |
| 110 | hs.append(h) |
| 111 | if features_adapter is not None: |
| 112 | assert len(features_adapter)==adapter_idx, 'Wrong features_adapter' |
| 113 | |
| 114 | h = self.middle_block(h, emb, context=context, batch_size=b, is_imgbatch=is_imgbatch) |
| 115 | for module in self.output_blocks: |
| 116 | h = torch.cat([h, hs.pop()], dim=1) |
| 117 | h = module(h, emb, context=context, batch_size=b, is_imgbatch=is_imgbatch) |
| 118 | h = h.type(x.dtype) |
| 119 | y = self.out(h) |
| 120 | |
| 121 | # reshape back to (b c t h w) |
| 122 | y = rearrange(y, '(b t) c h w -> b c t h w', b=b) |
| 123 | return y |
| 124 | |
| 125 | def spatial_forward_BasicTransformerBlock(self, x, context=None, mask=None): |
| 126 | if isinstance(context, dict): |
nothing calls this directly
no test coverage detected