File size: 744 Bytes
491eded |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
from typing import *
import torch
class GuidanceIntervalSamplerMixin:
"""
A mixin class for samplers that apply classifier-free guidance with interval.
"""
def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs):
if cfg_interval[0] <= t <= cfg_interval[1]:
pred = super()._inference_model(model, x_t, t, cond, **kwargs)
if 'masks' in kwargs:
kwargs['masks'] = torch.zeros_like(kwargs['masks'])
neg_pred = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
return (1 + cfg_strength) * pred - cfg_strength * neg_pred
else:
return super()._inference_model(model, x_t, t, cond, **kwargs)
|