| from diffusers.configuration_utils import register_to_config |
| from diffusers.guiders.guider_utils import BaseGuidance, GuiderOutput |
|
|
|
|
| class Pixal3DGuidance(BaseGuidance): |
| """Classifier-free guidance with Pixal3D's sigma interval and clean-sample (x0) rescaling. |
| |
| Guidance is active while the current sigma lies inside `guidance_interval`; outside it, or with |
| `guidance_scale=1`, only the conditional prediction is computed. The rescale matches the standard |
| deviation of the guided x0 estimate to the conditional one, which needs the current `latents`. |
| """ |
|
|
| _input_predictions = ["pred_cond", "pred_uncond"] |
|
|
| @register_to_config |
| def __init__( |
| self, |
| guidance_scale=7.5, |
| guidance_rescale=0.7, |
| guidance_interval=(0.6, 1.0), |
| sigma_min=1e-5, |
| enabled=True, |
| ): |
| super().__init__(enabled=enabled) |
|
|
| @property |
| def _scale(self): |
| sigma = self._timestep |
| low, high = self.config.guidance_interval |
| if self._enabled and sigma is not None and low <= sigma <= high: |
| return self.config.guidance_scale |
| return 1.0 |
|
|
| @property |
| def is_conditional(self): |
| return self._count_prepared == 1 |
|
|
| @property |
| def num_conditions(self): |
| return 1 if self._scale in (0, 1) else 2 |
|
|
| def _predictions(self): |
| |
| if self._scale == 0: |
| return ["pred_uncond"] |
| return self._input_predictions[: self.num_conditions] |
|
|
| def prepare_inputs(self, data): |
| return [self._prepare_batch(data, i, name) for i, name in self._identified_predictions()] |
|
|
| def prepare_inputs_from_block_state(self, data, input_fields): |
| return [ |
| self._prepare_batch_from_block_state(input_fields, data, i, name) |
| for i, name in self._identified_predictions() |
| ] |
|
|
| def _identified_predictions(self): |
| return [(self._input_predictions.index(name), name) for name in self._predictions()] |
|
|
| def __call__(self, data): |
| predictions = {getattr(d, self._identifier_key): d.noise_pred for d in data} |
| return self.forward(**predictions, latents=data[0].latents) |
|
|
| def forward(self, pred_cond=None, pred_uncond=None, latents=None): |
| if pred_cond is None: |
| return GuiderOutput(pred=pred_uncond) |
| if pred_uncond is None: |
| return GuiderOutput(pred=pred_cond) |
| latents = latents.to(pred_cond.device) |
| pred_uncond = pred_uncond.to(pred_cond.device) |
| scale = self._scale |
| result = scale * pred_cond + (1 - scale) * pred_uncond |
| if self.config.guidance_rescale > 0: |
| sigma_min = self.config.sigma_min |
| sigma = sigma_min + (1 - sigma_min) * self._timestep |
| clean_pos = (1 - sigma_min) * latents - sigma * pred_cond |
| clean_cfg = (1 - sigma_min) * latents - sigma * result |
| dims = list(range(1, latents.ndim)) |
| rescaled = clean_cfg * (clean_pos.std(dim=dims, keepdim=True) / clean_cfg.std(dim=dims, keepdim=True)) |
| clean = self.config.guidance_rescale * rescaled + (1 - self.config.guidance_rescale) * clean_cfg |
| result = ((1 - sigma_min) * latents - clean) / sigma |
| return GuiderOutput(pred=result, pred_cond=pred_cond, pred_uncond=pred_uncond) |
|
|
