(self, x, pos, global_cond, context=None, context_pos=None)
| 107 | |
| 108 | |
| 109 | def forward(self, x, pos, global_cond, context=None, context_pos=None): |
| 110 | b, c, h, w = x.shape |
| 111 | x_in = x |
| 112 | x = rearrange(x, 'b c h w -> b h w c') |
| 113 | context = rearrange(context, 'b c h w -> b h w c') |
| 114 | x = self.x_in_norm(x, global_cond) |
| 115 | |
| 116 | if self.do_self_attention: |
| 117 | #x to qkv |
| 118 | x_qkv = self.x_qkv_proj(x) |
| 119 | pos = rearrange(pos, "... h w e -> ... (h w) e").to(x_qkv.dtype) |
| 120 | x_theta = self.x_pos_emb(pos) |
| 121 | if use_flash_2(x_qkv): |
| 122 | x_qkv = rearrange(x_qkv, "n h w (t nh e) -> n (h w) t nh e", t=3, e=self.d_head) |
| 123 | x_qkv = scale_for_cosine_sim_qkv(x_qkv, self.x_scale, 1e-6) |
| 124 | x_theta = torch.stack((x_theta, x_theta, torch.zeros_like(x_theta)), dim=-3) |
| 125 | x_qkv = apply_rotary_emb_(x_qkv, x_theta) |
| 126 | x_q, x_k, x_v = x_qkv.chunk(3,dim=-3) |
| 127 | else: |
| 128 | print("we couldnt run flash2, maybe it's not installed or the input si not bfloat16") |
| 129 | exit(1) |
| 130 | else: |
| 131 | #x to q |
| 132 | x_q = self.x_q_proj(x) |
| 133 | pos = rearrange(pos, "... h w e -> ... (h w) e").to(x_q.dtype) |
| 134 | x_theta = self.x_pos_emb(pos) |
| 135 | if use_flash_2(x_q): |
| 136 | x_q = rearrange(x_q, "n h w (nh e) -> n (h w) nh e", e=self.d_head) |
| 137 | x_q = scale_for_cosine_sim_single(x_q, self.x_scale[:, None], 1e-6) |
| 138 | x_q=x_q.unsqueeze(2) #n (h w) 1 nh e |
| 139 | x_theta=x_theta.unsqueeze(1) |
| 140 | x_q = apply_rotary_emb_(x_q, x_theta) |
| 141 | else: |
| 142 | print("we couldnt run flash2, maybe it's not installed or the input si not bfloat16") |
| 143 | exit(1) |
| 144 | |
| 145 | |
| 146 | #context to kv |
| 147 | cond_kv = self.cond_kv_proj(context) |
| 148 | # print("cond_kv init",cond_kv.shape) |
| 149 | context_pos = rearrange(context_pos, "... h w e -> ... (h w) e").to(cond_kv.dtype) |
| 150 | cond_theta = self.cond_pos_emb(context_pos) |
| 151 | if use_flash_2(cond_kv): |
| 152 | cond_kv = rearrange(cond_kv, "n h w (t nh e) -> n (h w) t nh e", t=2, e=self.d_head) |
| 153 | cond_k, cond_v = cond_kv.unbind(2) # makes each n (h w) nh e |
| 154 | cond_k = scale_for_cosine_sim_single(cond_k, self.cond_scale[:, None], 1e-6) |
| 155 | cond_k=cond_k.unsqueeze(2) #n (h w) 1 nh e |
| 156 | cond_theta=cond_theta.unsqueeze(1) |
| 157 | cond_k = apply_rotary_emb_(cond_k, cond_theta) |
| 158 | cond_k=cond_k.squeeze(2) |
| 159 | else: |
| 160 | print("we couldnt run flash2, maybe it's not installed or the input si not bfloat16") |
| 161 | exit(1) |
| 162 | |
| 163 | #doing self attention by concating K and V between X and cond |
| 164 | if self.do_self_attention: |
| 165 | k = torch.cat([x_k, cond_k.unsqueeze(2)], dim=1) |
| 166 | v = torch.cat([x_v, cond_v.unsqueeze(2)], dim=1) |
nothing calls this directly
no test coverage detected