pom/v_diffusion.py

# v-diffusion codes for DDPM inpainting.  May not be compatible with k-diffusion.

# @SuspectT's inpainting codes, Feb 25 2024
# shared w/ me over Discord:
# "that's the v-diffusion inpainting with ddpm
# optimal settings were around 100 steps for the scheduler 
# (ts refering to timesteps here) and resamples was 4"

import torch
from torch import nn
from typing import Callable
from tqdm import trange
import math
import sys

# from kcrowson/v-diffusion-pytorch
def t_to_alpha_sigma(t):
    """Returns the scaling factors for the clean image and for the noise, given
    a timestep."""
    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)



#class DDPM(SamplerBase):
class DDPM():
    
    def __init__(self, model_fn: Callable = None):
        super().__init__()
    
    def _step(
        self, model_fn: Callable, x_t: torch.Tensor, step: int,
        t_now: torch.Tensor, t_next: torch.Tensor,
        callback: Callable, model_args,  **sampler_args ) -> torch.Tensor:
        
        alpha_now, sigma_now = t_to_alpha_sigma(t_now) # Get alpha / sigma for current timestep.
        alpha_next, sigma_next = t_to_alpha_sigma(t_next) # Get alpha / sigma for next timestep.
        
        v_t = model_fn(x_t, t_now.expand(x_t.shape[0]), **model_args) # Expand t to match batch_size which corresponds to x_t.shape[0]
        
        eps_t = x_t * sigma_now + v_t * alpha_now
        pred_t = x_t * alpha_now - v_t * sigma_now
        
        if callback is not None:
            callback({'step': step, 'x': x_t, 't': t_now, 'pred': pred_t, 'eps': eps_t})
            
        return (pred_t * alpha_next + eps_t * sigma_next)
    
    def _sample( self, model_fn: Callable, x_t: torch.Tensor, ts: torch.Tensor,
        callback: Callable, model_args, **sampler_args ) -> torch.Tensor:
        
        print("Using DDPM Sampler.")
        steps = ts.size(0)
        
        use_tqdm = sampler_args.get('use_tqdm')
        use_range = trange if (use_tqdm if (use_tqdm != None) else False) else range
        
        for step in use_range(steps - 1):
            x_t = self._step( model_fn, x_t, step, ts[step], ts[step + 1],
                lambda kwargs: callback(**dict(kwargs, steps=steps)) if(callback != None) else None,
                model_args  )
            
        return x_t
    
      
    def _inpaint(self,
        model_fn: Callable, audio_source: torch.Tensor, mask: torch.Tensor,
        ts: torch.Tensor,   resamples: int, callback: Callable, model_args, **sampler_args
        ) -> torch.Tensor:
        steps = ts.size(0)
        batch_size = audio_source.size(0)
        alphas, sigmas = t_to_alpha_sigma(ts)

        # SHH: rescale audio_source to zero mean and unit variance
        audio_source = (audio_source - audio_source.mean()) / audio_source.std()

        x_t = audio_source
        
        use_tqdm = sampler_args.get('use_tqdm')
        use_range = trange if (use_tqdm if (use_tqdm != None) else False) else range
        
        for step in use_range(steps - 1):
            print("step, audio_source.min, audio_source.max, alphas[step], sigmas[step] = ", step, audio_source.min(), audio_source.max(), alphas[step], sigmas[step])
            audio_source_noised = audio_source * alphas[step] + torch.randn_like(audio_source) * sigmas[step]
            print("step, audio_source_noised.min, audio_source_noised.max = ", step, audio_source_noised.min(), audio_source_noised.max())
            sigma_dt = torch.sqrt(sigmas[step] ** 2 - sigmas[step + 1] ** 2)
            
            for re in range(resamples):
                
                #x_t = audio_source_noised * mask + x_t * ~mask 
                x_t = audio_source_noised * mask + x_t * (1.0-mask) 
         
                # from ImageTransformerDenoiserModelV2:
                #        def forward(self, x, sigma, aug_cond=None, class_cond=None, mapping_cond=None):
                #v_t = model_fn(x_t, ts[step].expand(batch_size), **model_args)
                print("step, re, x_t.min, x_t.max , sigmas[step]= ", step, re, x_t.min(), x_t.max(), sigmas[step])
                v_t = model_fn(x_t, sigmas[step].expand(batch_size), aug_cond=None, class_cond=None, mapping_cond=None)
                print("step, re, v_t.min, v_t.max = ", step, re, v_t.min(), v_t.max())
                if v_t.isnan().any():
                    print("v_t has NaNs.")
                    sys.exit(0)
                
                eps_t = x_t * sigmas[step] + v_t * alphas[step]
                pred_t = x_t * alphas[step] - v_t * sigmas[step]
                
                if callback is not None:
                    callback({'steps': steps, 'step': step, 'x': x_t, 't': ts[step], 'pred': pred_t, 'eps': eps_t, 'res': re})
                
                if(re < resamples - 1):
                    x_t = pred_t * alphas[step] + eps_t * sigmas[step + 1] + sigma_dt * torch.randn_like(x_t)
                else:
                    x_t = pred_t * alphas[step + 1] + eps_t * sigmas[step + 1]

                print("step, re, v_t.min, v_t.max, x_t.min, x_t.max = ", step, re, v_t.min(), v_t.max(), x_t.min(), x_t.max())

                #sys.exit(0)

        return (audio_source * mask + x_t * (1.0-mask))


