import json import os from typing import Optional import wandb from src.model_training.transformers_compat import patch_transformers_hybrid_cache patch_transformers_hybrid_cache() from diffsynth.trainers.utils import ModelLogger as BaseModelLogger class ModelLogger(BaseModelLogger): """Compatibility wrapper for legacy training scripts.""" def __init__( self, output_path, remove_prefix_in_ckpt=None, state_dict_converter=lambda x: x, wandb_run_name=None, ckpt_interval=None, resume_step_count=0, save_full_model=False, context_drop_prob: float = 0.0, enable_video_sampling=False, sampling_interval_steps: int = 0, sampling_two_chunk_memory: bool = False, sampling_action_path: Optional[str] = None, sampling_two_chunk_action_path: Optional[str] = None, sampling_negative_prompt: str = "oversaturated colors, overexposed, static, blurry details", sampling_height: int = 352, sampling_width: int = 640, sampling_num_frames: int = 81, sampling_num_inference_steps: int = 50, context_memory_frames: int = 1, context_source: str = "replay", context_per_frame_vae: bool = False, ): super().__init__(output_path, remove_prefix_in_ckpt=remove_prefix_in_ckpt, state_dict_converter=state_dict_converter) self.wandb_run_name = wandb_run_name self.ckpt_interval = int(ckpt_interval) if ckpt_interval else None self.step_count = int(resume_step_count) self.save_full_model = bool(save_full_model) self.total_steps = None self.context_drop_prob = float(context_drop_prob) self.enable_video_sampling = bool(enable_video_sampling) self.sampling_interval_steps = int(sampling_interval_steps) self.sampling_two_chunk_memory = bool(sampling_two_chunk_memory) self.sampling_action_path = sampling_action_path self.sampling_two_chunk_action_path = sampling_two_chunk_action_path self.sampling_negative_prompt = sampling_negative_prompt self.sampling_height = int(sampling_height) self.sampling_width = int(sampling_width) self.sampling_num_frames = int(sampling_num_frames) self.sampling_num_inference_steps = int(sampling_num_inference_steps) self.context_memory_frames = int(context_memory_frames) self.context_source = context_source.strip().lower() self.context_per_frame_vae = bool(context_per_frame_vae) self.wandb_logger = None if self.wandb_run_name: self.wandb_logger = wandb.init(project="wan-cam", name=self.wandb_run_name, reinit=True) def _save_step_or_epoch_ckpt(self, accelerator, model, path: str): state_dict = None unwrapped = accelerator.unwrap_model(model) if self.save_full_model: # Save full DiT (including action/camera/memory modules), not whole pipeline. state_dict = accelerator.get_state_dict(unwrapped.pipe.dit) for module_name in ("spatial_memory_module", "spatial_memory_readout_module"): module = getattr(unwrapped, module_name, None) if module is not None: state_dict.update( { f"{module_name}.{name}": param for name, param in accelerator.get_state_dict(module).items() } ) if state_dict is None: full_state = accelerator.get_state_dict(model) state_dict = unwrapped.export_trainable_state_dict(full_state, remove_prefix=self.remove_prefix_in_ckpt) state_dict = self.state_dict_converter(state_dict) os.makedirs(self.output_path, exist_ok=True) accelerator.save(state_dict, path, safe_serialization=True) def _maybe_sample_paper_process(self, accelerator=None, model=None, current_batch=None): if not ( self.enable_video_sampling and self.sampling_two_chunk_memory and self.sampling_interval_steps > 0 and self.step_count % self.sampling_interval_steps == 0 and accelerator is not None and model is not None and current_batch is not None ): return from diffsynth import save_video from src.model_training.multichunk_sample_utils import ( run_two_chunk_memory_monitor, sync_pipe_memory_from_training_module, ) sample = current_batch[0] if isinstance(current_batch, list) else current_batch first_frame = sample["video"][0] unwrapped = accelerator.unwrap_model(model) pipe = unwrapped.pipe sync_pipe_memory_from_training_module(pipe, unwrapped) action0 = self.sampling_two_chunk_action_path or self.sampling_action_path action1 = self.sampling_action_path frames0, frames1, meta = run_two_chunk_memory_monitor( pipe, prompt=sample.get("prompt") or sample.get("description") or "A scene.", negative_prompt=self.sampling_negative_prompt, action_path=self.sampling_action_path, chunk0_action_path=action0, chunk1_action_path=action1, first_frame_pil=first_frame, context_memory_frames=self.context_memory_frames, chunk_frames=self.sampling_num_frames, h=self.sampling_height, w=self.sampling_width, seed=42 + self.step_count + accelerator.process_index, sigma_shift=5.0, num_inference_steps=self.sampling_num_inference_steps, cfg_scale=5.0, inference_noise_level=0.0, omit_context_actions=False, context_source=self.context_source, context_position=os.environ.get("CONTEXT_POSITION", "suffix"), context_per_frame_vae=self.context_per_frame_vae, device=pipe.device, log_prefix=f"[paper-sampling][step={self.step_count}]", ) out_dir = os.path.join(self.output_path, "paper_process_sampling") os.makedirs(out_dir, exist_ok=True) tag = f"step_{self.step_count:07d}_rank{accelerator.process_index}" save_video(list(frames0) + list(frames1), os.path.join(out_dir, f"{tag}_pred.mp4"), fps=15, quality=5) with open(os.path.join(out_dir, f"{tag}_meta.json"), "w", encoding="utf-8") as f: json.dump(meta, f, ensure_ascii=False, indent=2) def on_step_end(self, loss, accelerator=None, model=None, current_batch=None): self.step_count += 1 if self.wandb_logger is not None: if accelerator is None or accelerator.is_main_process: loss_v = float(loss.detach().float().item()) self.wandb_logger.log({"train/loss": loss_v, "step": self.step_count}) if accelerator is not None and accelerator.is_main_process: self._maybe_sample_paper_process(accelerator, model, current_batch) if accelerator is not None and self.enable_video_sampling and self.sampling_two_chunk_memory and self.sampling_interval_steps > 0: accelerator.wait_for_everyone() if ( self.ckpt_interval and accelerator is not None and model is not None and (self.step_count % self.ckpt_interval == 0) ): accelerator.wait_for_everyone() if accelerator.is_main_process: path = os.path.join(self.output_path, f"Step-{self.step_count}.safetensors") self._save_step_or_epoch_ckpt(accelerator, model, path) def on_epoch_end(self, accelerator, model, epoch_id): accelerator.wait_for_everyone() if accelerator.is_main_process: path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors") self._save_step_or_epoch_ckpt(accelerator, model, path) def finish(self): if self.wandb_logger is not None: wandb.finish()