model/streaming_training.py

import time
import torch
import torch.distributed as dist
from typing import Tuple, Dict, Any, Optional, List
from einops import rearrange
from pipeline.streaming_switch_training import StreamingSwitchTrainingPipeline

class StreamingTrainingModel:

    def __init__(self, base_model, config):
        self.base_model = base_model
        self.config = config
        self.device = base_model.device
        self.dtype = base_model.dtype
        self.image_or_video_shape = getattr(config, 'image_or_video_shape', None)
        self.chunk_size = getattr(config, 'streaming_chunk_size', 21)
        self.max_length = getattr(config, 'streaming_max_length', 57)
        self.possible_max_length = getattr(config, 'streaming_possible_max_length', None)
        self.min_new_frame = getattr(config, 'streaming_min_new_frame', 18)
        self.generator = base_model.generator
        self.fake_score = base_model.fake_score
        self.scheduler = base_model.scheduler
        self.denoising_loss_func = base_model.denoising_loss_func
        self.num_frame_per_block = base_model.num_frame_per_block
        self.frame_seq_length = getattr(base_model.inference_pipeline, 'frame_seq_length', 1560)
        self.inference_pipeline = base_model.inference_pipeline
        if self.inference_pipeline is None:
            base_model._initialize_inference_pipeline()
            self.inference_pipeline = base_model.inference_pipeline
        self.reset_state()

    def _process_first_frame_encoding(self, frames: torch.Tensor) -> torch.Tensor:
        total_frames = frames.shape[1]
        if total_frames <= 1:
            return frames
        process_frames = min(21, total_frames)
        with torch.no_grad():
            frames_to_decode = frames[:, :-(process_frames - 1), ...]
            pixels = self.base_model.vae.decode_to_pixel(frames_to_decode)
            last_frame_pixel = pixels[:, -1:, ...].to(self.dtype)
            last_frame_pixel = rearrange(last_frame_pixel, 'b t c h w -> b c t h w')
            image_latent = self.base_model.vae.encode_to_latent(last_frame_pixel).to(self.dtype)
        remaining_frames = frames[:, -(process_frames - 1):, ...]
        processed_frames = torch.cat([image_latent, remaining_frames], dim=1)
        return processed_frames

    def reset_state(self):
        self.state = {'current_length': 0, 'conditional_info': None, 'has_switched': False, 'previous_frames': None, 'temp_max_length': None, '_ei_chunk_count': 0, 'text_prompts': None, 'switch_text_prompts': None}
        self.inference_pipeline.clear_kv_cache()
        gen = self.inference_pipeline.generator
        from torch.distributed.fsdp import FullyShardedDataParallel as _FSDP
        if hasattr(gen, '_fsdp_wrapped_module'):
            wrapper = gen._fsdp_wrapped_module
            m = wrapper.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 inner is not None and getattr(inner, 'query_memory_encoder', None) is not None:
            inner.query_memory_encoder.reset(batch_size=self.image_or_video_shape[0] if self.image_or_video_shape else 1, device=self.device, dtype=self.dtype)
            inner._ei_prev_window_start = None

    def _should_switch_prompt(self, chunk_start_frame: int, chunk_size: int) -> bool:
        from pipeline.streaming_switch_training import StreamingSwitchTrainingPipeline
        if not isinstance(self.inference_pipeline, StreamingSwitchTrainingPipeline):
            return False
        if self.state.get('has_switched', False):
            return False
        switch_info = self.state['conditional_info'].get('switch_info', {})
        switch_frame_index = switch_info.get('switch_frame_index')
        if switch_frame_index is None:
            return False
        chunk_end_frame = chunk_start_frame + chunk_size
        should_switch = chunk_start_frame <= switch_frame_index < chunk_end_frame
        return should_switch

    def _get_current_conditional_dict(self, chunk_start_frame: int) -> dict:
        cond_info = self.state['conditional_info']
        switch_info = cond_info.get('switch_info', {})
        if switch_info:
            switch_frame_index = switch_info.get('switch_frame_index')
            if switch_frame_index is not None:
                if self.state.get('has_switched', False) or chunk_start_frame >= switch_frame_index:
                    return switch_info.get('switch_conditional_dict', cond_info['conditional_dict'])
        return cond_info['conditional_dict']

    def _get_current_text_prompts(self, chunk_start_frame: int) -> Optional[list]:
        text_prompts = self.state.get('text_prompts')
        if text_prompts is None:
            return None
        cond_info = self.state['conditional_info']
        switch_info = cond_info.get('switch_info', {})
        if switch_info:
            switch_text_prompts = self.state.get('switch_text_prompts')
            if switch_text_prompts is None:
                raise RuntimeError('[StreamingTrain-Model] switch_info present (conditional_dict switches) but switch_text_prompts is None. setup_sequence must receive switch_text_prompts whenever switch_conditional_dict is provided -- no silent fallback.')
            switch_frame_index = switch_info.get('switch_frame_index')
            if switch_frame_index is not None:
                if self.state.get('has_switched', False) or chunk_start_frame >= switch_frame_index:
                    return switch_text_prompts
        return text_prompts

    def _generate_chunk(self, noise_chunk: torch.Tensor, chunk_start_frame: int, requires_grad: bool=True) -> Tuple[torch.Tensor, Optional[int], Optional[int]]:
        current_conditional_dict = self._get_current_conditional_dict(chunk_start_frame)
        kwargs = {'noise': noise_chunk, 'conditional_dict': current_conditional_dict, 'current_start_frame': chunk_start_frame, 'requires_grad': requires_grad, 'return_sim_step': False}
        if isinstance(self.inference_pipeline, StreamingSwitchTrainingPipeline):
            switch_info = self.state['conditional_info'].get('switch_info', {})
            if switch_info and self._should_switch_prompt(chunk_start_frame, noise_chunk.shape[1]):
                if not dist.is_initialized() or dist.get_rank() == 0:
                    print(f"[StreamingTrain-Model] Switching prompt at frame {switch_info['switch_frame_index']}")
                relative_switch_index = max(0, switch_info['switch_frame_index'] - chunk_start_frame)
                kwargs['switch_frame_index'] = relative_switch_index
                kwargs['switch_conditional_dict'] = switch_info['switch_conditional_dict']
                if self.state['previous_frames'] is not None:
                    kwargs['switch_recache_frames'] = self.state['previous_frames']
                self.state['has_switched'] = True
        output, denoised_timestep_from, denoised_timestep_to = self.inference_pipeline.generate_chunk_with_cache(**kwargs)
        return (output, denoised_timestep_from, denoised_timestep_to)

    def setup_sequence(self, conditional_dict: Dict, unconditional_dict: Dict, initial_latent: Optional[torch.Tensor]=None, switch_conditional_dict: Optional[Dict]=None, switch_frame_index: Optional[int]=None, temp_max_length: Optional[int]=None, text_prompts: Optional[List]=None, switch_text_prompts: Optional[List]=None):
        from utils.debug_option import maybe_empty_cache
        maybe_empty_cache()
        batch_size = self.image_or_video_shape[0]
        if self.inference_pipeline.kv_cache1 is None:
            self.inference_pipeline._initialize_kv_cache(batch_size=batch_size, dtype=self.dtype, device=self.device)
        if self.inference_pipeline.crossattn_cache is None:
            self.inference_pipeline._initialize_crossattn_cache(batch_size=batch_size, dtype=self.dtype, device=self.device)
        self.reset_state()
        self.state['temp_max_length'] = temp_max_length
        self.state['text_prompts'] = text_prompts
        self.state['switch_text_prompts'] = switch_text_prompts
        if initial_latent is not None:
            self.state['current_length'] = initial_latent.shape[1]
        else:
            self.state['current_length'] = 0
        self.state['conditional_info'] = {'conditional_dict': conditional_dict, 'unconditional_dict': unconditional_dict}
        if switch_conditional_dict is not None and switch_frame_index is not None:
            self.state['conditional_info']['switch_info'] = {'switch_conditional_dict': switch_conditional_dict, 'switch_frame_index': switch_frame_index}
        if initial_latent is not None:
            timestep = torch.zeros([batch_size, initial_latent.shape[1]], device=self.device, dtype=torch.int64)
            with torch.no_grad():
                self.inference_pipeline.generator(noisy_image_or_video=initial_latent, conditional_dict=conditional_dict, timestep=timestep, kv_cache=self.inference_pipeline.kv_cache1, crossattn_cache=self.inference_pipeline.crossattn_cache, current_start=0)

    def can_generate_more(self) -> bool:
        current_length = self.state['current_length']
        temp_max_length = self.state.get('temp_max_length')
        can_generate = current_length < temp_max_length and current_length + self.min_new_frame <= temp_max_length
        return can_generate

    def generate_next_chunk(self, requires_grad: bool=True) -> Tuple[torch.Tensor, Dict[str, Any]]:
        if not self.can_generate_more():
            raise ValueError('Cannot generate more chunks')
        current_length = self.state['current_length']
        batch_size = self.image_or_video_shape[0]
        previous_frames = self.state.get('previous_frames')
        if previous_frames is not None:
            max_new_frames = min(self.state['temp_max_length'] - current_length + 1, self.chunk_size)
            possible_new_frames = list(range(self.min_new_frame, max_new_frames, 3))
            if dist.is_initialized():
                if dist.get_rank() == 0:
                    import random
                    selected_idx = random.randint(0, len(possible_new_frames) - 1)
                else:
                    selected_idx = 0
                selected_idx_tensor = torch.tensor(selected_idx, device=self.device, dtype=torch.int32)
                dist.broadcast(selected_idx_tensor, src=0)
                selected_idx = selected_idx_tensor.item()
            else:
                import random
                selected_idx = random.randint(0, len(possible_new_frames) - 1)
            new_frames_to_generate = possible_new_frames[selected_idx]
            overlap_frames = self.chunk_size - new_frames_to_generate
            if overlap_frames > 0 and overlap_frames <= previous_frames.shape[1]:
                overlap_frames_to_use = overlap_frames
            else:
                overlap_frames_to_use = 0
                new_frames_to_generate = self.chunk_size
        else:
            overlap_frames_to_use = 0
            new_frames_to_generate = self.chunk_size
        noise_chunk = torch.randn([batch_size, new_frames_to_generate, *self.image_or_video_shape[2:]], device=self.device, dtype=self.dtype)
        generated_new_frames, denoised_timestep_from, denoised_timestep_to = self._generate_chunk(noise_chunk=noise_chunk, chunk_start_frame=current_length, requires_grad=requires_grad)
        if previous_frames is not None:
            full_chunk = torch.cat([previous_frames, generated_new_frames], dim=1)
        else:
            full_chunk = generated_new_frames
        frames_to_save = full_chunk.detach().clone()[:, -self.chunk_size:, ...]
        if previous_frames is not None:
            full_chunk = self._process_first_frame_encoding(full_chunk)
        if previous_frames is not None:
            gradient_mask = torch.zeros_like(full_chunk, dtype=torch.bool)
            gradient_mask[:, overlap_frames_to_use:overlap_frames_to_use + new_frames_to_generate, ...] = True
        else:
            gradient_mask = torch.ones_like(full_chunk, dtype=torch.bool)
        self.state['current_length'] += new_frames_to_generate
        self.state['previous_frames'] = frames_to_save
        gen = self.inference_pipeline.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
        encoder = getattr(inner, 'query_memory_encoder', None) if inner else None
        if encoder is not None:
            chunk_count = self.state.get('_ei_chunk_count', 0) + 1
            self.state['_ei_chunk_count'] = chunk_count
            if chunk_count % encoder.bptt_clips == 0:
                encoder.detach_state()
        info = {'denoised_timestep_from': denoised_timestep_from, 'denoised_timestep_to': denoised_timestep_to, 'chunk_start_frame': current_length, 'chunk_frames': full_chunk.shape[1], 'new_frames_generated': new_frames_to_generate, 'current_length': self.state['current_length'], 'gradient_mask': gradient_mask, 'overlap_frames_used': overlap_frames_to_use}
        if not dist.is_initialized() or dist.get_rank() == 0:
            print(f"[StreamingTrain-Model] current_training_chunk: ({self.state['current_length'] - new_frames_to_generate} -> {self.state['current_length']})/{self.state['temp_max_length']}")
        return (full_chunk, info)

    def compute_generator_loss(self, chunk: torch.Tensor, chunk_info: Dict[str, Any]) -> Tuple[torch.Tensor, Dict[str, Any]]:
        _t_loss_start = time.time()
        chunk_start_frame = chunk_info['chunk_start_frame']
        conditional_dict = self._get_current_conditional_dict(chunk_start_frame)
        unconditional_dict = self.state['conditional_info']['unconditional_dict']
        gradient_mask = chunk_info.get('gradient_mask', None)
        text_prompts = self._get_current_text_prompts(chunk_start_frame)
        dmd_loss, dmd_log_dict = self.base_model.compute_distribution_matching_loss(image_or_video=chunk, conditional_dict=conditional_dict, unconditional_dict=unconditional_dict, gradient_mask=gradient_mask, denoised_timestep_from=chunk_info['denoised_timestep_from'], denoised_timestep_to=chunk_info['denoised_timestep_to'], text_prompts=text_prompts)
        dmd_log_dict.update({'loss_time': time.time() - _t_loss_start, 'new_frames_supervised': chunk_info.get('new_frames_generated', chunk.shape[1])})
        return (dmd_loss, dmd_log_dict)

    def _clear_cache_gradients(self):
        if hasattr(self.inference_pipeline, 'kv_cache1') and self.inference_pipeline.kv_cache1 is not None:
            for cache_block in self.inference_pipeline.kv_cache1:
                if 'k' in cache_block and cache_block['k'].requires_grad:
                    cache_block['k'] = cache_block['k'].detach()
                if 'v' in cache_block and cache_block['v'].requires_grad:
                    cache_block['v'] = cache_block['v'].detach()
        if hasattr(self.inference_pipeline, 'crossattn_cache') and self.inference_pipeline.crossattn_cache is not None:
            for cache_block in self.inference_pipeline.crossattn_cache:
                if 'k' in cache_block and cache_block['k'].requires_grad:
                    cache_block['k'] = cache_block['k'].detach()
                if 'v' in cache_block and cache_block['v'].requires_grad:
                    cache_block['v'] = cache_block['v'].detach()

    def compute_critic_loss(self, chunk: torch.Tensor, chunk_info: Dict[str, Any]) -> Tuple[torch.Tensor, Dict[str, Any]]:
        _t_loss_start = time.time()
        if chunk.requires_grad:
            chunk = chunk.detach()
        self._clear_cache_gradients()
        from utils.debug_option import maybe_empty_cache
        maybe_empty_cache()
        chunk_start_frame = chunk_info['chunk_start_frame']
        conditional_dict = self._get_current_conditional_dict(chunk_start_frame)
        gradient_mask = chunk_info.get('gradient_mask', None)
        batch_size, num_frame = chunk.shape[:2]
        denoised_timestep_from = chunk_info.get('denoised_timestep_from', None)
        denoised_timestep_to = chunk_info.get('denoised_timestep_to', None)
        min_timestep = denoised_timestep_to if getattr(self.base_model, 'ts_schedule', False) and denoised_timestep_to is not None else getattr(self.base_model, 'min_score_timestep')
        max_timestep = denoised_timestep_from if getattr(self.base_model, 'ts_schedule_max', False) and denoised_timestep_from is not None else getattr(self.base_model, 'num_train_timestep')
        critic_timestep = self.base_model._get_timestep(min_timestep=min_timestep, max_timestep=max_timestep, batch_size=batch_size, num_frame=num_frame, num_frame_per_block=getattr(self.base_model, 'num_frame_per_block', 3), uniform_timestep=True).to(self.device)
        if getattr(self.base_model, 'timestep_shift') > 1:
            timestep_shift = self.base_model.timestep_shift
            critic_timestep = timestep_shift * (critic_timestep / 1000) / (1 + (timestep_shift - 1) * (critic_timestep / 1000)) * 1000
        critic_timestep = critic_timestep.clamp(self.base_model.min_step, self.base_model.max_step)
        critic_noise = torch.randn_like(chunk)
        noisy_chunk = self.scheduler.add_noise(chunk.flatten(0, 1), critic_noise.flatten(0, 1), critic_timestep.flatten(0, 1)).unflatten(0, (batch_size, num_frame))
        _, pred_fake_image = self.fake_score(noisy_image_or_video=noisy_chunk, conditional_dict=conditional_dict, timestep=critic_timestep)
        denoising_loss_type = getattr(self.base_model.args, 'denoising_loss_type', 'mse')
        if denoising_loss_type == 'flow':
            from utils.wan_wrapper import WanDiffusionWrapper
            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(scheduler=self.scheduler, x0_pred=pred_fake_image.flatten(0, 1), xt=noisy_chunk.flatten(0, 1), timestep=critic_timestep.flatten(0, 1))
            pred_fake_noise = None
        else:
            flow_pred = None
            pred_fake_noise = self.scheduler.convert_x0_to_noise(x0=pred_fake_image.flatten(0, 1), xt=noisy_chunk.flatten(0, 1), timestep=critic_timestep.flatten(0, 1)).unflatten(0, (batch_size, num_frame))
        gradient_mask_flat = gradient_mask.flatten(0, 1) if gradient_mask is not None else None
        denoising_loss = self.denoising_loss_func(x=chunk.flatten(0, 1), x_pred=pred_fake_image.flatten(0, 1), noise=critic_noise.flatten(0, 1), noise_pred=pred_fake_noise, alphas_cumprod=self.scheduler.alphas_cumprod, timestep=critic_timestep.flatten(0, 1), flow_pred=flow_pred, gradient_mask=gradient_mask_flat)
        del conditional_dict, critic_noise, noisy_chunk, pred_fake_image
        if 'flow_pred' in locals():
            del flow_pred
        if 'pred_fake_noise' in locals():
            del pred_fake_noise
        critic_log_dict = {'loss_time': time.time() - _t_loss_start, 'new_frames_supervised': chunk_info.get('new_frames_generated', num_frame)}
        return (denoising_loss, critic_log_dict)

    def get_sequence_length(self) -> int:
        return self.state.get('current_length', 0)