distribution_matching ¶

def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, Any],
) -> None:
    super().__init__(cfg=cfg, role_models=role_models)
    mcfg = self.method_config

    student_sample_steps = get_optional_int(mcfg, "student_sample_steps", where="method.student_sample_steps")
    if student_sample_steps is None:
        student_sample_steps = 4
    if int(student_sample_steps) <= 0:
        raise ValueError("method.student_sample_steps must be positive, "
                         f"got {student_sample_steps}")
    self._student_sample_steps = int(student_sample_steps)

    use_mean_velocity_raw = mcfg.get("use_mean_velocity", True)
    if not isinstance(use_mean_velocity_raw, bool):
        raise ValueError("method.use_mean_velocity must be a bool, "
                         f"got {type(use_mean_velocity_raw).__name__}")
    self._use_mean_velocity = bool(use_mean_velocity_raw)

    # Optional pinned rollout schedule (descending, absolute t-units).
    # Falls back to dmd_denoising_steps when absent.
    raw_t_list = mcfg.get("t_list_override", None)
    if raw_t_list is None:
        self._t_list_override: list[float] | None = None
    else:
        if not isinstance(raw_t_list, list) or not raw_t_list:
            raise ValueError("method.t_list_override must be a non-empty list of "
                             f"floats when set, got {raw_t_list!r}")
        t_list = [float(x) for x in raw_t_list]
        for i in range(len(t_list) - 1):
            if t_list[i] < t_list[i + 1]:
                raise ValueError("method.t_list_override must be descending, "
                                 f"got {t_list!r}")
        self._t_list_override = t_list

    # Scoring conditioning: AnyFlow scores against r=0 for the DMD branch.
    score_r_raw = mcfg.get("dmd_score_r_value", 0.0)
    try:
        self._dmd_score_r = float(score_r_raw)
    except (TypeError, ValueError) as exc:
        raise ValueError("method.dmd_score_r_value must be numeric, "
                         f"got {score_r_raw!r}") from exc

    # Optional teacher guidance scale for the DMD loss (carry over from
    # DMD2Method's behavior; default 1.0).
    guidance = get_optional_float(mcfg, "real_score_guidance_scale", where="method.real_score_guidance_scale")
    self._real_score_guidance = float(guidance) if guidance is not None else 1.0

def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, ModelBase],
) -> None:
    super().__init__(cfg=cfg, role_models=role_models)
    if "student" not in role_models:
        raise ValueError("AnyFlowPretrainMethod requires role 'student'")
    if not self.student._trainable:
        raise ValueError("AnyFlowPretrainMethod requires student to be trainable")

    mcfg = self.method_config
    self._diffusion_ratio = float(
        get_optional_float(mcfg, "diffusion_ratio", where="method.diffusion_ratio") or 0.5)
    self._consistency_ratio = float(
        get_optional_float(mcfg, "consistency_ratio", where="method.consistency_ratio") or 0.25)
    if self._diffusion_ratio + self._consistency_ratio > 1.0:
        raise ValueError("method.diffusion_ratio + method.consistency_ratio must "
                         f"be <= 1, got {self._diffusion_ratio} + "
                         f"{self._consistency_ratio}")

    # δ: finite-difference step in absolute train-timestep units.
    epsilon = get_optional_int(mcfg, "epsilon", where="method.epsilon")
    self._fd_epsilon = float(epsilon) if epsilon is not None else 5.0

    # Loss weighting scheme (uniform / gaussian / beta08).
    raw_weight_type = mcfg.get("weight_type", "beta08")
    if not isinstance(raw_weight_type, str):
        raise ValueError("method.weight_type must be a string, got "
                         f"{type(raw_weight_type).__name__}")
    weight_type = raw_weight_type.strip().lower()
    if weight_type not in {"uniform", "gaussian", "beta08"}:
        raise ValueError("method.weight_type must be one of "
                         "{uniform, gaussian, beta08}, "
                         f"got {raw_weight_type!r}")
    self._weight_type = weight_type

    # Guidance fused into the training target (default 1.0 = unused).
    fg = get_optional_float(mcfg, "fuse_guidance_scale", where="method.fuse_guidance_scale")
    self._fuse_guidance_scale = float(fg) if fg is not None else 1.0
    if self._fuse_guidance_scale <= 0.0:
        raise ValueError("method.fuse_guidance_scale must be positive, "
                         f"got {self._fuse_guidance_scale}")

    # Flow-map scheduler — uses pipeline_config.flow_shift if present
    # and falls back to method.shift (and finally 1.0).
    shift = float(getattr(self.training_config.pipeline_config, "flow_shift", 0.0) or 0.0)
    if shift <= 0.0:
        shift_override = get_optional_float(mcfg, "shift", where="method.shift")
        shift = float(shift_override) if shift_override is not None else 1.0
    self._shift = shift

    # Lazy-imported to avoid circular imports on package load.
    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
        FlowMapEulerDiscreteScheduler, )
    self._flow_map_scheduler = FlowMapEulerDiscreteScheduler(
        num_train_timesteps=int(self.student.num_train_timesteps),
        shift=self._shift,
    )

    self.student.init_preprocessors(self.training_config)
    self._init_optimizer_and_scheduler()

