(self, batch: Dict)
| 173 | return image |
| 174 | |
| 175 | def shared_step(self, batch: Dict) -> Any: |
| 176 | x = self.get_input(batch) |
| 177 | if self.lr_scale is not None: |
| 178 | lr_x = F.interpolate(x, scale_factor=1 / self.lr_scale, mode="bilinear", align_corners=False) |
| 179 | lr_x = F.interpolate(lr_x, scale_factor=self.lr_scale, mode="bilinear", align_corners=False) |
| 180 | lr_z = self.encode_first_stage(lr_x, batch) |
| 181 | batch["lr_input"] = lr_z |
| 182 | |
| 183 | x = x.permute(0, 2, 1, 3, 4).contiguous() |
| 184 | if self.noised_image_input: |
| 185 | image = x[:, :, 0:1] |
| 186 | image = self.add_noise_to_first_frame(image) |
| 187 | image = self.encode_first_stage(image, batch) |
| 188 | |
| 189 | x = self.encode_first_stage(x, batch) # x [2, 3, 49, 480, 720] batch: fps [8,8] |
| 190 | x = x.permute(0, 2, 1, 3, 4).contiguous() |
| 191 | if self.noised_image_input: |
| 192 | image = image.permute(0, 2, 1, 3, 4).contiguous() |
| 193 | if self.noised_image_all_concat: |
| 194 | image = image.repeat(1, x.shape[1], 1, 1, 1) |
| 195 | else: |
| 196 | image = torch.concat([image, torch.zeros_like(x[:, 1:])], dim=1) |
| 197 | if random.random() < self.noised_image_dropout: |
| 198 | image = torch.zeros_like(image) |
| 199 | batch["concat_images"] = image |
| 200 | |
| 201 | gc.collect() |
| 202 | torch.cuda.empty_cache() |
| 203 | ''' |
| 204 | Here is to training |
| 205 | ''' |
| 206 | # breakpoint() |
| 207 | loss, loss_dict = self(x, batch) |
| 208 | return loss, loss_dict |
| 209 | |
| 210 | def get_input(self, batch): |
| 211 | return batch[self.input_key].to(self.dtype) |
nothing calls this directly
no test coverage detected