def alpha_sigma_to_t(alpha, sigma):
    """Returns a timestep, given the scaling factors for the clean image and for
    the noise."""
    return torch.atan2(sigma, alpha) / math.pi * 2  

def log_snr_to_alpha_sigma(log_snr):
    """Returns the scaling factors for the clean image and for the noise, given
    the log SNR for a timestep."""
    return log_snr.sigmoid().sqrt(), log_snr.neg().sigmoid().sqrt()

def get_ddpm_schedule(ddpm_t):
    """Returns timesteps for the noise schedule from the DDPM paper."""
    log_snr = -torch.special.expm1(1e-4 + 10 * ddpm_t**2).log()
    alpha, sigma = log_snr_to_alpha_sigma(log_snr)
    return alpha_sigma_to_t(alpha, sigma)


#class LogSchedule(SchedulerBase):
class LogSchedule():
    def __init__(self, device:torch.device = None):
        super().__init__(device)
        
    def create(self, steps: int, first: float = 1, last: float = 0, device: torch.device = None, scheduler_args = {'min_log_snr': -10, 'max_log_snr': 10}) -> torch.Tensor:
        ramp = torch.linspace(first, last, steps, device = device if (device != None) else self.device)
        min_log_snr = scheduler_args.get('min_log_snr')
        max_log_snr = scheduler_args.get('max_log_snr')
        return self.get_log_schedule(
            ramp,
            min_log_snr if min_log_snr!=None else -10,
            max_log_snr if max_log_snr!=None else 10,
        )
        
    def get_log_schedule(self, t, min_log_snr=-10, max_log_snr=10):
        log_snr = t * (min_log_snr - max_log_snr) + max_log_snr
        alpha = log_snr.sigmoid().sqrt()
        sigma = log_snr.neg().sigmoid().sqrt()
        return torch.atan2(sigma, alpha) / math.pi * 2  # this returns a timestep?


#class CrashSchedule(SchedulerBase):
class CrashSchedule():
    def __init__(self, device:torch.device = None):
        super().__init__(device)
    
    def create(self, steps: int, first: float = 1, last: float = 0, device: torch.device = None, scheduler_args = None) -> torch.Tensor:
        ramp = torch.linspace(first, last, steps, device = device if (device != None) else self.device)
        sigma = torch.sin(ramp * math.pi / 2) ** 2
        alpha = (1 - sigma**2) ** 0.5
        return torch.atan2(sigma, alpha) / math.pi * 2 # this returns a timestep?