fix cfg_channel (#217)

hpcaitech · Mar 25, 2024 · a37a189 · a37a189
1 parent b7cbae2
commit a37a189
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/opensora/schedulers/iddpm/__init__.py b/opensora/schedulers/iddpm/__init__.py
@@ -22,6 +22,7 @@ def __init__(
  rescale_learned_sigmas=False,
  diffusion_steps=1000,
  cfg_scale=4.0,
+ cfg_channel=None,
  ):
  betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
  if use_kl:
@@ -49,6 +50,7 @@ def __init__(
  )
 
  self.cfg_scale = cfg_scale
+ self.cfg_channel = cfg_channel
 
  def sample(
  self,
@@ -68,7 +70,7 @@ def sample(
  if additional_args is not None:
  model_args.update(additional_args)
 
- forward = partial(forward_with_cfg, model, cfg_scale=self.cfg_scale)
+ forward = partial(forward_with_cfg, model, cfg_scale=self.cfg_scale, cfg_channel=self.cfg_channel)
  samples = self.p_sample_loop(
  forward,
  z.shape,