(
self,
sample, timestep, encoder_hidden_states,
conditioning, processor_id, add_time_id, add_text_embeds,
tiled=False, tile_size=64, tile_stride=32,
unet:SDXLUNet=None,
**kwargs
)
| 149 | |
| 150 | |
| 151 | def forward( |
| 152 | self, |
| 153 | sample, timestep, encoder_hidden_states, |
| 154 | conditioning, processor_id, add_time_id, add_text_embeds, |
| 155 | tiled=False, tile_size=64, tile_stride=32, |
| 156 | unet:SDXLUNet=None, |
| 157 | **kwargs |
| 158 | ): |
| 159 | task_id = self.task_id[processor_id] |
| 160 | |
| 161 | # 1. time |
| 162 | t_emb = self.time_proj(timestep).to(sample.dtype) |
| 163 | t_emb = self.time_embedding(t_emb) |
| 164 | |
| 165 | time_embeds = self.add_time_proj(add_time_id) |
| 166 | time_embeds = time_embeds.reshape((add_text_embeds.shape[0], -1)) |
| 167 | add_embeds = torch.concat([add_text_embeds, time_embeds], dim=-1) |
| 168 | add_embeds = add_embeds.to(sample.dtype) |
| 169 | if unet is not None and unet.is_kolors: |
| 170 | add_embeds = unet.add_time_embedding(add_embeds) |
| 171 | else: |
| 172 | add_embeds = self.add_time_embedding(add_embeds) |
| 173 | |
| 174 | control_type = torch.zeros((sample.shape[0], 8), dtype=sample.dtype, device=sample.device) |
| 175 | control_type[:, task_id] = 1 |
| 176 | control_embeds = self.control_type_proj(control_type.flatten()) |
| 177 | control_embeds = control_embeds.reshape((sample.shape[0], -1)) |
| 178 | control_embeds = control_embeds.to(sample.dtype) |
| 179 | control_embeds = self.control_type_embedding(control_embeds) |
| 180 | time_emb = t_emb + add_embeds + control_embeds |
| 181 | |
| 182 | # 2. pre-process |
| 183 | height, width = sample.shape[2], sample.shape[3] |
| 184 | hidden_states = self.conv_in(sample) |
| 185 | hidden_states = self.fuse_condition_to_input(hidden_states, task_id, conditioning) |
| 186 | text_emb = encoder_hidden_states |
| 187 | if unet is not None and unet.is_kolors: |
| 188 | text_emb = unet.text_intermediate_proj(text_emb) |
| 189 | res_stack = [hidden_states] |
| 190 | |
| 191 | # 3. blocks |
| 192 | for i, block in enumerate(self.blocks): |
| 193 | if tiled and not isinstance(block, PushBlock): |
| 194 | _, _, inter_height, _ = hidden_states.shape |
| 195 | resize_scale = inter_height / height |
| 196 | hidden_states = TileWorker().tiled_forward( |
| 197 | lambda x: block(x, time_emb, text_emb, res_stack)[0], |
| 198 | hidden_states, |
| 199 | int(tile_size * resize_scale), |
| 200 | int(tile_stride * resize_scale), |
| 201 | tile_device=hidden_states.device, |
| 202 | tile_dtype=hidden_states.dtype |
| 203 | ) |
| 204 | else: |
| 205 | hidden_states, _, _, _ = block(hidden_states, time_emb, text_emb, res_stack) |
| 206 | |
| 207 | # 4. ControlNet blocks |
| 208 | controlnet_res_stack = [block(res) for block, res in zip(self.controlnet_blocks, res_stack)] |
nothing calls this directly
no test coverage detected