SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
35 lines
1.6 KiB
Python
35 lines
1.6 KiB
Python
"""Self-forcing — causal video distillation (extends DMD2): train against your own KV-cached causal
|
|
rollout.
|
|
|
|
The student rollout is the SHARED ``chunk_rollout`` loop (with the slab-KV cache) — exactly the loop
|
|
the engine streams at serve time — so the causal student is distilled under the runtime it will be
|
|
served on. The loss is DMD2's distribution-matching loss applied to the causal rollout latents. (The
|
|
demonstrated structure is the causal rollout + KV + DMD loss; this mini reuses the shared flow-match
|
|
sampler in place of the repo's dedicated SelfForcingFlowMatchScheduler.)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from v2._enums import ConsistencyLevel
|
|
from v2.request import DiffusionParams, TaskType, make_request
|
|
from v2.training.methods.base import new_instance
|
|
from v2.training.methods.dmd2 import DMD2Method
|
|
|
|
|
|
class SelfForcingMethod(DMD2Method):
|
|
name = "self_forcing"
|
|
consistency = ConsistencyLevel.C1
|
|
student_loop_id = "chunk_rollout" # the SHARED causal loop (KV-cached)
|
|
|
|
def _rollout_sample(self, prompt: str, seed: int) -> np.ndarray:
|
|
req = make_request(TaskType.T2V, self.student.card.model_id, prompt, diffusion=DiffusionParams(seed=seed))
|
|
res = self._rollout(req) # drives chunk_rollout + slab-KV
|
|
latents = res.outputs["latents"]
|
|
return np.asarray(latents, dtype="float32")
|
|
|
|
|
|
def build_self_forcing(causal_card, **kw) -> SelfForcingMethod:
|
|
"""Build a self-forcing method over a Wan-causal card (chunk_rollout loop)."""
|
|
return SelfForcingMethod(new_instance(causal_card), new_instance(causal_card), new_instance(causal_card), **kw)
|