| 187 | |
| 188 | |
| 189 | class STGGuider(comfy.samplers.CFGGuider): |
| 190 | def __init__( |
| 191 | self, model: ModelPatcher, cfg, stg_scale, rescale_scale: float = None |
| 192 | ): |
| 193 | model = model.clone() |
| 194 | super().__init__(model) |
| 195 | |
| 196 | self.stg_flag = STGFlag( |
| 197 | do_skip=False, |
| 198 | skip_layers=model.model_options["transformer_options"]["skip_block_list"], |
| 199 | ) |
| 200 | |
| 201 | self.patch_model(model, self.stg_flag) |
| 202 | |
| 203 | self.cfg = cfg |
| 204 | self.stg_scale = stg_scale |
| 205 | self.rescale_scale = rescale_scale |
| 206 | |
| 207 | @classmethod |
| 208 | def patch_model(cls, model: ModelPatcher, stg_flag: STGFlag): |
| 209 | transformer_blocks = cls.get_transformer_blocks(model) |
| 210 | |
| 211 | for i, block in enumerate(transformer_blocks): |
| 212 | model.set_model_patch_replace( |
| 213 | STGBlockWrapper(block, stg_flag, i), "dit", "double_block", i |
| 214 | ) |
| 215 | |
| 216 | @staticmethod |
| 217 | def get_transformer_blocks(model: ModelPatcher): |
| 218 | diffusion_model = model.get_model_object("diffusion_model") |
| 219 | key = "diffusion_model.transformer_blocks" |
| 220 | if diffusion_model.__class__.__name__ == "LTXVTransformer3D": |
| 221 | key = "diffusion_model.transformer.transformer_blocks" |
| 222 | return model.get_model_object(key) |
| 223 | |
| 224 | def set_conds(self, positive, negative): |
| 225 | self.inner_set_conds({"positive": positive, "negative": negative}) |
| 226 | |
| 227 | def predict_noise( |
| 228 | self, |
| 229 | x: torch.Tensor, |
| 230 | timestep: torch.Tensor, |
| 231 | model_options: dict = {}, |
| 232 | seed=None, |
| 233 | ): |
| 234 | # in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg |
| 235 | # but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly |
| 236 | |
| 237 | positive_cond = self.conds.get("positive", None) |
| 238 | negative_cond = self.conds.get("negative", None) |
| 239 | |
| 240 | if model_options.get("sigma_to_params_mapping", None) is not None: |
| 241 | cfg_value, stg_scale, stg_layer_skip_layer_indices, stg_rescale = ( |
| 242 | model_options["sigma_to_params_mapping"](timestep) |
| 243 | ) |
| 244 | self.stg_flag.skip_layers = stg_layer_skip_layer_indices |
| 245 | self.patch_model(self.model_patcher, self.stg_flag) |
| 246 | |