from .base_pipeline import BasePipeline import torch import random def _resolve_sigma_shift(pipe, inputs=None, default=5.0): if inputs is not None and isinstance(inputs, dict) and "sigma_shift" in inputs and inputs["sigma_shift"] is not None: return inputs["sigma_shift"] return getattr(pipe, "sigma_shift", default) def _configure_training_timesteps(pipe, inputs): mode = inputs.get("train_timestep_mode", getattr(pipe, "train_timestep_mode", "full")) sigma_shift = _resolve_sigma_shift(pipe, inputs) full_steps = int(inputs.get("train_full_num_inference_steps", getattr(pipe, "train_full_num_inference_steps", 1000))) inference_steps = int(inputs.get("train_inference_num_steps", getattr(pipe, "train_inference_num_steps", inputs.get("num_inference_steps", 20)))) mixed_prob = float(inputs.get("train_timestep_mixed_inference_prob", getattr(pipe, "train_timestep_mixed_inference_prob", 0.8))) active_mode = mode if mode == "mixed": active_mode = "inference" if torch.rand(1).item() < mixed_prob else "full" num_steps = inference_steps if active_mode == "inference" else full_steps pipe.scheduler.set_timesteps(num_steps, training=True, shift=sigma_shift) return active_mode, num_steps def _align_aug_latents_to_clean(clean_latents, aug_latents, inputs): if aug_latents is None: return None if aug_latents.shape == clean_latents.shape: return aug_latents aug_aligned = clean_latents.clone() t_aug = aug_latents.shape[2] t_clean = clean_latents.shape[2] anchor_latent_vid = inputs.get("anchor_latent_vid", None) if anchor_latent_vid is not None and anchor_latent_vid.shape[1] > 0: num_ancs = int(anchor_latent_vid.shape[1]) aug_aligned[:, :, num_ancs:num_ancs + t_aug, :, :] = aug_latents else: aug_aligned[:, :, t_clean - t_aug:, :, :] = aug_latents return aug_aligned def _resolve_trajectory_correction_config(pipe: BasePipeline): schedule = getattr(pipe, "trajectory_correction_schedule", None) weight = getattr(pipe, "trajectory_correction_weight", None) if schedule is not None: return schedule, float(1.0 if weight is None else weight) # Backward compatibility for old checkpoints/scripts that encoded schedule # choices as negative traj_correction_alpha values. legacy_alpha = float(getattr(pipe, "traj_correction_alpha", 1.0)) if legacy_alpha >= 0: return "fixed", legacy_alpha legacy_schedule_map = { -1.0: "sigma_ramp", -2.0: "augmented_velocity", -3.0: "scheduler_weight", -4.0: "gaussian_timestep", } return legacy_schedule_map.get(legacy_alpha, "fixed"), 1.0 def _resolve_trajectory_correction_alpha(pipe: BasePipeline, timestep): schedule, weight = _resolve_trajectory_correction_config(pipe) if schedule == "fixed": return weight if isinstance(timestep, torch.Tensor): timestep_cpu = timestep.detach().to(device=pipe.scheduler.timesteps.device, dtype=pipe.scheduler.timesteps.dtype) else: timestep_cpu = torch.tensor([timestep], device=pipe.scheduler.timesteps.device, dtype=pipe.scheduler.timesteps.dtype) if schedule == "sigma_ramp": # Monotonic with noise level: low-noise -> smaller alpha, high-noise -> larger alpha. timesteps = pipe.scheduler.timesteps sigmas = pipe.scheduler.sigmas timestep_values = timestep_cpu.reshape(-1) timestep_index = torch.argmin((timesteps - timestep_values[0]).abs()) sigma_t = sigmas[timestep_index].to(dtype=torch.float32) sigma_min = torch.min(sigmas).to(dtype=torch.float32) sigma_max = torch.max(sigmas).to(dtype=torch.float32) sigma_range = (sigma_max - sigma_min).clamp_min(1e-8) alpha = (sigma_t - sigma_min) / sigma_range return (alpha * weight).to(dtype=pipe.torch_dtype, device=pipe.device) if schedule == "scheduler_weight": alpha = pipe.scheduler.training_weight(timestep_cpu.reshape(-1)[:1]) return (alpha * weight).to(dtype=pipe.torch_dtype, device=pipe.device) if schedule == "gaussian_timestep": timesteps = pipe.scheduler.timesteps.to(dtype=torch.float32) steps = float(getattr(pipe.scheduler, 'num_train_timesteps', 1000)) gaussian = torch.exp(-2 * ((timesteps - steps / 2) / steps) ** 2) timestep_values = timestep_cpu.reshape(-1).to(dtype=timesteps.dtype) timestep_index = torch.argmin((timesteps - timestep_values[0]).abs()) gaussian_timestep_value = gaussian[timestep_index] gaussian_min = torch.min(gaussian) gaussian_max = torch.max(gaussian) gaussian_range = (gaussian_max - gaussian_min).clamp_min(1e-8) alpha = (gaussian_timestep_value - gaussian_min) / gaussian_range return (alpha * weight).to(dtype=pipe.torch_dtype, device=pipe.device) raise ValueError(f"Unknown trajectory_correction_schedule: {schedule}") def _resolve_adaptive_aug_strength(pipe: BasePipeline, timestep): if not bool(getattr(pipe, 'adaptive_aug', False)): return 1.0 timesteps = pipe.scheduler.timesteps sigmas = pipe.scheduler.sigmas if isinstance(timestep, torch.Tensor): timestep_values = timestep.detach().to(device=timesteps.device, dtype=timesteps.dtype).reshape(-1) else: timestep_values = torch.tensor([timestep], device=timesteps.device, dtype=timesteps.dtype) timestep_index = torch.argmin((timesteps - timestep_values[0]).abs()) sigma_t = sigmas[timestep_index].to(dtype=torch.float32) gamma = max(float(getattr(pipe, 'adaptive_aug_gamma', 2.0)), 1e-6) min_strength = float(getattr(pipe, 'adaptive_aug_min_strength', 0.0)) min_strength = max(0.0, min(1.0, min_strength)) # Low-noise repair targets are the most unstable. Fade augmentation out below # sigma_low, restore it by sigma_high, and shape the transition with gamma. sigma_low = float(getattr(pipe, 'adaptive_aug_sigma_low', 0.05)) sigma_high = float(getattr(pipe, 'adaptive_aug_sigma_high', 0.30)) sigma_low = max(0.0, sigma_low) sigma_high = max(sigma_low + 1e-6, sigma_high) p = ((sigma_t - sigma_low) / (sigma_high - sigma_low)).clamp(0.0, 1.0) strength = min_strength + (1.0 - min_strength) * torch.pow(p, gamma) return strength.to(dtype=pipe.torch_dtype, device=pipe.device) def _rand_aug_allowed_for_timestep(pipe: BasePipeline, timestep): threshold = float(getattr(pipe, 'rand_aug_sigma_threshold', 0.5)) if threshold <= 0.0: return True timesteps = pipe.scheduler.timesteps sigmas = pipe.scheduler.sigmas if isinstance(timestep, torch.Tensor): timestep_values = timestep.detach().to(device=timesteps.device, dtype=timesteps.dtype).reshape(-1) else: timestep_values = torch.tensor([timestep], device=timesteps.device, dtype=timesteps.dtype) timestep_index = torch.argmin((timesteps - timestep_values[0]).abs()) sigma_t = sigmas[timestep_index].to(dtype=torch.float32) return bool((sigma_t >= threshold).item()) def _build_video_aug_training_pair(pipe: BasePipeline, inputs, noise, timestep, aug_latents=None): clean_latents = inputs["input_latents"] x_t_clean = pipe.scheduler.add_noise(clean_latents, noise, timestep) v_clean_t = pipe.scheduler.training_target(clean_latents, noise, timestep) if aug_latents is None: return x_t_clean, v_clean_t aug_latents = _align_aug_latents_to_clean(clean_latents, aug_latents, inputs) aug_strength = _resolve_adaptive_aug_strength(pipe, timestep) aug_latents = clean_latents + (aug_latents - clean_latents) * aug_strength x_t_aug = pipe.scheduler.add_noise(aug_latents, noise, timestep) # Special mode: construct velocity target directly from augmented noisy sample. trajectory_schedule, _ = _resolve_trajectory_correction_config(pipe) if trajectory_schedule == "augmented_velocity": if isinstance(timestep, torch.Tensor): timestep_values = timestep.detach().to(device=pipe.scheduler.timesteps.device, dtype=pipe.scheduler.timesteps.dtype).reshape(-1) else: timestep_values = torch.tensor([timestep], device=pipe.scheduler.timesteps.device, dtype=pipe.scheduler.timesteps.dtype) timestep_index = torch.argmin((pipe.scheduler.timesteps - timestep_values[0]).abs()) sigma_t = pipe.scheduler.sigmas[timestep_index].to(dtype=pipe.torch_dtype, device=pipe.device) # sigma_t = pipe.scheduler.sigmas[timestep_index].to(dtype=pipe.torch_dtype, device=pipe.device).clamp_min(1e-8) v_t = (x_t_aug.float() - clean_latents.float()) / sigma_t return x_t_aug, v_t alpha = _resolve_trajectory_correction_alpha(pipe, timestep) v_t = v_clean_t.float() + (x_t_aug.float() - x_t_clean.float()) * alpha return x_t_aug, v_t def _apply_motion_latent_noise(pipe: BasePipeline, inputs, noise, timestep, start_motion): if not getattr(pipe, 'add_noise_to_motion_latent', False): return if "y" not in inputs or inputs["y"] is None or inputs.get("clean_motion_latent") is None: return clean_motion_latents = inputs["clean_motion_latent"] timestep_cpu = timestep.cpu() timestep_index = torch.argmin((pipe.scheduler.timesteps - timestep_cpu).abs()) sigma = pipe.scheduler.sigmas[timestep_index].to(dtype=pipe.torch_dtype, device=pipe.device) if getattr(pipe, 'motion_latent_shared_noise', False): motion_noise = noise[0, :, :clean_motion_latents.shape[1], :, :].to(device=pipe.device, dtype=pipe.torch_dtype) else: motion_noise = torch.randn_like(clean_motion_latents) noisy_motion_latents = ((1.0 - sigma) * clean_motion_latents + sigma * motion_noise).unsqueeze(0) num_anchor_latents = start_motion num_motion_latents = getattr(pipe, 'num_motion_latents', 1) inputs["y"][:, 4:, num_anchor_latents:num_anchor_latents + num_motion_latents, :, :] = noisy_motion_latents.to(inputs["y"].dtype) def _run_training_pass(pipe: BasePipeline, inputs, timestep, noise, start_motion, aug_latents=None, compute_loss=True): run_inputs = dict(inputs) if run_inputs.get("y") is not None: run_inputs["y"] = run_inputs["y"].clone() run_inputs["noise"] = noise run_inputs["latents"], training_target = _build_video_aug_training_pair( pipe, run_inputs, noise, timestep, aug_latents=aug_latents, ) models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} noise_pred = pipe.model_fn(**models, **run_inputs, timestep=timestep) loss = None if compute_loss: loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float()) loss = loss * pipe.scheduler.training_weight(timestep) return run_inputs, training_target, noise_pred, loss def _compute_x1_pred_aug_latents(pipe: BasePipeline, inputs, timestep, noise, start_motion, with_grad=False): if with_grad: run_inputs, _, noise_pred, loss = _run_training_pass( pipe, inputs, timestep, noise, start_motion, compute_loss=True, ) x_1_pred = pipe.scheduler.step(noise_pred, timestep, run_inputs["latents"], to_final=True, self_corr=False) return x_1_pred.detach(), loss with torch.no_grad(): run_inputs, _, noise_pred, _ = _run_training_pass( pipe, inputs, timestep, noise, start_motion, compute_loss=False, ) x_1_pred = pipe.scheduler.step(noise_pred, timestep, run_inputs["latents"], to_final=True, self_corr=False) return x_1_pred.detach(), None def FlowMatchSFTLoss(pipe: BasePipeline, **inputs): preset_timestep = inputs.get("train_sampled_timestep", None) if preset_timestep is None: _configure_training_timesteps(pipe, inputs) elif not hasattr(pipe.scheduler, "timesteps") or len(pipe.scheduler.timesteps) == 0: # Fallback for unexpected call paths that did not pre-configure scheduler. _configure_training_timesteps(pipe, inputs) max_timestep_boundary = int(inputs.get("max_timestep_boundary", 1) * len(pipe.scheduler.timesteps)) min_timestep_boundary = int(inputs.get("min_timestep_boundary", 0) * len(pipe.scheduler.timesteps)) anchor_latent_vid = inputs.get("anchor_latent_vid", None) start_motion = 0 if anchor_latent_vid is not None and anchor_latent_vid.shape[1] > 0: inputs["input_latents"] = torch.cat([anchor_latent_vid.unsqueeze(0), inputs["input_latents"]], dim=2) start_motion = anchor_latent_vid.shape[1] if preset_timestep is None: timestep_id = torch.randint(min_timestep_boundary, max_timestep_boundary, (1,)) timestep = pipe.scheduler.timesteps[timestep_id].to(dtype=pipe.torch_dtype, device=pipe.device) else: if isinstance(preset_timestep, torch.Tensor): timestep = preset_timestep.to(dtype=pipe.torch_dtype, device=pipe.device).reshape(-1)[:1] else: timestep = torch.tensor([preset_timestep], dtype=pipe.torch_dtype, device=pipe.device) noise = torch.randn_like(inputs["input_latents"]) video_aug_prob = float(getattr(pipe, "video_aug_prob", 1.0)) mode_flags = { "rand_aug": bool(getattr(pipe, "rand_aug", False)), "mixed_aug": bool(getattr(pipe, "mixed_aug", False)), "use_self_pred_aug": bool(getattr(pipe, "use_self_pred_aug", False)), } enabled_modes = [name for name, enabled in mode_flags.items() if enabled] if len(enabled_modes) > 1: raise ValueError("Enable only one augmentation mode at a time: rand_aug, mixed_aug, or use_self_pred_aug.") selected_mode = None if enabled_modes and torch.rand(1).item() < video_aug_prob: only_mode = enabled_modes[0] if only_mode == "mixed_aug": selected_mode = random.choice(["rand_aug", "use_self_pred_aug"]) else: selected_mode = only_mode if selected_mode == "use_self_pred_aug": first_backprop = bool(getattr(pipe, "self_pred_aug_backprop_first", False)) self_pred_aug, first_loss = _compute_x1_pred_aug_latents( pipe, inputs, timestep, noise, start_motion, with_grad=first_backprop, ) _, _, _, second_loss = _run_training_pass( pipe, inputs, timestep, noise, start_motion, aug_latents=self_pred_aug, compute_loss=True, ) if first_backprop and first_loss is not None: return 0.5 * (first_loss + second_loss) return second_loss rand_aug_latents = None if selected_mode == "rand_aug" and _rand_aug_allowed_for_timestep(pipe, timestep): rand_aug_latents = inputs.get("input_latents_aug", None) _, _, _, loss = _run_training_pass( pipe, inputs, timestep, noise, start_motion, aug_latents=rand_aug_latents, compute_loss=True, ) return loss def DirectDistillLoss(pipe: BasePipeline, **inputs): pipe.scheduler.set_timesteps(inputs["num_inference_steps"], shift=_resolve_sigma_shift(pipe, inputs)) pipe.scheduler.training = True models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} for progress_id, timestep in enumerate(pipe.scheduler.timesteps): timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) noise_pred = pipe.model_fn(**models, **inputs, timestep=timestep, progress_id=progress_id) inputs["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs) loss = torch.nn.functional.mse_loss(inputs["latents"].float(), inputs["input_latents"].float()) return loss class TrajectoryImitationLoss(torch.nn.Module): def __init__(self): super().__init__() self.initialized = False def initialize(self, device): import lpips # TODO: remove it self.loss_fn = lpips.LPIPS(net='alex').to(device) self.initialized = True def fetch_trajectory(self, pipe: BasePipeline, timesteps_student, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale): trajectory = [inputs_shared["latents"].clone()] pipe.scheduler.set_timesteps(num_inference_steps, target_timesteps=timesteps_student, shift=_resolve_sigma_shift(pipe, inputs_shared)) models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} for progress_id, timestep in enumerate(pipe.scheduler.timesteps): timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) noise_pred = pipe.cfg_guided_model_fn( pipe.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) inputs_shared["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred.detach(), **inputs_shared) trajectory.append(inputs_shared["latents"].clone()) return pipe.scheduler.timesteps, trajectory def align_trajectory(self, pipe: BasePipeline, timesteps_teacher, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale): loss = 0 pipe.scheduler.set_timesteps(num_inference_steps, training=True, shift=_resolve_sigma_shift(pipe, inputs_shared)) models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} for progress_id, timestep in enumerate(pipe.scheduler.timesteps): timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) progress_id_teacher = torch.argmin((timesteps_teacher - timestep).abs()) inputs_shared["latents"] = trajectory_teacher[progress_id_teacher] noise_pred = pipe.cfg_guided_model_fn( pipe.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) sigma = pipe.scheduler.sigmas[progress_id] sigma_ = 0 if progress_id + 1 >= len(pipe.scheduler.timesteps) else pipe.scheduler.sigmas[progress_id + 1] if progress_id + 1 >= len(pipe.scheduler.timesteps): latents_ = trajectory_teacher[-1] else: progress_id_teacher = torch.argmin((timesteps_teacher - pipe.scheduler.timesteps[progress_id + 1]).abs()) latents_ = trajectory_teacher[progress_id_teacher] target = (latents_ - inputs_shared["latents"]) / (sigma_ - sigma) loss = loss + torch.nn.functional.mse_loss(noise_pred.float(), target.float()) * pipe.scheduler.training_weight(timestep) return loss def compute_regularization(self, pipe: BasePipeline, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, num_inference_steps, cfg_scale): inputs_shared["latents"] = trajectory_teacher[0] pipe.scheduler.set_timesteps(num_inference_steps, shift=_resolve_sigma_shift(pipe, inputs_shared)) models = {name: getattr(pipe, name) for name in pipe.in_iteration_models} for progress_id, timestep in enumerate(pipe.scheduler.timesteps): timestep = timestep.unsqueeze(0).to(dtype=pipe.torch_dtype, device=pipe.device) noise_pred = pipe.cfg_guided_model_fn( pipe.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) inputs_shared["latents"] = pipe.step(pipe.scheduler, progress_id=progress_id, noise_pred=noise_pred.detach(), **inputs_shared) image_pred = pipe.vae_decoder(inputs_shared["latents"]) image_real = pipe.vae_decoder(trajectory_teacher[-1]) loss = self.loss_fn(image_pred.float(), image_real.float()) return loss def forward(self, pipe: BasePipeline, inputs_shared, inputs_posi, inputs_nega): if not self.initialized: self.initialize(pipe.device) with torch.no_grad(): pipe.scheduler.set_timesteps(8, shift=_resolve_sigma_shift(pipe, inputs_shared)) timesteps_teacher, trajectory_teacher = self.fetch_trajectory(inputs_shared["teacher"], pipe.scheduler.timesteps, inputs_shared, inputs_posi, inputs_nega, 50, 2) timesteps_teacher = timesteps_teacher.to(dtype=pipe.torch_dtype, device=pipe.device) loss_1 = self.align_trajectory(pipe, timesteps_teacher, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, 8, 1) loss_2 = self.compute_regularization(pipe, trajectory_teacher, inputs_shared, inputs_posi, inputs_nega, 8, 1) loss = loss_1 + loss_2 return loss