code/models/tt_dit/solvers/euler.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /solvers /euler.py

Add files using upload-large-folder tool

629b123 verified 11 days ago

2.59 kB

# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import ttnn

from .base import Solver
from .schedule import Schedule

if TYPE_CHECKING:
    from collections.abc import Sequence

    from diffusers.schedulers import FlowMatchEulerDiscreteScheduler


class EulerSolver(Solver):
    def __init__(self, *, scheduler: FlowMatchEulerDiscreteScheduler | None = None) -> None:
        super().__init__(scheduler=scheduler)

        # Only meaningful with a scheduler; without one every schedule is taken as given.
        self._default_shift = scheduler.shift if scheduler is not None else None

    def set_schedule(
        self,
        num_inference_steps: int | None = None,
        *,
        sigmas: Sequence[float] | None = None,
        shift: float | None = None,
        **kwargs: Any,
    ) -> None:
        """Derive and adopt the schedule for one denoising run.
        Args:
            num_inference_steps: Number of denoising steps.
            sigmas: Explicit sigma schedule. Required when the solver has no scheduler, and
                then taken exactly as given; otherwise forwarded to `set_timesteps`.
            shift: Flow shift for this run only; the scheduler's construction-time value is
                used when omitted. Only accepted when the solver has a scheduler.
            kwargs: Forwarded to the scheduler's `set_timesteps`.
        """
        if self._scheduler is None:
            if sigmas is None or num_inference_steps is not None or shift is not None or kwargs:
                msg = "a scheduler-less EulerSolver accepts only `sigmas`"
                raise ValueError(msg)
            self._schedule = Schedule.from_sigmas(sigmas)
            return

        # set_shift writes `_shift`, which is what set_timesteps reads. The config copy is
        # currently unused, but keep it in sync so anything re-deriving from the config
        # (from_config, save_config) sees the shift actually in effect.
        value = self._default_shift if shift is None else shift
        self._scheduler.set_shift(value)
        self._scheduler.register_to_config(shift=value)
        self._schedule = Schedule.from_scheduler(self._scheduler, num_inference_steps, sigmas=sigmas, **kwargs)

    def step(self, *, step: int, latent: ttnn.Tensor, velocity_pred: ttnn.Tensor) -> ttnn.Tensor:
        sigma_curr = self.sigmas[step]
        sigma_next = self.sigmas[step + 1]

        return latent + (sigma_next - sigma_curr) * velocity_pred

	# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC

	# SPDX-License-Identifier: Apache-2.0

	from __future__ import annotations

	from typing import TYPE_CHECKING, Any

	import ttnn

	from .base import Solver
	from .schedule import Schedule

	if TYPE_CHECKING:
	from collections.abc import Sequence

	from diffusers.schedulers import FlowMatchEulerDiscreteScheduler


	class EulerSolver(Solver):
	def __init__(self, *, scheduler: FlowMatchEulerDiscreteScheduler \| None = None) -> None:
	super().__init__(scheduler=scheduler)

	# Only meaningful with a scheduler; without one every schedule is taken as given.
	self._default_shift = scheduler.shift if scheduler is not None else None

	def set_schedule(
	self,
	num_inference_steps: int \| None = None,
	*,
	sigmas: Sequence[float] \| None = None,
	shift: float \| None = None,
	**kwargs: Any,
	) -> None:
	"""Derive and adopt the schedule for one denoising run.

	Args:
	num_inference_steps: Number of denoising steps.
	sigmas: Explicit sigma schedule. Required when the solver has no scheduler, and
	then taken exactly as given; otherwise forwarded to `set_timesteps`.
	shift: Flow shift for this run only; the scheduler's construction-time value is
	used when omitted. Only accepted when the solver has a scheduler.
	kwargs: Forwarded to the scheduler's `set_timesteps`.
	"""
	if self._scheduler is None:
	if sigmas is None or num_inference_steps is not None or shift is not None or kwargs:
	msg = "a scheduler-less EulerSolver accepts only `sigmas`"
	raise ValueError(msg)
	self._schedule = Schedule.from_sigmas(sigmas)
	return

	# set_shift writes `_shift`, which is what set_timesteps reads. The config copy is
	# currently unused, but keep it in sync so anything re-deriving from the config
	# (from_config, save_config) sees the shift actually in effect.
	value = self._default_shift if shift is None else shift
	self._scheduler.set_shift(value)
	self._scheduler.register_to_config(shift=value)
	self._schedule = Schedule.from_scheduler(self._scheduler, num_inference_steps, sigmas=sigmas, **kwargs)

	def step(self, *, step: int, latent: ttnn.Tensor, velocity_pred: ttnn.Tensor) -> ttnn.Tensor:
	sigma_curr = self.sigmas[step]
	sigma_next = self.sigmas[step + 1]

	return latent + (sigma_next - sigma_curr) * velocity_pred