| 318 | |
| 319 | |
| 320 | class STGGuiderAdvanced(comfy.samplers.CFGGuider): |
| 321 | def __init__( |
| 322 | self, |
| 323 | model: ModelPatcher, |
| 324 | sigma_list, |
| 325 | cfg_list, |
| 326 | stg_scale_list, |
| 327 | stg_rescale_list, |
| 328 | stg_layers_indices_list, |
| 329 | skip_steps_sigma_threshold, |
| 330 | cfg_star_rescale, |
| 331 | apply_apg, |
| 332 | apg_cfg_scale, |
| 333 | eta, |
| 334 | norm_threshold, |
| 335 | ): |
| 336 | model = model.clone() |
| 337 | super().__init__(model) |
| 338 | |
| 339 | self.stg_flag = STGFlag( |
| 340 | do_skip=False, |
| 341 | skip_layers=model.model_options["transformer_options"].get( |
| 342 | "skip_block_list" |
| 343 | ), |
| 344 | ) |
| 345 | |
| 346 | self.sigma_list = sigma_list |
| 347 | self.cfg_list = cfg_list |
| 348 | self.stg_scale_list = stg_scale_list |
| 349 | self.stg_rescale_list = stg_rescale_list |
| 350 | self.stg_layers_indices_list = stg_layers_indices_list |
| 351 | self.skip_steps_sigma_threshold = skip_steps_sigma_threshold |
| 352 | self.cfg_star_rescale = cfg_star_rescale |
| 353 | self.apply_apg = apply_apg |
| 354 | self.apg_cfg_scale = apg_cfg_scale |
| 355 | self.eta = eta |
| 356 | self.norm_threshold = norm_threshold |
| 357 | STGGuider.patch_model(model, self.stg_flag) |
| 358 | |
| 359 | def sigma_to_params_mapping(self, sigma): |
| 360 | # Find the closest higher sigma value and return corresponding cfg |
| 361 | higher_sigmas = [s for s in self.sigma_list if s >= sigma] |
| 362 | if not higher_sigmas: |
| 363 | closest_idx = -1 # Return last cfg if no higher sigma exists |
| 364 | else: |
| 365 | closest_higher = min(higher_sigmas) |
| 366 | closest_idx = self.sigma_list.index(closest_higher) |
| 367 | return ( |
| 368 | self.cfg_list[closest_idx], |
| 369 | self.stg_scale_list[closest_idx], |
| 370 | self.stg_rescale_list[closest_idx], |
| 371 | self.stg_layers_indices_list[closest_idx], |
| 372 | ) |
| 373 | |
| 374 | def set_conds(self, positive, negative): |
| 375 | self.inner_set_conds({"positive": positive, "negative": negative}) |
| 376 | |
| 377 | def predict_noise( |