(
self,
hidden_states,
controlnet_conditioning,
timestep, prompt_emb, pooled_prompt_emb, guidance, text_ids, image_ids=None,
processor_id=None,
tiled=False, tile_size=128, tile_stride=64,
**kwargs
)
| 57 | |
| 58 | |
| 59 | def forward( |
| 60 | self, |
| 61 | hidden_states, |
| 62 | controlnet_conditioning, |
| 63 | timestep, prompt_emb, pooled_prompt_emb, guidance, text_ids, image_ids=None, |
| 64 | processor_id=None, |
| 65 | tiled=False, tile_size=128, tile_stride=64, |
| 66 | **kwargs |
| 67 | ): |
| 68 | if image_ids is None: |
| 69 | image_ids = self.prepare_image_ids(hidden_states) |
| 70 | |
| 71 | conditioning = self.time_embedder(timestep, hidden_states.dtype) + self.pooled_text_embedder(pooled_prompt_emb) |
| 72 | if self.guidance_embedder is not None: |
| 73 | guidance = guidance * 1000 |
| 74 | conditioning = conditioning + self.guidance_embedder(guidance, hidden_states.dtype) |
| 75 | prompt_emb = self.context_embedder(prompt_emb) |
| 76 | if self.controlnet_mode_embedder is not None: # Different from FluxDiT |
| 77 | processor_id = torch.tensor([self.mode_dict[processor_id]], dtype=torch.int) |
| 78 | processor_id = repeat(processor_id, "D -> B D", B=1).to(text_ids.device) |
| 79 | prompt_emb = torch.concat([self.controlnet_mode_embedder(processor_id), prompt_emb], dim=1) |
| 80 | text_ids = torch.cat([text_ids[:, :1], text_ids], dim=1) |
| 81 | image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1)) |
| 82 | |
| 83 | hidden_states = self.patchify(hidden_states) |
| 84 | hidden_states = self.x_embedder(hidden_states) |
| 85 | controlnet_conditioning = self.patchify(controlnet_conditioning) # Different from FluxDiT |
| 86 | hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_conditioning) # Different from FluxDiT |
| 87 | |
| 88 | controlnet_res_stack = [] |
| 89 | for block, controlnet_block in zip(self.blocks, self.controlnet_blocks): |
| 90 | hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb) |
| 91 | controlnet_res_stack.append(controlnet_block(hidden_states)) |
| 92 | |
| 93 | controlnet_single_res_stack = [] |
| 94 | hidden_states = torch.cat([prompt_emb, hidden_states], dim=1) |
| 95 | for block, controlnet_block in zip(self.single_blocks, self.controlnet_single_blocks): |
| 96 | hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb) |
| 97 | controlnet_single_res_stack.append(controlnet_block(hidden_states[:, prompt_emb.shape[1]:])) |
| 98 | |
| 99 | controlnet_res_stack = self.align_res_stack_to_original_blocks(controlnet_res_stack, 19, hidden_states[:, prompt_emb.shape[1]:]) |
| 100 | controlnet_single_res_stack = self.align_res_stack_to_original_blocks(controlnet_single_res_stack, 38, hidden_states[:, prompt_emb.shape[1]:]) |
| 101 | |
| 102 | return controlnet_res_stack, controlnet_single_res_stack |
| 103 | |
| 104 | |
| 105 | @staticmethod |
nothing calls this directly
no test coverage detected