# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_guidance.ipynb.

# %% auto 0
__all__ = ['DynamicCFG']

# %% ../nbs/01_guidance.ipynb 2
from .kdiff import *
from .schedules import name2schedule
from .normalizers import name2norm

# %% ../nbs/01_guidance.ipynb 3
class DynamicCFG:
    def __init__(self, norm_name, schedule_name):
        self.normalzier = name2norm[norm_name]()
        self.scheduler = name2schedule[schedule_name]()

    def guide(self, uncond, cond, ts):
        """Applies dynamic Classifier-free Guidance.
        """
        # set the conditional and unconditional vectors
        self.normalzier.set_latents(u=uncond, t=cond)
        # compute the guidate update vector: (cond - uncond)
        self.normalzier.compute_update()

        # apply optional pre-processing
        if self.normalzier.has_preproc:
            self.normalzier.pre_proc()

        # get the current, scheduled guidance value 
        guide_scale = self.scheduler.value_at(ts)
        # run classifier-free guidance 
        pred = self.normalzier.u + (guide_scale * self.normalzier.diff)
        self.normalzier.set_pred(pred)

        # apply optional post-processing
        if self.normalzier.has_postproc:
            self.normalzier.post_proc()

        return self.normalzier.get_pred()

    def update_sched_kwargs(self, other_kwargs):
        self.scheduler.update_sched_kwargs(other_kwargs)
    def set_timesteps(self, num_steps):
        self.scheduler.set_num_steps(num_steps)
        self.scheduler.set_guidance_schedule()
