# PID distillation model — inference subset of the DMD2-distilled student. # # The training-time teacher / fake_score / discriminator / DMD-loss machinery has been # stripped; what remains is the student net (`self.net`) plus the few-step sampler # (`_get_t_list`, `_student_sample_loop`, `_velocity_to_x0`) consumed by # `generate_samples_from_batch`. from __future__ import annotations import logging from collections import OrderedDict from contextlib import nullcontext from typing import Optional import attrs import torch from pid._src.models.pid_model import PidModel, PidModelConfig logger = logging.getLogger(__name__) @attrs.define(slots=False) class PidDistillModelConfig(PidModelConfig): """Inference config for the distilled student.""" # Few-step student schedule. student_timestep: float = 1.0 student_sample_steps: int = 1 student_sample_type: str = "sde" student_t_list: Optional[list] = None student_input_mode: str = "teacher_forcing" class PidDistillModel(PidModel): """Inference-only PID distilled student.""" def __init__(self, config: PidDistillModelConfig): # Stubs left in place so any parent code that probes for these attributes # gets None instead of AttributeError. self.teacher = None self.fake_score = None self.discriminator = None super().__init__(config) # --------------------------------------------------------------------- # Net output ↔ (x0, velocity) conversion # --------------------------------------------------------------------- def _net_output_to_x0( self, x_t: torch.Tensor, net_output: torch.Tensor, t: torch.Tensor, prediction_type: str, ) -> torch.Tensor: if prediction_type == "x0": return net_output.to(x_t.dtype) if prediction_type == "velocity": original_dtype = x_t.dtype s = [x_t.shape[0]] + [1] * (x_t.ndim - 1) t_shaped = t.double().view(*s) return (x_t.double() - t_shaped * net_output.double()).to(original_dtype) raise ValueError(f"Invalid prediction_type: {prediction_type}") def _net_output_to_velocity( self, x_t: torch.Tensor, net_output: torch.Tensor, t: torch.Tensor, prediction_type: str, ) -> torch.Tensor: if prediction_type == "velocity": return net_output if prediction_type == "x0": original_dtype = x_t.dtype s = [x_t.shape[0]] + [1] * (x_t.ndim - 1) t_shaped = t.double().view(*s).clamp(min=5e-2) return ((x_t.double() - net_output.double()) / t_shaped).to(original_dtype) raise ValueError(f"Invalid prediction_type: {prediction_type}") def _velocity_to_x0(self, x_t: torch.Tensor, net_output: torch.Tensor, t: torch.Tensor) -> torch.Tensor: return self._net_output_to_x0(x_t, net_output, t, self.config.prediction_type) # --------------------------------------------------------------------- # Multi-step student sampler # --------------------------------------------------------------------- def _get_t_list(self, device, num_steps: Optional[int] = None) -> torch.Tensor: target_steps = num_steps if num_steps is not None else self.config.student_sample_steps if self.config.student_t_list is not None: full_t = torch.tensor(self.config.student_t_list, device=device, dtype=torch.float32) if target_steps != self.config.student_sample_steps: indices = torch.linspace(0, len(full_t) - 1, target_steps + 1).round().long() t_list = full_t[indices] else: t_list = full_t else: t_list = torch.linspace( self.config.student_timestep, 0.0, target_steps + 1, device=device, dtype=torch.float32, ) assert abs(t_list[-1].item()) < 1e-6, "t_list must end at 0" if num_steps is not None: logger.info(f"[distill inference] num_steps={num_steps}, t_list={t_list.tolist()}") return t_list def _student_sample_loop( self, noise: torch.Tensor, t_list: torch.Tensor, caption_embs: torch.Tensor, lq_video_or_image: Optional[torch.Tensor], lq_latent: Optional[torch.Tensor], degrade_sigma_tensor: Optional[torch.Tensor], generator: Optional[torch.Generator] = None, ) -> torch.Tensor: B = noise.shape[0] timescale = self.fm_trainer.timescale autocast_ctx = torch.autocast("cuda", dtype=self.autocast_dtype) if self.autocast_dtype else nullcontext() x = noise net = self.net with autocast_ctx: for t_cur, t_next in zip(t_list[:-1], t_list[1:]): t_cur_batch = t_cur.expand(B) t_cur_scaled = t_cur_batch * timescale v_pred = net( x, t_cur_scaled, caption_embs, lq_video_or_image=lq_video_or_image, lq_latent=lq_latent, degrade_sigma=degrade_sigma_tensor, ) if t_next.item() > 0: if self.config.student_sample_type == "ode": v_for_step = self._net_output_to_velocity(x, v_pred, t_cur_batch, self.config.prediction_type) dt = t_next - t_cur x = x + dt * v_for_step else: x0_pred = self._velocity_to_x0(x, v_pred, t_cur_batch) eps_infer = torch.randn( x0_pred.shape, device=x0_pred.device, dtype=x0_pred.dtype, generator=generator, ) s = [B] + [1] * (x.ndim - 1) t_next_bcast = t_next.reshape(1).expand(s) x = (1.0 - t_next_bcast) * x0_pred + t_next_bcast * eps_infer else: x = self._velocity_to_x0(x, v_pred, t_cur_batch) return x # --------------------------------------------------------------------- # Inference entry point # --------------------------------------------------------------------- @torch.no_grad() def generate_samples_from_batch( self, data_batch: dict, guidance: float = None, cfg_scale: float = None, num_steps: int = None, seed: int = 0, image_size=None, shift: float = None, is_negative_prompt: bool = False, **kwargs, ): # Encode any missing LQ_latent via the frozen VAE so callers can pass either # LQ_video_or_image or LQ_latent. if "LQ_latent" not in data_batch and "LQ_video_or_image" in data_batch and self.vae_encoder is not None: data_batch["LQ_latent"] = ( self.encode_lq_latent(data_batch["LQ_video_or_image"]).contiguous().to(**self.tensor_kwargs) ) if "degrade_sigma" not in data_batch and "LQ_latent" in data_batch: B = data_batch["LQ_latent"].shape[0] data_batch["degrade_sigma"] = torch.zeros(B, device=data_batch["LQ_latent"].device, dtype=torch.float32) x0_key = self.config.input_data_key if image_size is None and x0_key in data_batch: x0_shape = data_batch[x0_key].shape img_h, img_w = x0_shape[-2], x0_shape[-1] else: image_size = image_size or self.config.image_size if isinstance(image_size, (list, tuple)): img_h, img_w = int(image_size[0]), int(image_size[1]) else: img_h = img_w = int(image_size) # Determine shift: explicit arg > SD3-style dynamic_shift (if configured) > config default. # The 4-step distilled sampler doesn't consume `shift` directly (it uses # student_t_list), but we keep the precedence ladder symmetric with the # non-distilled inference path in case future call sites read it. if shift is None and self.config.dynamic_shift is not None: import math _ds = self.config.dynamic_shift shift = _ds["base_shift"] * math.sqrt(max(img_h, img_w) / _ds["base_image_size_for_shift_calc"]) captions = data_batch[self.config.input_caption_key] if isinstance(captions, str): captions = [captions] B = len(captions) if self.config.use_fixed_prompt: captions = [self.config.fixed_positive_prompt] * B caption_embs, _ = self._encode_text_raw(captions) caption_embs = caption_embs.to(**self.tensor_kwargs) lq_video_or_image = None lq_latent = None if self.config.lq_condition_type in ("image", "image_latent"): lq_video_or_image = data_batch.get("LQ_video_or_image") if lq_video_or_image is not None: lq_video_or_image = lq_video_or_image.to(**self.tensor_kwargs) if self.config.lq_condition_type in ("latent", "image_latent"): lq_latent = data_batch.get("LQ_latent") if lq_latent is not None: lq_latent = lq_latent.to(**self.tensor_kwargs) sigma_val = data_batch.get("degrade_sigma", 0.0) if isinstance(sigma_val, torch.Tensor): degrade_sigma_tensor = sigma_val.to(device="cuda", dtype=torch.float32).reshape(-1) if degrade_sigma_tensor.numel() == 1: degrade_sigma_tensor = degrade_sigma_tensor.expand(B).contiguous() assert degrade_sigma_tensor.shape == (B,), ( f"data_batch['degrade_sigma'] expected [B={B}], got {tuple(degrade_sigma_tensor.shape)}" ) elif isinstance(sigma_val, (list, tuple)): degrade_sigma_tensor = torch.tensor(sigma_val, device="cuda", dtype=torch.float32) assert degrade_sigma_tensor.shape == (B,), ( f"data_batch['degrade_sigma'] expected length {B}, got {len(sigma_val)}" ) else: degrade_sigma_tensor = torch.full((B,), float(sigma_val), device="cuda", dtype=torch.float32) gen = torch.Generator(device="cuda").manual_seed(int(seed)) noise = torch.randn(B, 3, img_h, img_w, device="cuda", generator=gen) autocast_ctx = torch.autocast("cuda", dtype=self.autocast_dtype) if self.autocast_dtype else nullcontext() net = self.net net.eval() effective_steps = num_steps if num_steps is not None else self.config.student_sample_steps if effective_steps == 1: t_student = torch.full((B,), self.config.student_timestep, device="cuda", dtype=torch.float32) t_student_scaled = t_student * self.fm_trainer.timescale with autocast_ctx: v_student = net( noise, t_student_scaled, caption_embs, lq_video_or_image=lq_video_or_image, lq_latent=lq_latent, degrade_sigma=degrade_sigma_tensor, ) x0_student = self._velocity_to_x0(noise, v_student, t_student) else: t_list = self._get_t_list(device=torch.device("cuda"), num_steps=num_steps) x0_student = self._student_sample_loop( noise, t_list, caption_embs, lq_video_or_image, lq_latent, degrade_sigma_tensor, generator=gen, ) return x0_student.clamp(-1, 1).unsqueeze(2) # --------------------------------------------------------------------- # Checkpoint helpers (only the student `net.` prefix matters at inference) # --------------------------------------------------------------------- def model_dict(self) -> dict: return {"net": self.net} def state_dict(self, *args, **kwargs): return self.net.state_dict(prefix="net.") def load_state_dict(self, state_dict, strict=True, assign=False, **kwargs): _net_sd = OrderedDict() for k, v in state_dict.items(): if k.startswith("net.") and not k.startswith("net_ema."): _net_sd[k[len("net.") :]] = v elif k.startswith("net_ema.") or k.startswith("fake_score.") or k.startswith("discriminator."): continue else: _net_sd[k] = v missing, unexpected = self.net.load_state_dict(_net_sd, strict=False, assign=assign) if missing: lq_missing = [k for k in missing if "lq_proj" in k] other_missing = [k for k in missing if "lq_proj" not in k] if lq_missing: logger.info(f"Expected missing LQ keys ({len(lq_missing)} keys)") if other_missing and strict: logger.warning(f"Missing keys in net: {other_missing}") if unexpected: logger.warning(f"Unexpected keys in net: {unexpected}") def on_train_start(self, memory_format=torch.preserve_format) -> None: super().on_train_start(memory_format)