"""Public RDPO high-noise fine-tuning entry for SCoPE. Training reproduces the released recipe: only the high-noise expert is optimized, timesteps are sampled from ``[0.9, 1.0)``, and the data mixture is a ``ConcatDataset`` of the four native loaders in :mod:`scope.data`. The run is described by a single YAML config (see configs/train_rdpo_high_only.yaml). """ from __future__ import annotations import argparse from pathlib import Path from typing import Any import pytorch_lightning as pl import torch import yaml from pytorch_lightning.callbacks import ModelCheckpoint from torch.utils.data import DataLoader from torch.utils.data.dataloader import default_collate from scope.config import SCOPE_MODEL_ID, InferenceConfig from scope.data import build_training_dataset from scope.patch import enable_scope_grad from scope.weights import load_pipeline, resolve_model_dir _TRAINABLE_KEYWORDS = ["plucker_pe", "self_attn", "norm3", "ffn"] # PIL first frames cannot be default-collated, so they are kept as a list. _LIST_KEYS = ("first_frame_pil",) def load_config(path: Path) -> dict[str, Any]: config = yaml.safe_load(Path(path).read_text(encoding="utf-8")) if not isinstance(config, dict): raise ValueError(f"Training config must be a mapping: {path}") return config def _collate(batch: list[dict[str, Any]]) -> dict[str, Any]: output: dict[str, Any] = {} for key in batch[0]: values = [sample[key] for sample in batch] output[key] = values if key in _LIST_KEYS else default_collate(values) return output class SCoPEFineTuner(pl.LightningModule): """Fine-tune the released SCoPE model in the RDPO high-noise regime.""" def __init__( self, model_path: str, learning_rate: float, weight_decay: float, height: int, width: int, num_frames: int, ) -> None: super().__init__() self.save_hyperparameters() self.pipe = None def setup(self, stage: str | None = None) -> None: if self.pipe is not None: return inference_config = InferenceConfig( height=self.hparams.height, width=self.hparams.width, num_frames=self.hparams.num_frames, ) model_dir = resolve_model_dir(self.hparams.model_path) self.pipe = load_pipeline(model_dir, inference_config) self.pipe.i2v_vae_condition_mode = "official_zero" enable_scope_grad(self.pipe, _TRAINABLE_KEYWORDS, expert="high_noise_model") object.__setattr__(self, "dit", self.pipe.dit) object.__setattr__(self, "dit2", self.pipe.dit2) self.pipe.scheduler.set_timesteps( self.pipe.scheduler.num_train_timesteps, training=True, ) def training_step(self, batch: dict[str, Any], batch_index: int) -> torch.Tensor: del batch_index if len(batch["first_frame_pil"]) != 1: raise ValueError("SCoPE A14B training currently requires batch_size=1 per GPU") self.pipe.device = self.device self.pipe.load_models_to_device(["vae"]) video = batch["video"].to(device=self.device, dtype=self.pipe.torch_dtype) with torch.inference_mode(): latents = self.pipe.vae.single_encode(video, self.device) latents = latents.to(device=self.device, dtype=self.pipe.torch_dtype).detach() self.pipe.load_models_to_device(["text_encoder"]) with torch.inference_mode(): context = self.pipe.prompter.encode_prompt( batch["caption"], positive=True, device=self.device ) conditioning = self.pipe.build_i2v_conditioning( input_image=batch["first_frame_pil"][0], num_frames=self.hparams.num_frames, height=self.hparams.height, width=self.hparams.width, ) if conditioning is None: raise RuntimeError("The selected model does not expose Wan2.2 I2V conditioning") camera = { "pose": batch["pose"].to(device=self.device, dtype=self.pipe.torch_dtype), "x_fov": batch["x_fov"].to(device=self.device, dtype=self.pipe.torch_dtype), "xi": batch["xi"].to(device=self.device, dtype=self.pipe.torch_dtype), } self.pipe.load_models_to_device(["dit", "dit2"]) loss = self.pipe.training_loss( input_latents=latents, noise=torch.randn_like(latents), context=context, height=self.hparams.height, width=self.hparams.width, camera_control_panshot=camera, y=conditioning, first_frame_latents=conditioning[:, 4:, 0:1].clone(), min_timestep_boundary=0.9, max_timestep_boundary=1.0, switch_DiT_boundary=0.9, use_gradient_checkpointing=True, use_gradient_checkpointing_offload=False, ) self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True) return loss def configure_optimizers(self) -> torch.optim.Optimizer: parameters = [ parameter for parameter in self.pipe.dit.parameters() if parameter.requires_grad ] if not parameters: raise RuntimeError("No trainable parameters found in the high-noise expert") return torch.optim.AdamW( parameters, lr=self.hparams.learning_rate, weight_decay=self.hparams.weight_decay, ) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--config", type=Path, default=Path("configs/train_rdpo_high_only.yaml"), help="Training YAML (see configs/train_rdpo_high_only.yaml).", ) parser.add_argument("--model_path", default=None, help="Override config model_path.") parser.add_argument("--output_dir", type=Path, default=None, help="Override config output_dir.") parser.add_argument("--num_gpus", type=int, default=None, help="Override trainer.num_gpus.") parser.add_argument("--max_steps", type=int, default=None, help="Override trainer.max_steps.") parser.add_argument("--resume_from_checkpoint", type=Path, default=None) return parser def build_strategy(num_gpus: int) -> "str | Any": if num_gpus <= 1: return "auto" from pytorch_lightning.strategies import FSDPStrategy from torch.distributed.fsdp import ShardingStrategy from diffsynth.models.wan_video_dit import DiTBlock return FSDPStrategy( sharding_strategy=ShardingStrategy.FULL_SHARD, auto_wrap_policy={DiTBlock}, state_dict_type="sharded", use_orig_params=True, ) def main() -> None: args = build_parser().parse_args() config = load_config(args.config) data_config = config["data"] trainer_config = config.get("trainer", {}) optimizer_config = config.get("optimizer", {}) model_path = args.model_path or config.get("model_path", SCOPE_MODEL_ID) output_dir = Path(args.output_dir or config.get("output_dir", "outputs/training")) num_gpus = args.num_gpus or int(trainer_config.get("num_gpus", 8)) max_steps = args.max_steps or int(trainer_config.get("max_steps", 10_000)) pl.seed_everything(int(config.get("seed", 42)), workers=True) dataset = build_training_dataset(data_config) num_workers = int(data_config.get("num_workers", 4)) loader = DataLoader( dataset, batch_size=1, shuffle=True, num_workers=num_workers, collate_fn=_collate, pin_memory=True, persistent_workers=num_workers > 0, ) checkpoint = ModelCheckpoint( dirpath=output_dir / "checkpoints", filename="scope-{step:06d}", every_n_train_steps=int(trainer_config.get("save_every_n_steps", 1_000)), save_top_k=-1, save_last=True, save_on_train_epoch_end=False, ) model = SCoPEFineTuner( model_path=model_path, learning_rate=float(optimizer_config.get("learning_rate", 2e-5)), weight_decay=float(optimizer_config.get("weight_decay", 1e-2)), height=int(data_config.get("height", 480)), width=int(data_config.get("width", 832)), num_frames=int(data_config.get("num_frames", 81)), ) trainer = pl.Trainer( accelerator="gpu", devices=num_gpus, strategy=build_strategy(num_gpus), precision=str(trainer_config.get("precision", "bf16-mixed")), max_steps=max_steps, gradient_clip_val=float(trainer_config.get("gradient_clip_val", 1.0)), default_root_dir=output_dir, callbacks=[checkpoint], log_every_n_steps=int(trainer_config.get("log_every_n_steps", 10)), ) trainer.fit(model, train_dataloaders=loader, ckpt_path=args.resume_from_checkpoint) if __name__ == "__main__": main()