code/src/model_training/training_modules/model_logger.py

import json
import math
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,
        sampling_sigma_shift: float = 5.0,
    ):
        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)
        # Noise-schedule shift used by the periodic sampling monitor. Must match
        # the TRAINING shift (set in train.py from --timestep_shift), else the
        # monitor samples from a noise schedule the model never trained on ->
        # structure roughly OK but colors blown out / distorted.
        self.sampling_sigma_shift = float(sampling_sigma_shift)
        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",
                "geometry_spatial_memory_module",
                "geometry_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
        # Sampling runs the inference pipe.__call__, which (a) overwrites the
        # training scheduler with an inference schedule and (b) spikes GPU
        # memory on the main rank. Wrap it so a sampling failure can never abort
        # training, and always restore training scheduler state + free cache.
        try:
            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=self.sampling_sigma_shift,
                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)
        except Exception as e:
            # Sampling is a monitoring side-effect; never let it kill training.
            import traceback
            print(f"[paper-sampling][step={self.step_count}] FAILED (training continues): {e}")
            traceback.print_exc()
        finally:
            # Restore training scheduler + free the sampling step's GPU cache on
            # the main rank so the next training step (target VAE encode) is clean.
            _train_module = getattr(unwrapped, "restore_after_sampling", None)
            if callable(_train_module):
                _train_module()
            else:
                import gc
                gc.collect()
                if __import__("torch").cuda.is_available():
                    __import__("torch").cuda.empty_cache()

    def on_step_end(self, loss, accelerator=None, model=None, current_batch=None,
                    aux_loss=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())
                metrics = {"train/loss": loss_v, "step": self.step_count}
                # CGLA SSE load-balancing aux loss (summed per-block in the
                # training loop; passed in so dummy/skip steps don't log stale
                # values). Logged only when CGLA is active (cgla_aux_loss_weight>0).
                if aux_loss is not None and not math.isnan(aux_loss):
                    metrics["train/aux"] = float(aux_loss)
                self.wandb_logger.log(metrics)
        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()