(self, x, batch)
| 159 | self.first_stage_model = model |
| 160 | |
| 161 | def forward(self, x, batch): #CJH Here is to pass the loss function |
| 162 | # breakpoint() |
| 163 | loss = self.loss_fn(self.model, self.denoiser, self.conditioner, self.clip_img_encoder, x, batch) |
| 164 | loss_mean = loss.mean() |
| 165 | loss_dict = {"loss": loss_mean} |
| 166 | return loss_mean, loss_dict |
| 167 | |
| 168 | def add_noise_to_first_frame(self, image): |
| 169 | sigma = torch.normal(mean=-3.0, std=0.5, size=(image.shape[0],)).to(self.device) |
nothing calls this directly
no outgoing calls
no test coverage detected