| 132 | |
| 133 | |
| 134 | class InitialLayer(nn.Module): |
| 135 | def __init__(self, model): |
| 136 | super().__init__() |
| 137 | self.input_proj = model.input_proj |
| 138 | self.t_embedding = model.t_embedding |
| 139 | self.adaln_proj = model.adaln_proj |
| 140 | self.llm_cond_norm = model.llm_cond_norm |
| 141 | self.llm_cond_proj = model.llm_cond_proj |
| 142 | self.embed_image_indicator = model.embed_image_indicator |
| 143 | self.model = [model] |
| 144 | |
| 145 | def __getattr__(self, name): |
| 146 | return getattr(self.model[0], name) |
| 147 | |
| 148 | # Must NOT use autocast here or model output is so degraded it can't gen a coherent image. |
| 149 | @torch.compiler.disable |
| 150 | def forward(self, inputs): |
| 151 | for item in inputs: |
| 152 | if torch.is_floating_point(item): |
| 153 | item.requires_grad_(True) |
| 154 | x_chunk, timesteps, context_chunk, attn_mask_chunk = inputs |
| 155 | t_chunk = 1.0 - timesteps |
| 156 | bs, c, gh, gw = x_chunk.shape |
| 157 | |
| 158 | # This is only the conditional pathway |
| 159 | B = x_chunk.shape[0] |
| 160 | device = x_chunk.device |
| 161 | img_tokens = self._img_to_tokens(x_chunk) |
| 162 | L_img = img_tokens.shape[1] |
| 163 | L_text = context_chunk.shape[1] |
| 164 | L = L_text + L_img |
| 165 | latent_dim = img_tokens.shape[-1] |
| 166 | |
| 167 | x_full = torch.zeros(B, L, latent_dim, dtype=img_tokens.dtype, device=device) |
| 168 | x_full[:, L_text:] = img_tokens |
| 169 | |
| 170 | text_pos = torch.arange(L_text, device=device).view(-1, 1).expand(L_text, 3) |
| 171 | img_pos = self._image_position_ids(gh, gw, device) |
| 172 | position_ids = torch.cat([text_pos, img_pos], dim=0).unsqueeze(0).expand(B, L, 3) |
| 173 | |
| 174 | indicator = torch.empty(B, L, dtype=torch.long, device=device) |
| 175 | indicator[:, :L_text] = LLM_TOKEN_INDICATOR |
| 176 | indicator[:, L_text:] = OUTPUT_IMAGE_INDICATOR |
| 177 | |
| 178 | segment_ids = torch.ones(B, L, dtype=torch.long, device=device) |
| 179 | pad = (attn_mask_chunk == 0) |
| 180 | segment_ids[:, :L_text][pad] = SEQUENCE_PADDING_INDICATOR |
| 181 | indicator[:, :L_text][pad] = 0 |
| 182 | # Block-diagonal mask from segment ids: (B, 1, L, L), True = attend. |
| 183 | attn_mask = (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1) |
| 184 | |
| 185 | # backbone |
| 186 | llm_features = context_chunk |
| 187 | x = x_full |
| 188 | t = t_chunk |
| 189 | |
| 190 | indicator = indicator.to(torch.long) |
| 191 | output_image_mask = (indicator == OUTPUT_IMAGE_INDICATOR).to(x.dtype).unsqueeze(-1) |