(
self,
sample, timestep, encoder_hidden_states, conditioning,
tiled=False, tile_size=64, tile_stride=32,
**kwargs
)
| 94 | self.global_pool = global_pool |
| 95 | |
| 96 | def forward( |
| 97 | self, |
| 98 | sample, timestep, encoder_hidden_states, conditioning, |
| 99 | tiled=False, tile_size=64, tile_stride=32, |
| 100 | **kwargs |
| 101 | ): |
| 102 | # 1. time |
| 103 | time_emb = self.time_proj(timestep).to(sample.dtype) |
| 104 | time_emb = self.time_embedding(time_emb) |
| 105 | time_emb = time_emb.repeat(sample.shape[0], 1) |
| 106 | |
| 107 | # 2. pre-process |
| 108 | height, width = sample.shape[2], sample.shape[3] |
| 109 | hidden_states = self.conv_in(sample) + self.controlnet_conv_in(conditioning) |
| 110 | text_emb = encoder_hidden_states |
| 111 | res_stack = [hidden_states] |
| 112 | |
| 113 | # 3. blocks |
| 114 | for i, block in enumerate(self.blocks): |
| 115 | if tiled and not isinstance(block, PushBlock): |
| 116 | _, _, inter_height, _ = hidden_states.shape |
| 117 | resize_scale = inter_height / height |
| 118 | hidden_states = TileWorker().tiled_forward( |
| 119 | lambda x: block(x, time_emb, text_emb, res_stack)[0], |
| 120 | hidden_states, |
| 121 | int(tile_size * resize_scale), |
| 122 | int(tile_stride * resize_scale), |
| 123 | tile_device=hidden_states.device, |
| 124 | tile_dtype=hidden_states.dtype |
| 125 | ) |
| 126 | else: |
| 127 | hidden_states, _, _, _ = block(hidden_states, time_emb, text_emb, res_stack) |
| 128 | |
| 129 | # 4. ControlNet blocks |
| 130 | controlnet_res_stack = [block(res) for block, res in zip(self.controlnet_blocks, res_stack)] |
| 131 | |
| 132 | # pool |
| 133 | if self.global_pool: |
| 134 | controlnet_res_stack = [res.mean(dim=(2, 3), keepdim=True) for res in controlnet_res_stack] |
| 135 | |
| 136 | return controlnet_res_stack |
| 137 | |
| 138 | @staticmethod |
| 139 | def state_dict_converter(): |
nothing calls this directly
no test coverage detected