SolitaryThinker
4e331eb7ff
Mechanical conversion of every relative import under v2/ to an absolute v2.* path (from .x / ..x / ...x -> from v2.<pkg>.x) so imports are unambiguous, grep-able, and stable when code is copied/moved between entrypoints (VideoGenerator, CLI, server). Surgical prefix-only rewrite: only the 'from <dots><module>' prefix changed — import names, parentheses, multi-line formatting, comments, and ordering are byte-for-byte preserved (no collapsing, no reorder, no unrelated reformatting). - 471 imports across 125 files; v2/tests/ was already absolute (untouched). - Validated: all 125 files compile, every 'from v2.* import' target resolves to a real module/package, zero relative imports remain (full sweep), CPU mini suite green (202 passed, 2 skipped).
70 lines
2.7 KiB
Python
70 lines
2.7 KiB
Python
"""TrainingMethod base (design_v3 §10; mirrors the repo's train/methods contract).
|
||
|
||
The load-bearing principle: every method's *rollout* drives the SAME loop the engine serves
|
||
(via ``_rollout`` → ``rollout_loop``), so there is one numerics surface. Methods differ only in
|
||
loss math, roles, and capture — not in a second sampler. ``manages_optimization`` matches the
|
||
landed ``TrainingMethod.manages_optimization()`` seam (RL owns its sample→score→inner-train cadence).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
|
||
from v2._enums import ConsistencyLevel, ExecutionProfile
|
||
from v2.cache import CacheManager
|
||
from v2.card import load_card
|
||
from v2.models.common import text_encode_node_fn
|
||
from v2.training.rollout import rollout_loop
|
||
|
||
|
||
def new_instance(card) -> Any:
|
||
"""A fresh resident instance of a card (a training role: student/teacher/critic/old/reference)."""
|
||
return load_card(card, cache_manager=CacheManager.from_card(card), validate=False)
|
||
|
||
|
||
def predict_x0(velocity: np.ndarray, noised: np.ndarray, sigma: float) -> np.ndarray:
|
||
"""Flow prediction → clean sample: x0 = x_σ − σ·v (Wan ``x0 = x_t − σ·model_output``)."""
|
||
return noised - sigma * np.asarray(velocity, dtype="float32")
|
||
|
||
|
||
class TrainingMethod(ABC):
|
||
name: str = "base"
|
||
consistency: ConsistencyLevel = ConsistencyLevel.C1
|
||
student_loop_id: str = "diffusion_denoise"
|
||
|
||
def __init__(self, student_instance: Any, *, lr: float = 0.05):
|
||
self.student = student_instance
|
||
self.lr = lr
|
||
self.iteration = 0
|
||
|
||
@property
|
||
def student_dit(self) -> Any:
|
||
return self.student.component("transformer")
|
||
|
||
def manages_optimization(self) -> bool:
|
||
return False
|
||
|
||
def get_grad_clip_targets(self, iteration: int = 0) -> dict[str, Any]:
|
||
return {"student": self.student_dit}
|
||
|
||
def consistency_level(self) -> ConsistencyLevel:
|
||
return self.consistency
|
||
|
||
# --- rollout helpers: the SAME loop the engine serves ---------------------- #
|
||
def _encode_slots(self, request, instance: Any = None) -> dict:
|
||
slots: dict = {}
|
||
text_encode_node_fn(instance or self.student, slots, request, None)
|
||
return slots
|
||
|
||
def _rollout(self, request, *, instance: Any = None, loop_id: str | None = None):
|
||
inst = instance or self.student
|
||
slots = self._encode_slots(request, inst)
|
||
return rollout_loop(inst, loop_id or self.student_loop_id, request, slots=slots,
|
||
profile=ExecutionProfile.ROLLOUT)
|
||
|
||
@abstractmethod
|
||
def train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
||
"""Return (loss_map, metrics). loss_map values are floats; metrics includes grad_norm/*."""
|