Spaces:

diffusers-modular
/

flux3-image-klein-bbox

Running on Zero

flux3-image-klein-bbox / diffusers /guiders /ltx2_guidance.py

FLUX.3 layout-control canvas (custom HTML) — initial

3db1cbc verified 22 days ago

8.93 kB

# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from typing import TYPE_CHECKING

import torch

from ..configuration_utils import register_to_config
from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg


if TYPE_CHECKING:
    from ..modular_pipelines.modular_pipeline import BlockState


class LTX2Guidance(BaseGuidance):
    """
    Single-modality guider for LTX-2.X. Combines up to three terms via the delta formulation:
      - classifier-free guidance (CFG),
      - spatio-temporal guidance (STG), an extra pass with a set of transformer blocks perturbed,
      - modality-isolation guidance, an extra pass with A2V/V2A cross-attention disabled.
    One instance drives one modality, so a multimodal checkpoint carries several: LTX-2.X pairs a video `guider`
    with an `audio_guider`, each holding its own scales. The final guidance estimate is computed in whatever space
    the caller feeds it — the LTX-2 denoiser converts velocity->x0 before calling this and back afterwards, so
    `forward` operates on x0 predictions.
    Unlike `SkipLayerGuidance`, which applies its perturbations with hooks, `LTX2Guidance` expects its non-CFG
    perturbations to be configurable through model forward arguments (`spatio_temporal_guidance_blocks` and
    `isolate_modalities` in `LTX2VideoTransformer3DModel`). The calling denoiser block is responsible for setting
    those per pass identifier; see `LTX2LoopDenoiser` for a usage example.
    Args:
        guidance_scale (`float`, defaults to `1.0`):
            CFG scale for this modality. The CFG pass is skipped entirely at `1.0`.
        stg_scale (`float`, defaults to `0.0`):
            Spatio-temporal guidance scale for this modality. The STG pass is skipped entirely at `0.0`.
        modality_scale (`float`, defaults to `1.0`):
            Modality-isolation guidance scale for this modality. That pass is skipped entirely at `1.0`.
        guidance_rescale (`float`, defaults to `0.0`):
            Rescaling factor to prevent overexposure from high guidance scales. Based on [Common Diffusion Noise
            Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Range: 0.0 (no rescaling)
            to 1.0 (full rescaling).
        spatio_temporal_guidance_blocks (`list[int]`, *optional*):
            Transformer blocks to perturb on the STG pass. How the value is consumed is up to the denoiser block:
            LTX-2's perturbs whole blocks in a single forward that feeds every modality, so it reads this off the video
            guider alone — setting it on the audio guider has no effect, and that guider joins the same STG forward
            through its own `stg_scale`.
        start (`float`, defaults to `0.0`):
            Fraction of denoising steps (0.0-1.0) after which guidance starts.
        stop (`float`, defaults to `1.0`):
            Fraction of denoising steps (0.0-1.0) after which guidance stops.
        enabled (`bool`, defaults to `True`):
            Whether this guider applies guidance at all. When `False`, only the conditional pass runs.
    """

    # `pred_cond` is the base; the others are the guidance passes.
    _input_predictions = ["pred_cond", "pred_uncond", "pred_cond_stg", "pred_cond_modality"]

    @register_to_config
    def __init__(
        self,
        guidance_scale: float = 1.0,
        stg_scale: float = 0.0,
        modality_scale: float = 1.0,
        guidance_rescale: float = 0.0,
        spatio_temporal_guidance_blocks: list[int] | None = None,
        start: float = 0.0,
        stop: float = 1.0,
        enabled: bool = True,
    ):
        super().__init__(start, stop, enabled)
        self.guidance_scale = guidance_scale
        self.stg_scale = stg_scale
        self.modality_scale = modality_scale
        self.guidance_rescale = guidance_rescale
        self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks

    # --- per-term enabled checks (public: the denoiser ORs these across the two guiders) --------------------
    def _within_window(self) -> bool:
        if self._num_inference_steps is None:
            return True
        start_step = int(self._start * self._num_inference_steps)
        stop_step = int(self._stop * self._num_inference_steps)
        return start_step <= self._step < stop_step

    def is_cfg_enabled(self) -> bool:
        return self._enabled and self._within_window() and not math.isclose(self.guidance_scale, 1.0)

    def is_stg_enabled(self) -> bool:
        return self._enabled and self._within_window() and not math.isclose(self.stg_scale, 0.0)

    def is_modality_enabled(self) -> bool:
        return self._enabled and self._within_window() and not math.isclose(self.modality_scale, 1.0)

    def active_predictions(self) -> list[str]:
        preds = ["pred_cond"]
        if self.is_cfg_enabled():
            preds.append("pred_uncond")
        if self.is_stg_enabled():
            preds.append("pred_cond_stg")
        if self.is_modality_enabled():
            preds.append("pred_cond_modality")
        return preds

    @property
    def num_conditions(self) -> int:
        return len(self.active_predictions())

    @property
    def is_conditional(self) -> bool:
        # Not load-bearing here: the denoiser owns pass execution and carries the STG/modality flags on the
        # batches, so we never drive `prepare_models`/hooks off `_count_prepared` the way SkipLayerGuidance does.
        return True

    # Each pass reads its own slot ([cond, uncond, stg, modality]) of every `guider_input_fields` tuple, selected
    # by this per-identifier index. Distinct slots (rather than cond/uncond only) let STG/modality pull the same
    # conditioning as `pred_cond` while the denoiser overrides their per-pass model flags after preparation.
    _PREDICTION_INDEX = {"pred_cond": 0, "pred_uncond": 1, "pred_cond_stg": 2, "pred_cond_modality": 3}

    def prepare_inputs(self, guider_inputs: dict) -> list["BlockState"]:
        # One identifier-tagged batch per active pass. Every value in `guider_inputs` is a 4-tuple indexed by
        # `_PREDICTION_INDEX` ([cond, uncond, stg, modality]); each pass reads its own slot, so the per-pass model
        # flags (`spatio_temporal_guidance_blocks`, `isolate_modalities`) are carried exactly like the encoder
        # inputs. The denoiser fills each returned batch's `noise_pred` before calling the guider to combine.
        return [
            self._prepare_batch(guider_inputs, self._PREDICTION_INDEX[pred], pred)
            for pred in self.active_predictions()
        ]

    def prepare_inputs_from_block_state(self, data: "BlockState", input_fields: dict) -> list["BlockState"]:
        # One identifier-tagged batch per active pass. Each value in `input_fields` maps a transformer argument to a
        # 4-tuple of block-state attribute names indexed by `_PREDICTION_INDEX` ([cond, uncond, stg, modality]); the
        # base helper reads the pass's slot off `data`. The denoiser then sets the per-pass model flags and fills
        # each batch's `noise_pred` before calling the guider to combine.
        return [
            self._prepare_batch_from_block_state(input_fields, data, self._PREDICTION_INDEX[pred], pred)
            for pred in self.active_predictions()
        ]

    def forward(
        self,
        pred_cond: torch.Tensor,
        pred_uncond: torch.Tensor | None = None,
        pred_cond_stg: torch.Tensor | None = None,
        pred_cond_modality: torch.Tensor | None = None,
    ) -> GuiderOutput:
        # Delta formulation (matches the reference LTX-2 combine, applied in the x0 space the denoiser feeds):
        #   x0_cond + (gs-1)*(cond-uncond) + stg*(cond-stg) + (mod-1)*(cond-modality), then optional rescale.
        pred = pred_cond
        if self.is_cfg_enabled():
            pred = pred + (self.guidance_scale - 1.0) * (pred_cond - pred_uncond)
        if self.is_stg_enabled():
            pred = pred + self.stg_scale * (pred_cond - pred_cond_stg)
        if self.is_modality_enabled():
            pred = pred + (self.modality_scale - 1.0) * (pred_cond - pred_cond_modality)
        if self.guidance_rescale > 0.0:
            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)

	# Copyright 2026 The HuggingFace Team. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	import math
	from typing import TYPE_CHECKING

	import torch

	from ..configuration_utils import register_to_config
	from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg


	if TYPE_CHECKING:
	from ..modular_pipelines.modular_pipeline import BlockState


	class LTX2Guidance(BaseGuidance):
	"""
	Single-modality guider for LTX-2.X. Combines up to three terms via the delta formulation:
	- classifier-free guidance (CFG),
	- spatio-temporal guidance (STG), an extra pass with a set of transformer blocks perturbed,
	- modality-isolation guidance, an extra pass with A2V/V2A cross-attention disabled.

	One instance drives one modality, so a multimodal checkpoint carries several: LTX-2.X pairs a video `guider`
	with an `audio_guider`, each holding its own scales. The final guidance estimate is computed in whatever space
	the caller feeds it — the LTX-2 denoiser converts velocity->x0 before calling this and back afterwards, so
	`forward` operates on x0 predictions.

	Unlike `SkipLayerGuidance`, which applies its perturbations with hooks, `LTX2Guidance` expects its non-CFG
	perturbations to be configurable through model forward arguments (`spatio_temporal_guidance_blocks` and
	`isolate_modalities` in `LTX2VideoTransformer3DModel`). The calling denoiser block is responsible for setting
	those per pass identifier; see `LTX2LoopDenoiser` for a usage example.

	Args:
	guidance_scale (`float`, defaults to `1.0`):
	CFG scale for this modality. The CFG pass is skipped entirely at `1.0`.
	stg_scale (`float`, defaults to `0.0`):
	Spatio-temporal guidance scale for this modality. The STG pass is skipped entirely at `0.0`.
	modality_scale (`float`, defaults to `1.0`):
	Modality-isolation guidance scale for this modality. That pass is skipped entirely at `1.0`.
	guidance_rescale (`float`, defaults to `0.0`):
	Rescaling factor to prevent overexposure from high guidance scales. Based on [Common Diffusion Noise
	Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Range: 0.0 (no rescaling)
	to 1.0 (full rescaling).
	spatio_temporal_guidance_blocks (`list[int]`, optional):
	Transformer blocks to perturb on the STG pass. How the value is consumed is up to the denoiser block:
	LTX-2's perturbs whole blocks in a single forward that feeds every modality, so it reads this off the video
	guider alone — setting it on the audio guider has no effect, and that guider joins the same STG forward
	through its own `stg_scale`.
	start (`float`, defaults to `0.0`):
	Fraction of denoising steps (0.0-1.0) after which guidance starts.
	stop (`float`, defaults to `1.0`):
	Fraction of denoising steps (0.0-1.0) after which guidance stops.
	enabled (`bool`, defaults to `True`):
	Whether this guider applies guidance at all. When `False`, only the conditional pass runs.
	"""

	# `pred_cond` is the base; the others are the guidance passes.
	_input_predictions = ["pred_cond", "pred_uncond", "pred_cond_stg", "pred_cond_modality"]

	@register_to_config
	def __init__(
	self,
	guidance_scale: float = 1.0,
	stg_scale: float = 0.0,
	modality_scale: float = 1.0,
	guidance_rescale: float = 0.0,
	spatio_temporal_guidance_blocks: list[int] \| None = None,
	start: float = 0.0,
	stop: float = 1.0,
	enabled: bool = True,
	):
	super().__init__(start, stop, enabled)
	self.guidance_scale = guidance_scale
	self.stg_scale = stg_scale
	self.modality_scale = modality_scale
	self.guidance_rescale = guidance_rescale
	self.spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks

	# --- per-term enabled checks (public: the denoiser ORs these across the two guiders) --------------------
	def _within_window(self) -> bool:
	if self._num_inference_steps is None:
	return True
	start_step = int(self._start * self._num_inference_steps)
	stop_step = int(self._stop * self._num_inference_steps)
	return start_step <= self._step < stop_step

	def is_cfg_enabled(self) -> bool:
	return self._enabled and self._within_window() and not math.isclose(self.guidance_scale, 1.0)

	def is_stg_enabled(self) -> bool:
	return self._enabled and self._within_window() and not math.isclose(self.stg_scale, 0.0)

	def is_modality_enabled(self) -> bool:
	return self._enabled and self._within_window() and not math.isclose(self.modality_scale, 1.0)

	def active_predictions(self) -> list[str]:
	preds = ["pred_cond"]
	if self.is_cfg_enabled():
	preds.append("pred_uncond")
	if self.is_stg_enabled():
	preds.append("pred_cond_stg")
	if self.is_modality_enabled():
	preds.append("pred_cond_modality")
	return preds

	@property
	def num_conditions(self) -> int:
	return len(self.active_predictions())

	@property
	def is_conditional(self) -> bool:
	# Not load-bearing here: the denoiser owns pass execution and carries the STG/modality flags on the
	# batches, so we never drive `prepare_models`/hooks off `_count_prepared` the way SkipLayerGuidance does.
	return True

	# Each pass reads its own slot ([cond, uncond, stg, modality]) of every `guider_input_fields` tuple, selected
	# by this per-identifier index. Distinct slots (rather than cond/uncond only) let STG/modality pull the same
	# conditioning as `pred_cond` while the denoiser overrides their per-pass model flags after preparation.
	_PREDICTION_INDEX = {"pred_cond": 0, "pred_uncond": 1, "pred_cond_stg": 2, "pred_cond_modality": 3}

	def prepare_inputs(self, guider_inputs: dict) -> list["BlockState"]:
	# One identifier-tagged batch per active pass. Every value in `guider_inputs` is a 4-tuple indexed by
	# `_PREDICTION_INDEX` ([cond, uncond, stg, modality]); each pass reads its own slot, so the per-pass model
	# flags (`spatio_temporal_guidance_blocks`, `isolate_modalities`) are carried exactly like the encoder
	# inputs. The denoiser fills each returned batch's `noise_pred` before calling the guider to combine.
	return [
	self._prepare_batch(guider_inputs, self._PREDICTION_INDEX[pred], pred)
	for pred in self.active_predictions()
	]

	def prepare_inputs_from_block_state(self, data: "BlockState", input_fields: dict) -> list["BlockState"]:
	# One identifier-tagged batch per active pass. Each value in `input_fields` maps a transformer argument to a
	# 4-tuple of block-state attribute names indexed by `_PREDICTION_INDEX` ([cond, uncond, stg, modality]); the
	# base helper reads the pass's slot off `data`. The denoiser then sets the per-pass model flags and fills
	# each batch's `noise_pred` before calling the guider to combine.
	return [
	self._prepare_batch_from_block_state(input_fields, data, self._PREDICTION_INDEX[pred], pred)
	for pred in self.active_predictions()
	]

	def forward(
	self,
	pred_cond: torch.Tensor,
	pred_uncond: torch.Tensor \| None = None,
	pred_cond_stg: torch.Tensor \| None = None,
	pred_cond_modality: torch.Tensor \| None = None,
	) -> GuiderOutput:
	# Delta formulation (matches the reference LTX-2 combine, applied in the x0 space the denoiser feeds):
	# x0_cond + (gs-1)(cond-uncond) + stg(cond-stg) + (mod-1)*(cond-modality), then optional rescale.
	pred = pred_cond
	if self.is_cfg_enabled():
	pred = pred + (self.guidance_scale - 1.0) * (pred_cond - pred_uncond)
	if self.is_stg_enabled():
	pred = pred + self.stg_scale * (pred_cond - pred_cond_stg)
	if self.is_modality_enabled():
	pred = pred + (self.modality_scale - 1.0) * (pred_cond - pred_cond_modality)
	if self.guidance_rescale > 0.0:
	pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
	return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)