pipeline/streaming_switch_training.py

from pipeline.streaming_training import StreamingTrainingPipeline
from typing import List, Optional, Tuple
import torch
import torch.distributed as dist

class StreamingSwitchTrainingPipeline(StreamingTrainingPipeline):

    def __init__(self, *args, **kwargs):
        _apr_enabled = kwargs.pop('apr_enabled', False)
        _apr_alpha_max = kwargs.pop('apr_alpha_max', 0.8)
        _apr_d_window = kwargs.pop('apr_d_window', None)
        _apr_blend_sink = kwargs.pop('apr_blend_sink', False)
        _global_sink = kwargs.pop('global_sink', False)
        super().__init__(*args, **kwargs)
        self.global_sink = _global_sink
        self.apr_enabled = _apr_enabled
        self.apr_alpha_max = float(max(0.0, min(0.8, _apr_alpha_max)))
        self.apr_d_window = _apr_d_window
        self.apr_blend_sink = _apr_blend_sink
        if not dist.is_initialized() or dist.get_rank() == 0:
            print(f'[StreamingSwitchTrainingPipeline] global_sink={self.global_sink} (config-driven, Opt 14 kwargs.pop fix active)')

    def generate_chunk_with_cache(self, noise: torch.Tensor, conditional_dict: dict, *, current_start_frame: int=0, requires_grad: bool=True, switch_frame_index: Optional[int]=None, switch_conditional_dict: Optional[dict]=None, switch_recache_frames: Optional[torch.Tensor]=None, return_sim_step: bool=False) -> Tuple[torch.Tensor, Optional[int], Optional[int]]:
        if switch_conditional_dict is None or switch_frame_index is None:
            return super().generate_chunk_with_cache(noise=noise, conditional_dict=conditional_dict, current_start_frame=current_start_frame, requires_grad=requires_grad, return_sim_step=return_sim_step)
        batch_size, chunk_frames, num_channels, height, width = noise.shape
        assert chunk_frames % self.num_frame_per_block == 0
        num_blocks = chunk_frames // self.num_frame_per_block
        all_num_frames = [self.num_frame_per_block] * num_blocks
        output = torch.zeros_like(noise)
        num_denoising_steps = len(self.denoising_step_list)
        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
        if not requires_grad:
            start_gradient_frame_index = chunk_frames
        else:
            start_gradient_frame_index = switch_frame_index
        local_start_frame = 0
        self.generator.model.local_attn_size = int(self.local_attn_size)
        self._set_all_modules_max_attention_size(int(self.local_attn_size))
        using_second = False
        cond_in_use = conditional_dict
        for block_index, current_num_frames in enumerate(all_num_frames):
            if not using_second and local_start_frame >= switch_frame_index:
                self._recache_after_switch(output[:, :local_start_frame, ...], current_start_frame + local_start_frame, switch_conditional_dict, local_start_frame, switch_recache_frames)
                cond_in_use = switch_conditional_dict
                using_second = True
            noisy_input = noise[:, local_start_frame:local_start_frame + current_num_frames]
            for step_idx, current_timestep in enumerate(self.denoising_step_list):
                exit_flag = step_idx == exit_flags[0] if self.same_step_across_blocks else step_idx == exit_flags[block_index]
                timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep
                if not exit_flag:
                    with torch.no_grad():
                        _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame + local_start_frame) * self.frame_seq_length)
                        if step_idx < len(self.denoising_step_list) - 1:
                            next_timestep = self.denoising_step_list[step_idx + 1]
                            noisy_input = self.scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)).unflatten(0, denoised_pred.shape[:2])
                else:
                    enable_grad = local_start_frame >= start_gradient_frame_index
                    context_manager = torch.enable_grad() if enable_grad else torch.no_grad()
                    with context_manager:
                        _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame + local_start_frame) * self.frame_seq_length)
                    break
            output[:, local_start_frame:local_start_frame + current_num_frames] = denoised_pred
            context_timestep = torch.ones_like(timestep) * self.context_noise
            context_noisy = self.scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), context_timestep.flatten(0, 1)).unflatten(0, denoised_pred.shape[:2])
            with torch.no_grad():
                self.generator(noisy_image_or_video=context_noisy, conditional_dict=cond_in_use, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame + local_start_frame) * self.frame_seq_length)
            local_start_frame += current_num_frames
        if not self.same_step_across_blocks:
            denoised_timestep_from, denoised_timestep_to = (None, None)
        elif exit_flags[0] == len(self.denoising_step_list) - 1:
            denoised_timestep_to = 0
            denoised_timestep_from = 1000 - torch.argmin((self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
        else:
            denoised_timestep_to = 1000 - torch.argmin((self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
            denoised_timestep_from = 1000 - torch.argmin((self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
        if return_sim_step:
            return (output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1)
        return (output, denoised_timestep_from, denoised_timestep_to)

    def _recache_after_switch(self, output, current_start_frame, new_conditional_dict, local_start_frame=None, switch_recache_frames=None):
        sink_tok_size = None
        sink_backup = None
        sink_norm_before = None
        if self.global_sink and current_start_frame > 0:
            gen = self.generator
            from torch.distributed.fsdp import FullyShardedDataParallel as _FSDP
            if hasattr(gen, '_fsdp_wrapped_module'):
                w = gen._fsdp_wrapped_module
                m = w.model
                _inner_tmp = m._fsdp_wrapped_module if isinstance(m, _FSDP) else m
                if hasattr(_inner_tmp, 'base_model') and hasattr(_inner_tmp.base_model, 'model'):
                    _inner_tmp = _inner_tmp.base_model.model
            elif hasattr(gen, 'model'):
                _inner_tmp = gen.model
            else:
                _inner_tmp = None
            if _inner_tmp is not None:
                sink_tok_size = _inner_tmp.blocks[0].self_attn.sink_size * self.frame_seq_length
                sink_backup = [{'k': self.kv_cache1[i]['k'][:, :sink_tok_size].clone(), 'v': self.kv_cache1[i]['v'][:, :sink_tok_size].clone()} for i in range(self.num_transformer_blocks)]
                sink_norm_before = torch.norm(self.kv_cache1[0]['k'][:, :sink_tok_size]).float().item()
                if not dist.is_initialized() or dist.get_rank() == 0:
                    print(f'[Recache-Train] switch@frame={current_start_frame}, global_sink=True, sink_tok={sink_tok_size}, sink_norm(block0) BEFORE = {sink_norm_before:.4f}')
        elif not self.global_sink:
            if not dist.is_initialized() or dist.get_rank() == 0:
                print(f'[Recache-Train] switch@frame={current_start_frame}, global_sink=False (sink will be recached with new prompt)')
        apr_old_cache = None
        if self.apr_enabled and current_start_frame > 0:
            apr_old_cache = [{'k': self.kv_cache1[i]['k'].clone(), 'v': self.kv_cache1[i]['v'].clone()} for i in range(self.num_transformer_blocks)]
        if not self.global_sink:
            for block_idx in range(self.num_transformer_blocks):
                cache = self.kv_cache1[block_idx]
                cache['k'].zero_()
                cache['v'].zero_()
        for blk in self.crossattn_cache:
            blk['k'].zero_()
            blk['v'].zero_()
            blk['is_init'] = False
        gen = self.generator
        from torch.distributed.fsdp import FullyShardedDataParallel as _FSDP
        if hasattr(gen, '_fsdp_wrapped_module'):
            w = gen._fsdp_wrapped_module
            m = w.model
            inner = m._fsdp_wrapped_module if isinstance(m, _FSDP) else m
            if hasattr(inner, 'base_model') and hasattr(inner.base_model, 'model'):
                inner = inner.base_model.model
        elif hasattr(gen, 'model'):
            inner = gen.model
        else:
            inner = None
        if current_start_frame == 0:
            return
        if switch_recache_frames is not None:
            frames_to_recache = torch.cat([switch_recache_frames, output], dim=1)[:, -21:, ...]
            num_recache_frames = frames_to_recache.shape[1]
        elif local_start_frame is not None:
            num_recache_frames = min(local_start_frame, 21)
            frames_to_recache = output[:, -num_recache_frames:]
        else:
            num_recache_frames = min(current_start_frame, 21)
            frames_to_recache = output[:, -num_recache_frames:]
        batch_size, num_recache_frames, c, h, w = frames_to_recache.shape
        device = frames_to_recache.device
        block_mask = self.generator.model._prepare_blockwise_causal_attn_mask(device=device, num_frames=num_recache_frames, frame_seqlen=self.frame_seq_length, num_frame_per_block=self.num_frame_per_block, local_attn_size=21)
        context_timestep = torch.ones([batch_size, num_recache_frames], device=device, dtype=torch.int64) * self.context_noise
        self.generator.model.block_mask = block_mask
        _has_memory = inner is not None and (getattr(inner, 'query_memory_encoder', None) is not None or getattr(inner, 'sink_memory', None) is not None)
        if _has_memory:
            frame_seqlen = self.frame_seq_length
            sink_frames = inner.blocks[0].self_attn.sink_size
            max_attn = inner.blocks[0].self_attn.max_attention_size
            recent_window_frames = (max_attn - sink_frames * frame_seqlen) // frame_seqlen
            inner._ei_prev_window_start = max(sink_frames, current_start_frame - recent_window_frames)
        with torch.no_grad():
            self.generator(noisy_image_or_video=frames_to_recache, conditional_dict=new_conditional_dict, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame - num_recache_frames) * self.frame_seq_length)
        if sink_backup is not None:
            for i in range(self.num_transformer_blocks):
                self.kv_cache1[i]['k'][:, :sink_tok_size] = sink_backup[i]['k']
                self.kv_cache1[i]['v'][:, :sink_tok_size] = sink_backup[i]['v']
            sink_norm_after = torch.norm(self.kv_cache1[0]['k'][:, :sink_tok_size]).float().item()
            if not dist.is_initialized() or dist.get_rank() == 0:
                delta = abs(sink_norm_after - sink_norm_before)
                status = 'PRESERVED ✓' if delta < 0.001 else f'MISMATCH ✗ (delta={delta:.4e})'
                print(f'[Recache-Train] global_sink=True: sink_norm AFTER = {sink_norm_after:.4f} [{status}]')
                try:
                    import model.streaming_training as _st_mod
                    setattr(_st_mod, '_last_recache_sink_delta', delta)
                except Exception:
                    pass
            del sink_backup
        for blk in self.crossattn_cache:
            blk['k'].zero_()
            blk['v'].zero_()
            blk['is_init'] = False
        if self.apr_enabled and apr_old_cache is not None:
            self._apply_apr_blend(old_cache=apr_old_cache, recache_start_frame=current_start_frame - num_recache_frames, current_start_frame=current_start_frame, num_recache_frames=num_recache_frames, inner=inner)
            del apr_old_cache
        if inner is not None and getattr(inner, 'query_memory_encoder', None) is not None:
            enc = inner.query_memory_encoder
            if getattr(enc, 'memory_recache', False):
                enc.reset(batch_size=frames_to_recache.shape[0], device=frames_to_recache.device, dtype=torch.bfloat16)
                last_blk = self.num_transformer_blocks - 1
                cache = self.kv_cache1[last_blk]
                frame_seqlen = self.frame_seq_length
                sink_tok = inner.blocks[0].self_attn.sink_size * frame_seqlen
                local_end = cache['local_end_index'].item()
                if local_end > sink_tok:
                    recache_k = cache['k'][:, sink_tok:local_end].clone()
                    recache_v = cache['v'][:, sink_tok:local_end].clone()
                    sink_k = cache['k'][:, :sink_tok].clone() if sink_tok > 0 else None
                    sink_v = cache['v'][:, :sink_tok].clone() if sink_tok > 0 else None
                    enc.update(recache_k, recache_v, sink_k, sink_v)
        if inner is not None and getattr(inner, 'sink_memory', None) is not None:
            sm = inner.sink_memory
            sm.reset()
            sink_frames_count = inner.blocks[0].self_attn.sink_size
            sink_output = output[:, :sink_frames_count] if output.shape[1] >= sink_frames_count else None
            if sink_output is None:
                if switch_recache_frames is not None and switch_recache_frames.shape[1] >= sink_frames_count:
                    sink_output = switch_recache_frames[:, :sink_frames_count]
            if sink_output is not None:
                device = sink_output.device
                batch_size_local = sink_output.shape[0]
                sink_ts = torch.ones([batch_size_local, sink_frames_count], device=device, dtype=torch.int64) * self.context_noise
                n_heads = inner.blocks[0].self_attn.num_heads
                head_dim = inner.blocks[0].self_attn.head_dim
                sink_cache_size = sink_frames_count * self.frame_seq_length
                temp_kv = [{'k': torch.zeros([batch_size_local, sink_cache_size, n_heads, head_dim], device=device, dtype=torch.bfloat16), 'v': torch.zeros([batch_size_local, sink_cache_size, n_heads, head_dim], device=device, dtype=torch.bfloat16), 'global_end_index': torch.zeros(1, dtype=torch.long, device=device), 'local_end_index': torch.zeros(1, dtype=torch.long, device=device)} for _ in range(self.num_transformer_blocks)]
                temp_crossattn = [{'k': torch.zeros_like(self.crossattn_cache[0]['k']), 'v': torch.zeros_like(self.crossattn_cache[0]['v']), 'is_init': False} for _ in range(self.num_transformer_blocks)]
                with torch.no_grad():
                    self.generator(noisy_image_or_video=sink_output, conditional_dict=new_conditional_dict, timestep=sink_ts, kv_cache=temp_kv, crossattn_cache=temp_crossattn, current_start=0)
                frame_seqlen_local = self.frame_seq_length
                sink_frames_local = inner.blocks[0].self_attn.sink_size
                max_attn_local = inner.blocks[0].self_attn.max_attention_size
                recent_win = (max_attn_local - sink_frames_local * frame_seqlen_local) // frame_seqlen_local
                inner._ei_prev_window_start = max(sink_frames_local, current_start_frame - recent_win)
            if not dist.is_initialized() or dist.get_rank() == 0:
                print(f'[SinkMem] reset + re-captured on prompt switch at frame {current_start_frame}')

    def _apply_apr_blend(self, *, old_cache, recache_start_frame, current_start_frame, num_recache_frames, inner):
        frame_seqlen = self.frame_seq_length
        sink_size = inner.blocks[0].self_attn.sink_size
        sink_tok = sink_size * frame_seqlen
        D = self.apr_d_window if self.apr_d_window is not None else num_recache_frames
        D = max(1, int(D))
        alpha_max = self.apr_alpha_max
        n_frames = num_recache_frames
        with torch.no_grad():
            for blk_idx in range(self.num_transformer_blocks):
                new_k = self.kv_cache1[blk_idx]['k']
                new_v = self.kv_cache1[blk_idx]['v']
                old_k = old_cache[blk_idx]['k']
                old_v = old_cache[blk_idx]['v']
                if self.apr_blend_sink and (not self.global_sink):
                    for f in range(sink_size):
                        d_t = current_start_frame - f
                        alpha = min(alpha_max, max(0.0, 1.0 - d_t / D))
                        s0, s1 = (f * frame_seqlen, (f + 1) * frame_seqlen)
                        new_k[:, s0:s1] = (1 - alpha) * old_k[:, s0:s1] + alpha * new_k[:, s0:s1]
                        new_v[:, s0:s1] = (1 - alpha) * old_v[:, s0:s1] + alpha * new_v[:, s0:s1]
                base = sink_tok
                for f in range(n_frames):
                    absolute_frame = recache_start_frame + f
                    d_t = current_start_frame - absolute_frame
                    alpha = min(alpha_max, max(0.0, 1.0 - d_t / D))
                    s0, s1 = (base + f * frame_seqlen, base + (f + 1) * frame_seqlen)
                    new_k[:, s0:s1] = (1 - alpha) * old_k[:, s0:s1] + alpha * new_k[:, s0:s1]
                    new_v[:, s0:s1] = (1 - alpha) * old_v[:, s0:s1] + alpha * new_v[:, s0:s1]
        if not dist.is_initialized() or dist.get_rank() == 0:
            print(f'[APR] switch@{current_start_frame}: blended {n_frames} frames (α_max={alpha_max}, D_window={D})')