Method p_losses

src/model.py:347–545

(self, x_start, cond, t, origin_prompt=None, image_input=None, mask=None, noise=None, mask_background_image=None)

        return opt

    def p_losses(self, x_start, cond, t, origin_prompt=None, image_input=None, mask=None, noise=None, mask_background_image=None):

        noise = default(noise, lambda: torch.randn_like(x_start))
        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
        model_output = self.apply_model(x_noisy, t, cond)

        loss_dict = {}
        prefix = 'train' if self.training else 'val'

        if self.parameterization == "x0":
            target = x_start
        elif self.parameterization == "eps":
            target = noise
        else:
            raise NotImplementedError()

        loss_simple = self.get_loss(model_output, target, mean=False)
        
        if mask is not None:
            loss_simple = (loss_simple*mask).sum([1, 2, 3])/mask.sum([1, 2, 3])
        else:
            loss_simple = loss_simple.mean([1, 2, 3])
        loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})

        logvar_t = (self.logvar.to(self.device))[t]
        loss = loss_simple / torch.exp(logvar_t) + logvar_t
        # loss = loss_simple / torch.exp(self.logvar) + self.logvar
        if self.learn_logvar:
            loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
            loss_dict.update({'logvar': self.logvar.data.mean()})

        loss = self.l_simple_weight * loss.mean()
        
        if self.controller is not None and ('<new1>' in origin_prompt[0] or '<new1>' in origin_prompt[1]):

            batch_encoding = self.cond_stage_model.tokenizer(origin_prompt, truncation=True, max_length=self.cond_stage_model.max_length, return_length=True,
                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
            tokens = batch_encoding["input_ids"].to(self.device)
            len_add_token = len(self.cond_stage_model.modifier_token_id)
            
            def get_token_indices(tokens, token_id):
                indices = tokens == token_id
                indices = (indices*1).unsqueeze(-1)
                index_token = (indices == 1).nonzero()
                index_token = index_token.tolist()
                return index_token
            
            index_token = get_token_indices(tokens, self.cond_stage_model.modifier_token_id[-len_add_token])
            index_token_new2 = get_token_indices(tokens, self.cond_stage_model.modifier_token_id[-1])

            if self.new3_subject:
                index_token_new3 = get_token_indices(tokens, self.cond_stage_model.modifier_token_id[-2])

            del batch_encoding
            
            out_new1, out_foreground, out_new2, out_subject2 = {0:[], 1:[]}, {0:[], 1:[]}, {0:[], 1:[]}, {0:[], 1:[]}
            if self.new3_subject:
                out_new3, out_subject3 = {0:[], 1:[]}, {0:[], 1:[]}