Upload RotEx inference code, demos, documentation, and pretrained checkpoints

40440d6 verified 21 days ago

13.4 kB

"""
State-of-the-art sampling techniques (2024-2025).
Includes Flow Matching with Optimal Transport and Consistency Models.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Callable
import math


class FlowMatching:
    """
    Flow Matching with Optimal Transport for improved trajectory learning.
    Based on recent advances in continuous normalizing flows.
    """
    def __init__(self, sigma_min: float = 0.001):
        """
        Args:
            sigma_min: Minimum noise level for numerical stability
        """
        self.sigma_min = sigma_min
    
    def compute_optimal_transport_path(self, x0: torch.Tensor, x1: torch.Tensor, 
                                      t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Compute optimal transport path from noise (x0) to data (x1).
        
        Args:
            x0: Source distribution (noise) [B, ...]
            x1: Target distribution (data) [B, ...]
            t: Time steps [B] in range [0, 1]
        Returns:
            x_t: Interpolated samples
            v_t: Vector field (velocity)
        """
        t = t.view(-1, *([1] * (x0.dim() - 1)))
        
        # Optimal transport interpolation (linear for Gaussian)
        x_t = (1 - t) * x0 + t * x1
        
        # Velocity is the derivative: v = dx/dt = x1 - x0
        v_t = x1 - x0
        
        return x_t, v_t
    
    def compute_conditional_flow(self, x0: torch.Tensor, x1: torch.Tensor,
                                t: torch.Tensor, sigma: Optional[float] = None) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Compute conditional flow with added noise for robustness.
        
        Args:
            x0: Noise [B, ...]
            x1: Data [B, ...]
            t: Time [B]
            sigma: Optional noise level (defaults to sigma_min * (1-t))
        Returns:
            x_t: Noisy interpolation
            v_t: Target velocity
        """
        if sigma is None:
            sigma = self.sigma_min * (1 - t.view(-1, *([1] * (x0.dim() - 1))))
        else:
            sigma = sigma * torch.ones_like(t).view(-1, *([1] * (x0.dim() - 1)))
        
        # Conditional path
        x_t, v_t = self.compute_optimal_transport_path(x0, x1, t)
        
        # Add Gaussian noise for conditional flow matching
        eps = torch.randn_like(x_t)
        x_t = x_t + sigma * eps
        
        return x_t, v_t
    
    def sample(self, model: Callable, x_init: torch.Tensor, 
              num_steps: int = 50, **model_kwargs) -> torch.Tensor:
        """
        Sample using Euler method for ODE solving.
        
        Args:
            model: Velocity prediction model
            x_init: Initial noise [B, ...]
            num_steps: Number of integration steps
            **model_kwargs: Additional arguments for model
        Returns:
            x_final: Generated samples
        """
        x = x_init
        dt = 1.0 / num_steps
        
        for step in range(num_steps):
            t_current = torch.full((x.shape[0],), step / num_steps, 
                                  device=x.device, dtype=x.dtype)
            
            # Predict velocity
            with torch.no_grad():
                v = model(x, timesteps=t_current, **model_kwargs)
            
            # Euler step
            x = x + v * dt
        
        return x
    
    def sample_heun(self, model: Callable, x_init: torch.Tensor,
                   num_steps: int = 50, **model_kwargs) -> torch.Tensor:
        """
        Higher-order sampling using Heun's method (2nd order).
        More accurate but requires 2x model evaluations.
        
        Args:
            model: Velocity prediction model
            x_init: Initial noise
            num_steps: Number of steps
            **model_kwargs: Model arguments
        Returns:
            x_final: Generated samples
        """
        x = x_init
        dt = 1.0 / num_steps
        
        for step in range(num_steps):
            t_current = torch.full((x.shape[0],), step / num_steps,
                                  device=x.device, dtype=x.dtype)
            
            # First stage
            with torch.no_grad():
                v1 = model(x, timesteps=t_current, **model_kwargs)
            
            x_temp = x + v1 * dt
            
            # Second stage
            t_next = torch.full((x.shape[0],), (step + 1) / num_steps,
                               device=x.device, dtype=x.dtype)
            
            with torch.no_grad():
                v2 = model(x_temp, timesteps=t_next, **model_kwargs)
            
            # Heun's method: average of two velocities
            x = x + (v1 + v2) * 0.5 * dt
        
        return x


class ConsistencyModel:
    """
    Consistency Models for fast, single-step generation.
    Based on "Consistency Models" (Song et al., 2023).
    """
    def __init__(self, sigma_min: float = 0.002, sigma_max: float = 80.0,
                 rho: float = 7.0):
        """
        Args:
            sigma_min: Minimum noise level
            sigma_max: Maximum noise level
            rho: Schedule parameter
        """
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.rho = rho
    
    def get_schedule(self, num_steps: int, device: torch.device) -> torch.Tensor:
        """
        Get noise schedule for consistency training.
        
        Args:
            num_steps: Number of discretization steps
            device: Device for tensors
        Returns:
            sigmas: Noise levels [num_steps + 1]
        """
        rho_inv = 1.0 / self.rho
        steps = torch.arange(num_steps + 1, device=device)
        sigmas = (self.sigma_max ** rho_inv + 
                 steps / num_steps * (self.sigma_min ** rho_inv - self.sigma_max ** rho_inv)) ** self.rho
        return sigmas
    
    def consistency_function(self, model: Callable, x: torch.Tensor,
                            sigma: torch.Tensor, **model_kwargs) -> torch.Tensor:
        """
        Apply consistency function: maps (x, sigma) to denoised output.
        
        Args:
            model: Denoising model
            x: Noisy input
            sigma: Noise level
            **model_kwargs: Model arguments
        Returns:
            denoised: Consistency function output
        """
        # Skip connection for sigma_min
        c_skip = self.sigma_min ** 2 / (sigma ** 2 + self.sigma_min ** 2)
        c_out = sigma * self.sigma_min / (sigma ** 2 + self.sigma_min ** 2).sqrt()
        
        # Predict denoised output
        with torch.no_grad():
            # Convert sigma to timestep for model
            t = sigma / self.sigma_max
            model_out = model(x, timesteps=t, **model_kwargs)
        
        # Apply consistency function
        denoised = c_skip * x + c_out * model_out
        
        return denoised
    
    def sample_onestep(self, model: Callable, shape: tuple,
                      device: torch.device, dtype: torch.dtype,
                      **model_kwargs) -> torch.Tensor:
        """
        Single-step generation from noise to data.
        
        Args:
            model: Consistency model
            shape: Output shape
            device: Device
            dtype: Data type
            **model_kwargs: Model arguments
        Returns:
            samples: Generated samples
        """
        # Sample from prior (high noise)
        x = torch.randn(*shape, device=device, dtype=dtype) * self.sigma_max
        sigma = torch.full((shape[0],), self.sigma_max, device=device, dtype=dtype)
        
        # One-step denoising
        samples = self.consistency_function(model, x, sigma, **model_kwargs)
        
        return samples
    
    def sample_multistep(self, model: Callable, shape: tuple,
                        num_steps: int, device: torch.device,
                        dtype: torch.dtype, **model_kwargs) -> torch.Tensor:
        """
        Multi-step generation for better quality.
        
        Args:
            model: Consistency model
            shape: Output shape
            num_steps: Number of steps
            device: Device
            dtype: Data type
            **model_kwargs: Model arguments
        Returns:
            samples: Generated samples
        """
        # Get noise schedule
        sigmas = self.get_schedule(num_steps, device)
        
        # Start from maximum noise
        x = torch.randn(*shape, device=device, dtype=dtype) * self.sigma_max
        
        # Iterative refinement
        for i in range(num_steps, 0, -1):
            sigma_current = sigmas[i]
            sigma_next = sigmas[i - 1]
            
            # Denoise to current level
            x = self.consistency_function(
                model, x, 
                torch.full((shape[0],), sigma_current, device=device, dtype=dtype),
                **model_kwargs
            )
            
            # Add noise for next iteration (except last)
            if i > 1:
                noise = torch.randn_like(x)
                x = x + noise * (sigma_next ** 2 - self.sigma_min ** 2).sqrt()
        
        return x


class ClassifierFreeGuidance:
    """
    Enhanced Classifier-Free Guidance for better conditioning.
    """
    def __init__(self, guidance_scale: float = 7.5, min_scale: float = 1.0):
        """
        Args:
            guidance_scale: Guidance strength
            min_scale: Minimum guidance scale
        """
        self.guidance_scale = guidance_scale
        self.min_scale = min_scale
    
    def apply_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
                      guidance_scale: Optional[float] = None) -> torch.Tensor:
        """
        Apply classifier-free guidance to model outputs.
        
        Args:
            cond_out: Conditional model output
            uncond_out: Unconditional model output
            guidance_scale: Optional override for guidance scale
        Returns:
            guided_out: Guided output
        """
        scale = guidance_scale if guidance_scale is not None else self.guidance_scale
        scale = max(scale, self.min_scale)
        
        # CFG formula: out = uncond + scale * (cond - uncond)
        guided_out = uncond_out + scale * (cond_out - uncond_out)
        
        return guided_out
    
    def apply_dynamic_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
                              t: torch.Tensor, min_t: float = 0.0, 
                              max_t: float = 1.0) -> torch.Tensor:
        """
        Apply time-dependent guidance that decreases over time.
        
        Args:
            cond_out: Conditional output
            uncond_out: Unconditional output
            t: Time steps [B]
            min_t: Time to start decreasing guidance
            max_t: Time to reach minimum guidance
        Returns:
            guided_out: Dynamically guided output
        """
        # Linearly decrease guidance scale from max to min over time
        t_normalized = (t - min_t) / (max_t - min_t)
        t_normalized = t_normalized.clamp(0, 1)
        
        # Scale goes from guidance_scale to min_scale
        scale = self.guidance_scale + (self.min_scale - self.guidance_scale) * t_normalized
        scale = scale.view(-1, *([1] * (cond_out.dim() - 1)))
        
        guided_out = uncond_out + scale * (cond_out - uncond_out)
        
        return guided_out


class AdaptiveSampler:
    """
    Adaptive sampling that combines multiple techniques for optimal quality/speed.
    """
    def __init__(self, use_flow_matching: bool = True,
                 use_consistency: bool = False,
                 use_cfg: bool = True):
        self.use_flow_matching = use_flow_matching
        self.use_consistency = use_consistency
        self.use_cfg = use_cfg
        
        if use_flow_matching:
            self.flow_matching = FlowMatching()
        if use_consistency:
            self.consistency_model = ConsistencyModel()
        if use_cfg:
            self.cfg = ClassifierFreeGuidance()
    
    def sample(self, model: Callable, x_init: torch.Tensor,
              num_steps: int = 50, method: str = "heun",
              cfg_scale: Optional[float] = None,
              **model_kwargs) -> torch.Tensor:
        """
        Adaptive sampling with multiple techniques.
        
        Args:
            model: Model for sampling
            x_init: Initial noise
            num_steps: Number of steps
            method: Sampling method ("euler", "heun", "consistency")
            cfg_scale: Optional CFG scale
            **model_kwargs: Model arguments
        Returns:
            samples: Generated samples
        """
        if method == "consistency" and self.use_consistency:
            return self.consistency_model.sample_multistep(
                model, x_init.shape, num_steps,
                x_init.device, x_init.dtype, **model_kwargs
            )
        elif method == "heun" and self.use_flow_matching:
            return self.flow_matching.sample_heun(
                model, x_init, num_steps, **model_kwargs
            )
        else:  # euler
            return self.flow_matching.sample(
                model, x_init, num_steps, **model_kwargs
            )

	"""
	State-of-the-art sampling techniques (2024-2025).
	Includes Flow Matching with Optimal Transport and Consistency Models.
	"""
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from typing import Optional, Callable
	import math


	class FlowMatching:
	"""
	Flow Matching with Optimal Transport for improved trajectory learning.
	Based on recent advances in continuous normalizing flows.
	"""
	def __init__(self, sigma_min: float = 0.001):
	"""
	Args:
	sigma_min: Minimum noise level for numerical stability
	"""
	self.sigma_min = sigma_min

	def compute_optimal_transport_path(self, x0: torch.Tensor, x1: torch.Tensor,
	t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
	"""
	Compute optimal transport path from noise (x0) to data (x1).

	Args:
	x0: Source distribution (noise) [B, ...]
	x1: Target distribution (data) [B, ...]
	t: Time steps [B] in range [0, 1]
	Returns:
	x_t: Interpolated samples
	v_t: Vector field (velocity)
	"""
	t = t.view(-1, ([1] (x0.dim() - 1)))

	# Optimal transport interpolation (linear for Gaussian)
	x_t = (1 - t) * x0 + t * x1

	# Velocity is the derivative: v = dx/dt = x1 - x0
	v_t = x1 - x0

	return x_t, v_t

	def compute_conditional_flow(self, x0: torch.Tensor, x1: torch.Tensor,
	t: torch.Tensor, sigma: Optional[float] = None) -> tuple[torch.Tensor, torch.Tensor]:
	"""
	Compute conditional flow with added noise for robustness.

	Args:
	x0: Noise [B, ...]
	x1: Data [B, ...]
	t: Time [B]
	sigma: Optional noise level (defaults to sigma_min * (1-t))
	Returns:
	x_t: Noisy interpolation
	v_t: Target velocity
	"""
	if sigma is None:
	sigma = self.sigma_min * (1 - t.view(-1, ([1] (x0.dim() - 1))))
	else:
	sigma = sigma * torch.ones_like(t).view(-1, ([1] (x0.dim() - 1)))

	# Conditional path
	x_t, v_t = self.compute_optimal_transport_path(x0, x1, t)

	# Add Gaussian noise for conditional flow matching
	eps = torch.randn_like(x_t)
	x_t = x_t + sigma * eps

	return x_t, v_t

	def sample(self, model: Callable, x_init: torch.Tensor,
	num_steps: int = 50, **model_kwargs) -> torch.Tensor:
	"""
	Sample using Euler method for ODE solving.

	Args:
	model: Velocity prediction model
	x_init: Initial noise [B, ...]
	num_steps: Number of integration steps
	**model_kwargs: Additional arguments for model
	Returns:
	x_final: Generated samples
	"""
	x = x_init
	dt = 1.0 / num_steps

	for step in range(num_steps):
	t_current = torch.full((x.shape[0],), step / num_steps,
	device=x.device, dtype=x.dtype)

	# Predict velocity
	with torch.no_grad():
	v = model(x, timesteps=t_current, **model_kwargs)

	# Euler step
	x = x + v * dt

	return x

	def sample_heun(self, model: Callable, x_init: torch.Tensor,
	num_steps: int = 50, **model_kwargs) -> torch.Tensor:
	"""
	Higher-order sampling using Heun's method (2nd order).
	More accurate but requires 2x model evaluations.

	Args:
	model: Velocity prediction model
	x_init: Initial noise
	num_steps: Number of steps
	**model_kwargs: Model arguments
	Returns:
	x_final: Generated samples
	"""
	x = x_init
	dt = 1.0 / num_steps

	for step in range(num_steps):
	t_current = torch.full((x.shape[0],), step / num_steps,
	device=x.device, dtype=x.dtype)

	# First stage
	with torch.no_grad():
	v1 = model(x, timesteps=t_current, **model_kwargs)

	x_temp = x + v1 * dt

	# Second stage
	t_next = torch.full((x.shape[0],), (step + 1) / num_steps,
	device=x.device, dtype=x.dtype)

	with torch.no_grad():
	v2 = model(x_temp, timesteps=t_next, **model_kwargs)

	# Heun's method: average of two velocities
	x = x + (v1 + v2) * 0.5 * dt

	return x


	class ConsistencyModel:
	"""
	Consistency Models for fast, single-step generation.
	Based on "Consistency Models" (Song et al., 2023).
	"""
	def __init__(self, sigma_min: float = 0.002, sigma_max: float = 80.0,
	rho: float = 7.0):
	"""
	Args:
	sigma_min: Minimum noise level
	sigma_max: Maximum noise level
	rho: Schedule parameter
	"""
	self.sigma_min = sigma_min
	self.sigma_max = sigma_max
	self.rho = rho

	def get_schedule(self, num_steps: int, device: torch.device) -> torch.Tensor:
	"""
	Get noise schedule for consistency training.

	Args:
	num_steps: Number of discretization steps
	device: Device for tensors
	Returns:
	sigmas: Noise levels [num_steps + 1]
	"""
	rho_inv = 1.0 / self.rho
	steps = torch.arange(num_steps + 1, device=device)
	sigmas = (self.sigma_max ** rho_inv +
	steps / num_steps * (self.sigma_min rho_inv - self.sigma_max rho_inv)) ** self.rho
	return sigmas

	def consistency_function(self, model: Callable, x: torch.Tensor,
	sigma: torch.Tensor, **model_kwargs) -> torch.Tensor:
	"""
	Apply consistency function: maps (x, sigma) to denoised output.

	Args:
	model: Denoising model
	x: Noisy input
	sigma: Noise level
	**model_kwargs: Model arguments
	Returns:
	denoised: Consistency function output
	"""
	# Skip connection for sigma_min
	c_skip = self.sigma_min 2 / (sigma 2 + self.sigma_min ** 2)
	c_out = sigma * self.sigma_min / (sigma 2 + self.sigma_min 2).sqrt()

	# Predict denoised output
	with torch.no_grad():
	# Convert sigma to timestep for model
	t = sigma / self.sigma_max
	model_out = model(x, timesteps=t, **model_kwargs)

	# Apply consistency function
	denoised = c_skip * x + c_out * model_out

	return denoised

	def sample_onestep(self, model: Callable, shape: tuple,
	device: torch.device, dtype: torch.dtype,
	**model_kwargs) -> torch.Tensor:
	"""
	Single-step generation from noise to data.

	Args:
	model: Consistency model
	shape: Output shape
	device: Device
	dtype: Data type
	**model_kwargs: Model arguments
	Returns:
	samples: Generated samples
	"""
	# Sample from prior (high noise)
	x = torch.randn(shape, device=device, dtype=dtype) self.sigma_max
	sigma = torch.full((shape[0],), self.sigma_max, device=device, dtype=dtype)

	# One-step denoising
	samples = self.consistency_function(model, x, sigma, **model_kwargs)

	return samples

	def sample_multistep(self, model: Callable, shape: tuple,
	num_steps: int, device: torch.device,
	dtype: torch.dtype, **model_kwargs) -> torch.Tensor:
	"""
	Multi-step generation for better quality.

	Args:
	model: Consistency model
	shape: Output shape
	num_steps: Number of steps
	device: Device
	dtype: Data type
	**model_kwargs: Model arguments
	Returns:
	samples: Generated samples
	"""
	# Get noise schedule
	sigmas = self.get_schedule(num_steps, device)

	# Start from maximum noise
	x = torch.randn(shape, device=device, dtype=dtype) self.sigma_max

	# Iterative refinement
	for i in range(num_steps, 0, -1):
	sigma_current = sigmas[i]
	sigma_next = sigmas[i - 1]

	# Denoise to current level
	x = self.consistency_function(
	model, x,
	torch.full((shape[0],), sigma_current, device=device, dtype=dtype),
	**model_kwargs
	)

	# Add noise for next iteration (except last)
	if i > 1:
	noise = torch.randn_like(x)
	x = x + noise * (sigma_next 2 - self.sigma_min 2).sqrt()

	return x


	class ClassifierFreeGuidance:
	"""
	Enhanced Classifier-Free Guidance for better conditioning.
	"""
	def __init__(self, guidance_scale: float = 7.5, min_scale: float = 1.0):
	"""
	Args:
	guidance_scale: Guidance strength
	min_scale: Minimum guidance scale
	"""
	self.guidance_scale = guidance_scale
	self.min_scale = min_scale

	def apply_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
	guidance_scale: Optional[float] = None) -> torch.Tensor:
	"""
	Apply classifier-free guidance to model outputs.

	Args:
	cond_out: Conditional model output
	uncond_out: Unconditional model output
	guidance_scale: Optional override for guidance scale
	Returns:
	guided_out: Guided output
	"""
	scale = guidance_scale if guidance_scale is not None else self.guidance_scale
	scale = max(scale, self.min_scale)

	# CFG formula: out = uncond + scale * (cond - uncond)
	guided_out = uncond_out + scale * (cond_out - uncond_out)

	return guided_out

	def apply_dynamic_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
	t: torch.Tensor, min_t: float = 0.0,
	max_t: float = 1.0) -> torch.Tensor:
	"""
	Apply time-dependent guidance that decreases over time.

	Args:
	cond_out: Conditional output
	uncond_out: Unconditional output
	t: Time steps [B]
	min_t: Time to start decreasing guidance
	max_t: Time to reach minimum guidance
	Returns:
	guided_out: Dynamically guided output
	"""
	# Linearly decrease guidance scale from max to min over time
	t_normalized = (t - min_t) / (max_t - min_t)
	t_normalized = t_normalized.clamp(0, 1)

	# Scale goes from guidance_scale to min_scale
	scale = self.guidance_scale + (self.min_scale - self.guidance_scale) * t_normalized
	scale = scale.view(-1, ([1] (cond_out.dim() - 1)))

	guided_out = uncond_out + scale * (cond_out - uncond_out)

	return guided_out


	class AdaptiveSampler:
	"""
	Adaptive sampling that combines multiple techniques for optimal quality/speed.
	"""
	def __init__(self, use_flow_matching: bool = True,
	use_consistency: bool = False,
	use_cfg: bool = True):
	self.use_flow_matching = use_flow_matching
	self.use_consistency = use_consistency
	self.use_cfg = use_cfg

	if use_flow_matching:
	self.flow_matching = FlowMatching()
	if use_consistency:
	self.consistency_model = ConsistencyModel()
	if use_cfg:
	self.cfg = ClassifierFreeGuidance()

	def sample(self, model: Callable, x_init: torch.Tensor,
	num_steps: int = 50, method: str = "heun",
	cfg_scale: Optional[float] = None,
	**model_kwargs) -> torch.Tensor:
	"""
	Adaptive sampling with multiple techniques.

	Args:
	model: Model for sampling
	x_init: Initial noise
	num_steps: Number of steps
	method: Sampling method ("euler", "heun", "consistency")
	cfg_scale: Optional CFG scale
	**model_kwargs: Model arguments
	Returns:
	samples: Generated samples
	"""
	if method == "consistency" and self.use_consistency:
	return self.consistency_model.sample_multistep(
	model, x_init.shape, num_steps,
	x_init.device, x_init.dtype, **model_kwargs
	)
	elif method == "heun" and self.use_flow_matching:
	return self.flow_matching.sample_heun(
	model, x_init, num_steps, **model_kwargs
	)
	else: # euler
	return self.flow_matching.sample(
	model, x_init, num_steps, **model_kwargs
	)