Method predict_noise

stg.py:227–317

(
        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,
    ):
        # 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)

        if model_options.get("sigma_to_params_mapping", None) is not None:
            cfg_value, stg_scale, stg_layer_skip_layer_indices, stg_rescale = (
                model_options["sigma_to_params_mapping"](timestep)
            )
            self.stg_flag.skip_layers = stg_layer_skip_layer_indices
            self.patch_model(self.model_patcher, self.stg_flag)

        else:
            cfg_value = self.cfg
            stg_scale = self.stg_scale
            stg_rescale = self.rescale_scale

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

        noise_pred_neg = 0
        noise_pred_perturbed = 0

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

        if not math.isclose(stg_scale, 0.0):
            try:
                model_options["transformer_options"]["ptb_index"] = 0
                self.stg_flag.do_skip = True
                noise_pred_perturbed = comfy.samplers.calc_cond_batch(
                    self.inner_model,
                    [positive_cond],
                    x,
                    timestep,
                    model_options,
                )[0]
            finally:
                self.stg_flag.do_skip = False