def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, ModelBase],
) -> None:
    super().__init__(cfg=cfg, role_models=role_models)

    if "student" not in role_models:
        raise ValueError("DMD2Method requires role 'student'")
    if "teacher" not in role_models:
        raise ValueError("DMD2Method requires role 'teacher'")
    if "critic" not in role_models:
        raise ValueError("DMD2Method requires role 'critic'")

    self.teacher = role_models["teacher"]
    self.critic = role_models["critic"]

    if not self.student._trainable:
        raise ValueError("DMD2Method requires student to be trainable")
    if self.teacher._trainable:
        raise ValueError("DMD2Method requires teacher to be "
                         "non-trainable")
    if not self.critic._trainable:
        raise ValueError("DMD2Method requires critic to be trainable")
    self._cfg_uncond = self._parse_cfg_uncond()
    self._rollout_mode = self._parse_rollout_mode()
    self._validate_preprocessed_data_type()
    self._configure_student_negative_conditioning()
    self._denoising_step_list: torch.Tensor | None = (None)
    (
        self._score_min_timestep,
        self._score_max_timestep,
    ) = self._parse_score_timestep_bounds()

    # Initialize preprocessors on student.
    self.student.init_preprocessors(self.training_config)

    self._init_optimizers_and_schedulers()

def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, ModelBase],
) -> None:
    super().__init__(
        cfg=cfg,
        role_models=role_models,
    )

    # Validate causal student.
    if not isinstance(self.student, CausalModelBase):
        raise ValueError("SelfForcingMethod requires a causal student "
                         "implementing CausalModelBase.")

    if self._rollout_mode != "simulate":
        raise ValueError("SelfForcingMethod only supports "
                         "method_config.rollout_mode='simulate'")

    mcfg = self.method_config

    chunk_size = get_optional_int(
        mcfg,
        "chunk_size",
        where="method_config.chunk_size",
    )
    if chunk_size is None:
        chunk_size = 3
    if chunk_size <= 0:
        raise ValueError("method_config.chunk_size must be a positive "
                         f"integer, got {chunk_size}")
    self._chunk_size = int(chunk_size)

    sample_type_raw = mcfg.get("student_sample_type", "sde")
    sample_type = _require_str(
        sample_type_raw,
        where="method_config.student_sample_type",
    )
    sample_type = sample_type.strip().lower()
    if sample_type not in {"sde", "ode"}:
        raise ValueError("method_config.student_sample_type must be one "
                         f"of {{sde, ode}}, got {sample_type_raw!r}")
    self._student_sample_type: Literal["sde", "ode"] = (
        sample_type  # type: ignore[assignment]
    )

    same_step_raw = mcfg.get("same_step_across_blocks", False)
    if same_step_raw is None:
        same_step_raw = False
    self._same_step_across_blocks = _require_bool(
        same_step_raw,
        where="method_config.same_step_across_blocks",
    )

    last_step_raw = mcfg.get("last_step_only", False)
    if last_step_raw is None:
        last_step_raw = False
    self._last_step_only = _require_bool(
        last_step_raw,
        where="method_config.last_step_only",
    )

    context_noise = get_optional_float(
        mcfg,
        "context_noise",
        where="method_config.context_noise",
    )
    if context_noise is None:
        context_noise = 0.0
    if context_noise < 0.0:
        raise ValueError("method_config.context_noise must be >= 0, "
                         f"got {context_noise}")
    self._context_noise = float(context_noise)

    enable_grad_raw = mcfg.get("enable_gradient_in_rollout", True)
    if enable_grad_raw is None:
        enable_grad_raw = True
    self._enable_gradient_in_rollout = _require_bool(
        enable_grad_raw,
        where="method_config.enable_gradient_in_rollout",
    )

    start_grad_frame = get_optional_int(
        mcfg,
        "start_gradient_frame",
        where="method_config.start_gradient_frame",
    )
    if start_grad_frame is None:
        start_grad_frame = 0
    if start_grad_frame < 0:
        raise ValueError("method_config.start_gradient_frame must be "
                         f">= 0, got {start_grad_frame}")
    self._start_gradient_frame = int(start_grad_frame)

    shift = float(getattr(
        self.training_config.pipeline_config,
        "flow_shift",
        0.0,
    ) or 0.0)
    self._sf_scheduler = SelfForcingFlowMatchScheduler(
        num_inference_steps=1000,
        num_train_timesteps=int(self.student.num_train_timesteps),
        shift=shift,
        sigma_min=0.0,
        extra_one_step=True,
        training=True,
    )

    self._sf_denoising_step_list: torch.Tensor | None = None

