Add RF

diffusion/rectified_flow.py

@@ -0,0 +1,322 @@
+import torch
+from tqdm import tqdm
+import numpy as np
+import math
+class RectifiedFlow():
+    def __init__(self, num_timesteps, warmup_timesteps = 10, noise_scale=1.0, init_type='gaussian', eps=1, sampling='logit', window_size=8):
+        """
+        eps: A `float` number. The smallest time step to sample from.
+        """
+        self.num_timesteps = num_timesteps
+
+        self.warmup_timesteps = warmup_timesteps*num_timesteps
+        self.T = 1000.
+        self.noise_scale = noise_scale
+        self.init_type = init_type
+        self.eps = eps
+
+        self.window_size = window_size
+
+        self.sampling = sampling
+
+    def logit(self, x):
+        return torch.log(x / (1 - x))
+
+    def logit_normal(self, x, mu=0, sigma=1):
+        return 1 / (sigma * math.sqrt(2 * torch.pi) * x * (1 - x)) * torch.exp(-(self.logit(x) - mu) ** 2 / (2 * sigma ** 2))
+
+    def training_loss(self, model, v, a, model_kwargs):
+        """
+        v: [B, T, C, H, W]
+        a: [B, T, N, F]
+        """
+
+        B,T = v.shape[:2]
+            
+        tw = torch.rand((v.shape[0],1), device=v.device)
+
+        window_indexes = torch.linspace(0, self.window_size-1, steps=self.window_size, device=v.device).unsqueeze(0).repeat(B,1)
+        
+        rollout = torch.bernoulli(torch.tensor(0.8).repeat(B).to(v.device)).bool()
+        t_rollout = (window_indexes+tw)/self.window_size
+        
+        t_pre_rollout = window_indexes/self.window_size + tw
+
+        t = torch.where(rollout.unsqueeze(1).repeat(1,self.window_size), t_rollout, t_pre_rollout)
+        t = 1 - t # swap 0 and 1, since 1 is full image and 0 is full noise
+
+        t = torch.clamp(t, 0+1e-6, 1-1e-6)
+
+        if self.sampling == 'logit':
+            weigths = self.logit_normal(t, mu=0, sigma=1)
+        else:
+            weigths = torch.ones_like(t)
+
+        B, T = t.shape
+
+        v_z0 = self.get_z0(v).to(v.device)
+        a_z0 = self.get_z0(a).to(a.device)
+
+        t_video = t.view(B,T,1,1,1).repeat(1,1,v.shape[2], v.shape[3], v.shape[4]) 
+        t_audio = t.view(B,T,1,1,1).repeat(1,1,a.shape[2], a.shape[3], a.shape[4])
+
+        perturbed_video = t_video*v + (1-t_video)*v_z0
+        perturbed_audio = t_audio*a + (1-t_audio)*a_z0
+
+        t_rf = t*(self.T-self.eps) + self.eps
+        score_v, score_a = model(perturbed_video, perturbed_audio, t_rf, **model_kwargs)
+
+        # score_v = [B, T, C, H, W]
+        # score_a = [B, T, N, F]
+        target_video = v - v_z0 # direction of the flow
+        target_audio = a - a_z0 # direction of the flow
+
+        loss_video = torch.square(score_v-target_video)
+        loss_audio = torch.square(score_a-target_audio)
+
+        loss_video = torch.mean(loss_video, dim=[2,3,4])
+        loss_audio = torch.mean(loss_audio, dim=[2,3,4])
+
+        #mask out the loss for the time steps that are greater than T
+
+        loss_video = loss_video * (weigths)
+        loss_video = torch.mean(loss_video) 
+
+        loss_audio = loss_audio * (weigths)
+        loss_audio = torch.mean(loss_audio)
+
+        return {"loss": (loss_video + loss_audio)}
+
+    def sample(self, model, v_z, a_z, model_kwargs, progress=True):
+        B = v_z.shape[0]
+
+        window_indexes = torch.linspace(0, self.window_size-1, steps=self.window_size, device=v_z.device).unsqueeze(0).repeat(B,1)
+
+
+        # warm up with different number of warmup timestep to be more precise
+        for i in tqdm(range(self.warmup_timesteps), disable=not progress):
+            dt, t_partial, t_rf = self.calculate_prerolling_timestep(window_indexes, i)
+
+            score_v, score_a = model(v_z, a_z, t_rf, **model_kwargs)
+
+            v_z = v_z.detach().clone() + dt*score_v
+            a_z = a_z.detach().clone() + dt*score_a
+
+        v_f = v_z[:,0]
+        a_f = a_z[:,0]
+
+        v_z = torch.cat([v_z[:,1:], torch.randn_like(v_z[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+        a_z = torch.cat([a_z[:,1:], torch.randn_like(a_z[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+
+        def yield_frame():
+            nonlocal v_z, a_z, window_indexes
+            yield (v_f, a_f)
+
+            dt = 1/(self.num_timesteps*self.window_size)
+
+            while True:
+                for i in range(self.num_timesteps):
+                    tw = (self.num_timesteps - i)/self.num_timesteps
+                    t = (window_indexes + tw)/self.window_size
+                    t = 1-t
+
+                    t_rf = t*(self.T-self.eps) + self.eps
+
+                    score_v, score_a = model(v_z, a_z, t_rf, **model_kwargs)
+
+                    v_z = v_z.detach().clone() + dt*score_v
+                    a_z = a_z.detach().clone() + dt*score_a
+
+                v = v_z[:,0]
+                a = a_z[:,0]
+
+                #remove the first element
+                v_noise = torch.randn_like(v_z[:,0]).unsqueeze(1)*self.noise_scale
+                a_noise = torch.randn_like(a_z[:,0]).unsqueeze(1)*self.noise_scale
+
+                v_z = torch.cat([v_z[:,1:],v_noise], dim=1)
+                a_z = torch.cat([a_z[:,1:],a_noise], dim=1)
+
+                yield (v, a)
+        
+        return yield_frame
+    
+    def sample_a2v(self, model, v_z, a, model_kwargs, scale=1, progress=True):
+        B = v_z.shape[0]
+        window_indexes = torch.linspace(0, self.window_size-1, steps=self.window_size, device=v_z.device).unsqueeze(0).repeat(B,1)
+
+        a_partial = a[:, :self.window_size]
+
+        a_noise = torch.randn_like(a, device=v_z.device)*self.noise_scale
+        a_noise_partial = a_noise[:, :self.window_size]
+
+
+        with torch.enable_grad():
+            # warm up with different number of warmup timestep to be more precise
+            for i in tqdm(range(self.warmup_timesteps), disable=not progress):
+                v_z = v_z.detach().requires_grad_(True)
+
+                dt, t_partial, t_rf = self.calculate_prerolling_timestep(window_indexes, i)
+
+                a_z = a_partial*t_partial + a_noise_partial*(1-t_partial)
+
+                score_v, score_a = model(v_z, a_z, t_rf, **model_kwargs)
+
+                loss = torch.square((a_partial-a_noise_partial)-score_a)
+                grad = torch.autograd.grad(loss.mean(), v_z)[0]
+
+                v_z = v_z.detach() + dt*score_v - ((t_partial+dt)!=1) * dt * grad * scale
+
+        v_f = v_z[:,0].detach()
+        v_z = torch.cat([v_z[:,1:], torch.randn_like(v_z[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+
+
+        def yield_frame():
+            nonlocal v_z, a, a_noise, window_indexes
+            yield v_f
+
+            dt = 1/(self.num_timesteps*self.window_size)
+
+            while True:
+                torch.cuda.empty_cache()
+
+                a = a[:,1:]
+                a_noise = a_noise[:,1:]
+                if a.shape[1] < self.window_size:
+                    a = torch.cat([a, torch.randn_like(a[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+                    a_noise = torch.cat([a_noise, torch.randn_like(a[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+
+                a_partial = a[:, :self.window_size]
+                a_noise_partial = a_noise[:, :self.window_size]
+
+                with torch.enable_grad():
+                    for i in range(self.num_timesteps):
+                        v_z = v_z.detach().requires_grad_(True)
+
+                        tw = (self.num_timesteps - i)/self.num_timesteps
+                        t = (window_indexes + tw)/self.window_size
+                        t = 1-t
+
+                        t_partial = t.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+                        t_rf = t*(self.T-self.eps) + self.eps
+
+                        a_z = a_partial*t_partial + torch.randn_like(a_partial, device=v_z.device)*self.noise_scale*(1-t_partial)
+
+                        score_v, score_a = model(v_z, a_z, t_rf, **model_kwargs)
+
+                        loss = torch.square((a_partial-a_noise_partial)-score_a)
+                        grad = torch.autograd.grad(loss.mean(), v_z)[0]
+
+                        v_z = v_z.detach() + dt*score_v - ((t_partial+dt)!=1) * dt * grad * scale
+
+                v = v_z[:,0].detach()
+
+                v_noise = torch.randn_like(v_z[:,0]).unsqueeze(1)*self.noise_scale
+                v_z = torch.cat([v_z[:,1:],v_noise], dim=1)
+                yield v
+
+        return yield_frame
+    
+    def sample_v2a(self, model, v, a_z, model_kwargs, scale=2, progress=True):
+        B = a_z.shape[0]
+        window_indexes = torch.linspace(0, self.window_size-1, steps=self.window_size, device=a_z.device).unsqueeze(0).repeat(B,1)
+
+        v_partial = v[:, :self.window_size]
+        v_noise = torch.randn_like(v, device=a_z.device)*self.noise_scale
+        v_noise_partial = v_noise[:, :self.window_size]
+
+        with torch.enable_grad():
+            # warm up with different number of warmup timestep to be more precise
+            for i in tqdm(range(self.warmup_timesteps), disable=not progress):
+                a_z = a_z.detach().requires_grad_(True)
+
+                dt, t_partial, t_rf = self.calculate_prerolling_timestep(window_indexes, i)
+
+                v_z = v_partial*t_partial + v_noise_partial*(1-t_partial)
+
+                score_v, score_a = model(v_z, a_z, t_rf, **model_kwargs)
+
+                loss = torch.square((v_partial-v_noise_partial)-score_v)
+                grad = torch.autograd.grad(loss.mean(), a_z)[0]
+
+                a_z = a_z.detach() + dt*score_a - ((t_partial + dt)!=1) * dt * grad * scale
+
+        a_f = a_z[:,0].detach()
+        a_z = torch.cat([a_z[:,1:], torch.randn_like(a_z[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+
+        def yield_frame():
+            nonlocal v, v_noise, a_z, window_indexes
+            yield a_f
+
+            dt = 1/(self.num_timesteps*self.window_size)
+            while True:
+                torch.cuda.empty_cache()
+                v = v[:,1:]
+                v_noise = v_noise[:,1:]
+                
+                if v.shape[1] < self.window_size:
+                    v = torch.cat([v, torch.randn_like(v[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+                    v_noise = torch.cat([v, torch.randn_like(v[:,0]).unsqueeze(1)*self.noise_scale], dim=1)
+
+                v_partial = v[:, :self.window_size]
+                v_noise_partial = v_noise[:, :self.window_size]
+
+                with torch.enable_grad():
+                    for i in range(self.num_timesteps):
+                        a_z = a_z.detach().requires_grad_(True)
+
+                        tw = (self.num_timesteps - i)/self.num_timesteps
+                        t = (window_indexes + tw)/self.window_size
+                        t = 1-t
+
+                        t_partial = t.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+                        t_rf = t*(self.T-self.eps) + self.eps
+
+                        v_z = v_partial*t_partial + v_noise_partial*(1-t_partial)
+
+                        score_v, score_a = model(v_z, a_z, t_rf, **model_kwargs)
+
+                        loss = torch.square((v_partial-v_noise_partial)-score_v)
+                        grad = torch.autograd.grad(loss.mean(), a_z)[0]
+
+                        a_z = a_z.detach() + dt*score_a - ((t_partial + dt)!=1) * dt * grad * scale
+
+                a = a_z[:,0].detach()
+
+                a_noise = torch.randn_like(a_z[:,0]).unsqueeze(1)*self.noise_scale
+                a_z = torch.cat([a_z[:,1:],a_noise], dim=1)
+
+
+                yield a
+
+        return yield_frame
+
+    def calculate_prerolling_timestep(self, window_indexes, i):
+        tw = (self.warmup_timesteps - i)/self.warmup_timesteps
+        tw_future = (self.warmup_timesteps - (i+1))/self.warmup_timesteps
+
+        t = window_indexes/self.window_size + tw
+        
+        #timestep for the next iteration, to calculate dt
+        t_future = window_indexes/self.window_size + tw_future
+
+        #Swap 0 with 1, 1 is full image, 0 is full noise
+        t = 1-t
+        t_future = 1 - t_future
+
+        t = torch.clamp(t, 0, 1)
+        t_future = torch.clamp(t_future, 0, 1)
+
+        dt = torch.abs(t_future-t).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) # [B, window_size, 1, 1, 1]
+
+        t_partial = t.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        t_rf= t*(self.T-self.eps) + self.eps
+        return dt,t_partial,t_rf
+    
+    def get_z0(self, batch, train=True):
+
+        if self.init_type == 'gaussian':
+            ### standard gaussian #+ 0.5
+            return torch.randn(batch.shape)*self.noise_scale
+        else:
+            raise NotImplementedError("INITIALIZATION TYPE NOT IMPLEMENTED")