Method predict_noise

stg.py:818–884

(
        self,
        x: torch.Tensor,
        timestep: torch.Tensor,
        model_options: dict = {},
        seed=None,
    )

        self.inner_set_conds({"positive": positive, "negative": negative})

    def predict_noise(
        self,
        x: torch.Tensor,
        timestep: torch.Tensor,
        model_options: dict = {},
        seed=None,
    ):
        # CFG zero init - skipping steps with timestep bigger than given threshold.

        # in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
        # but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
        positive_cond = self.conds.get("positive", None)
        negative_cond = self.conds.get("negative", None)

        noise_pred_pos = comfy.samplers.calc_cond_batch(
            self.inner_model,
            [positive_cond],
            x,
            timestep,
            model_options,
        )[0]

        if (
            self.previous_timestep is not None
            and timestep.item() > self.previous_timestep
        ):
            print("Resetting momentum buffer")
            self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)

        noise_pred_neg = 0
        if not math.isclose(self.cfg_scale, 1.0):
            noise_pred_neg = comfy.samplers.calc_cond_batch(
                self.inner_model,
                [negative_cond],
                x,
                timestep,
                model_options,
            )[0]

        apg_result = apg(
            noise_pred_pos,
            noise_pred_neg,
            self.cfg_scale,
            self.eta,
            self.norm_threshold,
        )

        # normally this would be done in cfg_function, but we skipped
        # that for efficiency: we can compute the noise predictions in
        # a single call to calc_cond_batch() (rather than two)
        # so we replicate the hook here
        for fn in model_options.get("sampler_post_cfg_function", []):
            args = {
                "denoised": apg_result,
                "cond": positive_cond,
                "uncond": negative_cond,
                "model": self.inner_model,
                "uncond_denoised": noise_pred_neg,