def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, Any],
) -> None:
    pipeline_config = getattr(cfg.training, "pipeline_config", None)
    vae_config = getattr(pipeline_config, "vae_config", None)
    if vae_config is not None:
        vae_config.load_encoder = True
        vae_config.load_decoder = True

    super().__init__(cfg=cfg, role_models=role_models)
    if not isinstance(self.student, CausalModelBase):
        raise ValueError("StreamingLongTuningMethod requires a causal student")

    mcfg = self.method_config
    default_chunk = get_optional_int(
        mcfg,
        "streaming_chunk_size",
        where="method.streaming_chunk_size",
    )
    default_max = get_optional_int(
        mcfg,
        "streaming_max_length",
        where="method.streaming_max_length",
    )
    default_num_latent_t = int(self.training_config.data.num_latent_t)
    self._distill_stages = parse_multi_phased_distill_schedule(
        mcfg.get("multi_phased_distill_schedule", None),
        default_num_latent_t=default_num_latent_t,
        default_streaming_chunk_size=default_chunk,
        default_streaming_max_length=default_max,
    )
    self._streaming_state: _StreamingState | None = None
    self._streaming_anchor_inject_k = int(
        get_optional_int(
            mcfg,
            "streaming_anchor_inject_k",
            where="method.streaming_anchor_inject_k",
        ) or 1)
    if self._streaming_anchor_inject_k < 1:
        raise ValueError("method.streaming_anchor_inject_k must be >= 1")
    reencode_anchor_raw = mcfg.get("streaming_reencode_overlap_anchor", True)
    if reencode_anchor_raw is None:
        reencode_anchor_raw = True
    self._streaming_reencode_overlap_anchor = _as_bool(
        reencode_anchor_raw,
        where="method.streaming_reencode_overlap_anchor",
    )
    require_full_blocks_raw = mcfg.get("streaming_require_full_blocks", False)
    if require_full_blocks_raw is None:
        require_full_blocks_raw = False
    self._streaming_require_full_blocks = _as_bool(
        require_full_blocks_raw,
        where="method.streaming_require_full_blocks",
    )
    real_score_cfg_raw = mcfg.get("real_score_cfg", True)
    if real_score_cfg_raw is None:
        real_score_cfg_raw = True
    self._real_score_cfg = _as_bool(
        real_score_cfg_raw,
        where="method.real_score_cfg",
    )
    self._score_min_timestep = self._score_timestep_from_ratio("min_timestep_ratio")
    self._score_max_timestep = self._score_timestep_from_ratio("max_timestep_ratio")

    max_stage_latents = max(s.num_latent_t for s in self._distill_stages)
    if max_stage_latents > default_num_latent_t:
        raise ValueError("training.data.num_latent_t must be at least the largest "
                         "multi-stage horizon so prepared conditioning is not trimmed "
                         "before streaming: "
                         f"num_latent_t={default_num_latent_t}, "
                         f"max_stage_latents={max_stage_latents}")