From 21d751fbb4bd025e645dd87306eb5f9ddb26b51d Mon Sep 17 00:00:00 2001 From: leffff Date: Wed, 15 Jul 2026 09:53:21 +0300 Subject: [PATCH] [bugfix]: address Kandinsky5 QAD PR review findings Fixes all P1/P2 items from the review of 7c7704da: - E2E test's data_path pointed at a nonexistent combined_parquet_dataset/ subdir; preprocess_kandinsky5_overfit.py writes directly to its output root. - Stage 2 loaded the raw DCP checkpoint-N/ directly, which isn't a diffusers model dir. Added dcp_to_diffusers --verify (strict-reload the export immediately) and wired the conversion into the E2E test, launcher scripts, and README between stage 1 and stage 2. - Kandinsky5DmdDenoisingStage called the transformer without set_forward_context, so LocalAttention's missing-context guard caught the AssertionError and silently fell back to plain SDPA -- Stage 2 validation never actually exercised ATTN_QAT_TRAIN. Wrapped the call to match the parent stage's contract and added a scoped backend assertion (ATTN_QAT_TRAIN/ATTN_QAT_INFER only, since other backends can legitimately fall back to auto-selection). - Exported DMD checkpoints kept the base T2V model_index.json, so VideoGenerator.from_pretrained resolved them to Kandinsky5T2VPipeline instead of the DMD sampler. Added Kandinsky5DMDConfig and documented the required override_pipeline_cls_name + pipeline_config combo. - The FSDP-training dtype-detection special case (scan params, default to bf16) applied unconditionally, so an explicit fp32 pipeline was silently cast to bf16 for plain inference. Extracted _resolve_target_dtype and gated the scan on isinstance(transformer, FSDPModule). - Preprocessing tokenized captions with a 512-token cap that doesn't account for the 129-token system template runtime strips off; now reads text_encoder_max_lengths from pipeline_config to match runtime. Also replaced unconditional item["cap"][0] indexing (silently takes the first character of a string caption) with a get_caption() helper that accepts/validates both the documented string schema and the list-of-variants schema some producers use. - test_kandinsky5_qat_attention_engages.py set FASTVIDEO_ATTENTION_BACKEND and cleared the attention backend selector cache without teardown, leaking ATTN_QAT_TRAIN into later tests in the same pytest process. Now uses monkeypatch + try/finally. Added regression tests: test_kandinsky5_target_dtype.py, test_kandinsky5_dmd_pipeline_resolution.py, test_kandinsky5_overfit_caption.py. Co-Authored-By: Claude Sonnet 5 --- .../kandinsky5/distill_dmd_qat.sh | 16 +- .../kandinsky5/dmd2_t2v_480p_qat.yaml | 5 + .../configs/fine_tuning/kandinsky5/README.md | 58 ++++++- .../fine_tuning/kandinsky5/finetune_qat.sh | 6 + fastvideo/configs/pipelines/__init__.py | 6 +- fastvideo/configs/pipelines/kandinsky5.py | 18 ++ .../preprocess_kandinsky5_overfit.py | 34 +++- fastvideo/pipelines/stages/kandinsky5.py | 128 ++++++++++++--- ...test_kandinsky5_dmd_pipeline_resolution.py | 63 +++++++ .../test_e2e_kandinsky5_dmd_t2v_overfit.py | 46 +++++- .../stages/test_kandinsky5_target_dtype.py | 67 ++++++++ .../test_kandinsky5_qat_attention_engages.py | 154 ++++++++++-------- .../test_kandinsky5_overfit_caption.py | 45 +++++ .../train/entrypoint/dcp_to_diffusers.py | 38 +++++ 14 files changed, 577 insertions(+), 107 deletions(-) create mode 100644 fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py create mode 100644 fastvideo/tests/stages/test_kandinsky5_target_dtype.py create mode 100644 fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py diff --git a/examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh b/examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh index 7ec6b4eac..4b1741d94 100755 --- a/examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh +++ b/examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh @@ -9,8 +9,20 @@ # Kandinsky5-specific handling needed). # # Run examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh (stage 1) -# first and point models.student/teacher/critic.init_from in -# dmd2_t2v_480p_qat.yaml at that checkpoint before launching this script. +# first. Its checkpoint-N/ is a raw DCP checkpoint (dcp/ + metadata/RNG +# state only) -- models.student/teacher/critic.init_from need a diffusers +# model directory (model_index.json + component subfolders), so convert it +# first: +# +# python -m fastvideo.train.entrypoint.dcp_to_diffusers \ +# --checkpoint outputs/kandinsky5_t2v_qat_finetune/checkpoint- \ +# --output-dir outputs/kandinsky5_t2v_qat_finetune/checkpoint--diffusers \ +# --role student --verify +# +# then point models.student/teacher/critic.init_from in dmd2_t2v_480p_qat.yaml +# at that exported checkpoint--diffusers directory (all three roles use +# the same weights; teacher/critic full-precision masking happens at load +# time, not by pointing at a different export) before launching this script. set -euo pipefail export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN diff --git a/examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml b/examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml index a4d6047ab..541c587f4 100644 --- a/examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml +++ b/examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml @@ -1,6 +1,11 @@ # DMD2 QAT-aware distillation: Kandinsky5 T2V Lite, 480p (teacher many-step # -> student 4-step), continuing from a stage-1 Attn-QAT finetune checkpoint. # +# init_from below must be a diffusers model directory (model_index.json + +# component subfolders), not the raw checkpoint-N/ stage 1 writes (dcp/ + +# metadata/RNG state only). Convert with fastvideo.train.entrypoint.dcp_to_diffusers +# first -- see distill_dmd_qat.sh's header for the exact command. +# # - Teacher: frozen, full-precision Kandinsky5 (loaded from the stage-1 QAT # checkpoint's base weights, or the original pretrained checkpoint -- # the _loading_teacher_critic_model gate in component_loader.py masks diff --git a/examples/train/configs/fine_tuning/kandinsky5/README.md b/examples/train/configs/fine_tuning/kandinsky5/README.md index 5b13d0ed2..3e30f07f5 100644 --- a/examples/train/configs/fine_tuning/kandinsky5/README.md +++ b/examples/train/configs/fine_tuning/kandinsky5/README.md @@ -62,9 +62,30 @@ during either training stage. ## Train stage 2 (QAT-aware DMD distillation) +Stage 1 writes a raw DCP checkpoint (`checkpoint-/dcp` + metadata/RNG +state) under `training.checkpoint.output_dir` +(`outputs/kandinsky5_t2v_qat_finetune/` by default) -- `models.*.init_from` +needs a diffusers model directory instead (`model_index.json` + component +subfolders), so convert it first: + +```bash +python -m fastvideo.train.entrypoint.dcp_to_diffusers \ + --checkpoint outputs/kandinsky5_t2v_qat_finetune/checkpoint- \ + --output-dir outputs/kandinsky5_t2v_qat_finetune/checkpoint--diffusers \ + --role student --verify +``` + +`--role student` exports the trained transformer weights (the only role +that matters here: student/teacher/critic in stage 2 all load from the same +checkpoint, and `_loading_teacher_critic_model` handles the full-precision +masking for teacher/critic at load time, not at export time). `--verify` +strictly reloads the exported transformer immediately, so a key-mapping bug +fails here instead of deep inside the stage-2 launch below. Then point +`models.student/teacher/critic.init_from` in `dmd2_t2v_480p_qat.yaml` at +that `checkpoint--diffusers` directory (already the placeholder value in +the checked-in config) and launch: + ```bash -# point models.student/teacher/critic.init_from in dmd2_t2v_480p_qat.yaml -# at the stage-1 checkpoint first bash examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh ``` @@ -90,15 +111,35 @@ their default layer lists (`fastvideo/layers/quantization/nvfp4_qat_config.py`, `fastvideo/layers/quantization/fp8_config.py`), so no `target_layers` override is needed: +`dcp_to_diffusers` copies the base T2V checkpoint's `model_index.json` +unchanged into every export (see "Train stage 2" above), so `_class_name` +still says the base T2V pipeline and the registry +(`fastvideo/registry.py`) has no way to auto-detect that a given directory +is actually a DMD (four-step re-noise sampler) export -- it will resolve to +`Kandinsky5T2VPipeline` and run the full-length sampler on DMD-distilled +weights. Pass `override_pipeline_cls_name` and a `Kandinsky5DMDConfig` +explicitly to select the right pipeline/sampler and get +`dmd_denoising_steps` set (`Kandinsky5T2VConfig`'s default of `None` makes +`Kandinsky5DmdDenoisingStage` raise immediately): + ```python from fastvideo import VideoGenerator +from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig from fastvideo.layers.quantization import get_quantization_config gen = VideoGenerator.from_pretrained( "path/to/kandinsky5_dmd_checkpoint", num_gpus=1, + override_pipeline_cls_name="Kandinsky5DMDPipeline", + pipeline_config=Kandinsky5DMDConfig(), transformer_quant=get_quantization_config("nvfp4_qat")(), use_fsdp_inference=False, ) +# Kandinsky5DmdDenoisingStage ignores request.sampling.num_inference_steps +# and guidance_scale -- it's a fixed 4-step, no-CFG sampler driven entirely +# by pipeline_config.dmd_denoising_steps (set above). Adjust +# Kandinsky5DMDConfig(dmd_denoising_steps=[...]) instead if the checkpoint +# was distilled/validated with a different schedule than the default +# [1000, 750, 500, 250]. gen.generate(request={"prompt": "...", "output": {"save_video": True}}) ``` @@ -114,7 +155,14 @@ to a supported dense backend while the FP4 linear layers still run. - `fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py` -- confirms `ATTN_QAT_TRAIN` actually engages (not a silent SDPA fallback) and that gradients flow through a forward+backward pass. +- `fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py` -- + confirms an unmodified export resolves to `Kandinsky5T2VPipeline` (the bug) + and that `override_pipeline_cls_name="Kandinsky5DMDPipeline"` fixes it (the + documented workaround above), plus `Kandinsky5DMDConfig`'s + `dmd_denoising_steps` default. - `fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py` -- a few - steps of both training stages on a single synthetic sample. No golden - reference video exists yet (unlike the Wan e2e test); once a real training - run is available, snapshot one and add an SSIM assertion to match. + steps of both training stages on a single synthetic sample, including the + `dcp_to_diffusers --verify` conversion between them (the same command + documented above). No golden reference video exists yet (unlike the Wan + e2e test); once a real training run is available, snapshot one and add an + SSIM assertion to match. diff --git a/examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh b/examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh index 7a8cf113c..0a194b622 100755 --- a/examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh +++ b/examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh @@ -15,6 +15,12 @@ # # Data: preprocess with fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py # (or an equivalent parquet dataset matching pyarrow_schema_t2v) first. +# +# Next: this writes a raw DCP checkpoint (checkpoint-N/dcp + metadata/RNG +# state), not something stage 2 can load directly. Convert it with +# fastvideo.train.entrypoint.dcp_to_diffusers before running +# ../../distribution_matching/kandinsky5/distill_dmd_qat.sh -- see that +# script's header for the exact command. set -euo pipefail export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN # <-- enables Attn-QAT training diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 374cb1931..19f8217d6 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -6,7 +6,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.hyworld import HYWorldConfig -from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig +from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig @@ -21,6 +21,6 @@ __all__ = [ "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", - "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", - "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name" + "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", + "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/kandinsky5.py b/fastvideo/configs/pipelines/kandinsky5.py index fba5ad44e..047704dbf 100644 --- a/fastvideo/configs/pipelines/kandinsky5.py +++ b/fastvideo/configs/pipelines/kandinsky5.py @@ -120,3 +120,21 @@ class Kandinsky5I2VConfig(Kandinsky5T2VConfig): super().__post_init__() # I2V needs the VAE encoder to encode the conditioning image. self.vae_config.load_encoder = True + + +@dataclass +class Kandinsky5DMDConfig(Kandinsky5T2VConfig): + """Kandinsky-5.0 DMD (few-step distilled) text-to-video pipeline configuration. + + Checkpoints exported by ``fastvideo.train.entrypoint.dcp_to_diffusers`` + copy their base T2V checkpoint's ``model_index.json`` unchanged, so + ``_class_name`` still says the base T2V pipeline and the registry + (``fastvideo/registry.py``) cannot auto-detect a DMD export -- pass this + config together with ``override_pipeline_cls_name="Kandinsky5DMDPipeline"`` + explicitly to ``VideoGenerator.from_pretrained``/``from_config`` (see + ``examples/train/configs/fine_tuning/kandinsky5/README.md``). Without it, + ``Kandinsky5T2VConfig``'s ``dmd_denoising_steps=None`` makes + ``Kandinsky5DmdDenoisingStage`` raise immediately. + """ + + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250]) diff --git a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py index 4a3f1ec17..22c6d77a9 100644 --- a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py @@ -100,6 +100,29 @@ def load_video(path: str, num_frames: int, height: int, width: int) -> torch.Ten return video +def get_caption(item: dict) -> str: + """Extract a caption from one ``videos2caption.json`` entry. + + The documented schema (``docs/training/data_preprocess.md``) stores + ``cap`` as a plain string; some producers (e.g. this repo's own e2e + test fixtures, mirroring ``preprocess_hunyuan_overfit.py``) instead + store a non-empty list of caption variants and use the first one. + Indexing unconditionally with ``item["cap"][0]`` silently takes the + first *character* of a string caption instead of erroring, so accept + and validate both forms explicitly here. + """ + cap = item.get("cap") + if isinstance(cap, str): + if not cap: + raise ValueError(f"Empty 'cap' string for entry: {item}") + return cap + if isinstance(cap, list): + if not cap or not isinstance(cap[0], str) or not cap[0]: + raise ValueError(f"'cap' list must be non-empty with a non-empty first string, got: {item}") + return cap[0] + raise ValueError(f"'cap' must be a string or a list of strings, got {type(cap).__name__}: {item}") + + def main() -> None: # FastVideo's native text-encoder layers (e.g. CLIPAttention's # QKVParallelLinear) are tensor-parallel-aware and assert the TP process @@ -145,6 +168,12 @@ def main() -> None: ).to(device).eval() qwen_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer")) qwen_tok_kwargs = dict(pipeline_config.text_encoder_configs[0].tokenizer_kwargs) + # TextEncodingStage overrides max_length from pipeline_config.text_encoder_max_lengths + # at runtime (fastvideo/pipelines/stages/text_encoding.py) -- Reason1Config's static + # tokenizer_kwargs default (text_len=512) doesn't account for the 129-token Kandinsky + # system template ENCODE_START_IDX strips off afterwards, so mirror the runtime value + # here or training conditions on fewer caption tokens than inference does. + qwen_tok_kwargs["max_length"] = pipeline_config.text_encoder_max_lengths[0] print("Loading CLIP text encoder...") clip_enc = TextEncoderLoader().load( @@ -153,12 +182,13 @@ def main() -> None: ).to(device).eval() clip_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer_2")) clip_tok_kwargs = dict(pipeline_config.text_encoder_configs[1].tokenizer_kwargs) + clip_tok_kwargs["max_length"] = pipeline_config.text_encoder_max_lengths[1] # --- Process each video --- records = [] for item in caption_data: video_name = item["path"] - caption = item["cap"][0] + caption = get_caption(item) video_path = os.path.join(DATA_DIR, "videos", video_name) print(f"\nProcessing: {video_name}") @@ -243,7 +273,7 @@ def main() -> None: # Write validation prompts for callback. # Wrap in "data" key -- ValidationDataset expects field="data". # Use "caption" field -- ValidationDataset aliases it to "prompt". - val_prompts = {"data": [{"caption": item["cap"][0]} for item in caption_data]} + val_prompts = {"data": [{"caption": get_caption(item)} for item in caption_data]} val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json") with open(val_path, "w") as f: json.dump(val_prompts, f, indent=2) diff --git a/fastvideo/pipelines/stages/kandinsky5.py b/fastvideo/pipelines/stages/kandinsky5.py index 6429b9c49..717eec613 100644 --- a/fastvideo/pipelines/stages/kandinsky5.py +++ b/fastvideo/pipelines/stages/kandinsky5.py @@ -10,7 +10,10 @@ import torch from diffusers.utils.torch_utils import randn_tensor from tqdm.auto import tqdm +import fastvideo.envs as envs +from fastvideo.attention import LocalAttention from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder +from fastvideo.attention.selector import backend_name_to_enum from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context @@ -27,6 +30,7 @@ from fastvideo.pipelines.stages.decoding import DecodingStage from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.platforms import AttentionBackendEnum from fastvideo.utils import PRECISION_TO_TYPE logger = init_logger(__name__) @@ -169,6 +173,96 @@ class Kandinsky5DenoisingStage(PipelineStage): seq_len = int(mask.sum(1).max().item()) return torch.arange(seq_len, device=device) + # Backends whose selection is all-or-nothing: fastvideo.platforms.cuda + # raises ImportError immediately if the kernel isn't available, rather + # than the generic "requested backend unsupported by this layer, falling + # back to automatic selection" path _cached_get_attn_backend takes for + # most other backend names (see fastvideo/attention/selector.py). Both + # values are also always in Kandinsky5ArchConfig's + # _supported_attention_backends, so that generic fallback never applies + # to them here. That combination makes an exact backend match a safe, + # false-positive-free signal *only* for these two -- any other + # FASTVIDEO_ATTENTION_BACKEND value (e.g. a backend meant for a + # different model family sharing the same process/env var) can + # legitimately resolve to something other than what was requested. + _STRICT_BACKENDS = frozenset({ + AttentionBackendEnum.ATTN_QAT_TRAIN, + AttentionBackendEnum.ATTN_QAT_INFER, + }) + + def _assert_local_attention_backend_engaged(self) -> None: + """Guard against ``LocalAttention`` silently falling back to SDPA. + + ``Kandinsky5Attention.forward`` catches the ``AssertionError`` + ``LocalAttention`` raises when no pipeline forward context is set and + falls back to plain ``F.scaled_dot_product_attention`` -- silently + skipping whichever kernel ``FASTVIDEO_ATTENTION_BACKEND`` requested + (the fake-quantized ``ATTN_QAT_TRAIN``/``ATTN_QAT_INFER`` kernels). + ``LocalAttention.backend`` is resolved once at module construction, + independent of whether a forward context is later set, so this only + catches a backend that failed to resolve to what the env var + requested -- it does not prove forward context was present for a + given forward call. Pair it with always wrapping the actual + transformer call in ``set_forward_context`` (see the denoising loops + below), which is what prevents the runtime fallback. + + Scoped to ``_STRICT_BACKENDS``: other backend names can legitimately + differ from what ``LocalAttention.backend`` resolved to (silent, + logged, intentional fallback when unsupported by this layer -- see + ``_STRICT_BACKENDS``'s comment), so asserting on those would flag + expected behavior as a bug. + """ + backend_env = envs.FASTVIDEO_ATTENTION_BACKEND + if not backend_env: + return + expected = backend_name_to_enum(backend_env) + if expected is None or expected not in self._STRICT_BACKENDS: + return + for module in self.transformer.modules(): + if isinstance(module, LocalAttention): + assert module.backend == expected, ( + f"Kandinsky5 local attention resolved to backend {module.backend}, expected " + f"{expected} from FASTVIDEO_ATTENTION_BACKEND={backend_env}. This likely means " + "LocalAttention's missing-forward-context guard silently fell back to SDPA.") + return + + def _resolve_target_dtype(self, fastvideo_args: FastVideoArgs) -> torch.dtype: + """Resolve the transformer's actual compute dtype. + + Trust ``pipeline_config.dit_precision`` directly for a normal, + non-FSDP load -- ``TransformerLoader.load()`` asserts every + parameter matches it exactly (``fastvideo/models/loader/ + component_loader.py``), so this holds for standalone T2V/I2V + inference regardless of which precision (including fp32) was + requested. + + FSDP2-wrapped modules are the one case where that assertion doesn't + carry forward: ``maybe_load_fsdp_model`` hardcodes its + ``MixedPrecisionPolicy`` to ``param_dtype=torch.bfloat16`` for every + FSDP-wrapped load, independent of ``dit_precision`` -- this covers + both the live transformer ``ValidationCallback`` reuses from + training (whose ``pipeline_config.dit_precision`` reflects the + fp32 master-weight load dtype, not the actual bf16 compute dtype) + and multi-GPU ``use_fsdp_inference=True`` runs. Detect the real + compute dtype from the parameters only in that case, mirroring the + Cosmos 2.5 denoising stage. + """ + declared_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] + try: + from torch.distributed.fsdp import FSDPModule + except Exception: # pragma: no cover - FSDP not always available + return declared_dtype + + if not isinstance(self.transformer, FSDPModule): + return declared_dtype + + target_dtype = torch.bfloat16 # safe default + for p in self.transformer.parameters(): + if p.dtype != torch.float32: + target_dtype = p.dtype + break + return target_dtype + @staticmethod def fast_sta_nabla( T: int, @@ -270,20 +364,10 @@ class Kandinsky5DenoisingStage(PipelineStage): loader = TransformerLoader() self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) fastvideo_args.model_loaded["transformer"] = True + self._assert_local_attention_backend_engaged() device = get_local_torch_device() - # FSDP-wrapped models (e.g. the live transformer reused by - # ValidationCallback during training) may have their true compute - # dtype disagree with pipeline_config.dit_precision -- the training - # config forces dit_precision to "fp32" (master-weight load dtype) - # even though the actual FSDP2 mixed-precision policy runs the model - # in bf16. Detect the real weight dtype instead, mirroring the - # Cosmos 2.5 denoising stage. - target_dtype = torch.bfloat16 # safe default - for p in self.transformer.parameters(): - if p.dtype != torch.float32: - target_dtype = p.dtype - break + target_dtype = self._resolve_target_dtype(fastvideo_args) autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast latents = batch.latents num_channels = getattr( @@ -444,15 +528,10 @@ class Kandinsky5DmdDenoisingStage(Kandinsky5DenoisingStage): loader = TransformerLoader() self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) fastvideo_args.model_loaded["transformer"] = True + self._assert_local_attention_backend_engaged() device = get_local_torch_device() - # See the parent stage's forward() for why this can't be - # PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]. - target_dtype = torch.bfloat16 # safe default - for p in self.transformer.parameters(): - if p.dtype != torch.float32: - target_dtype = p.dtype - break + target_dtype = self._resolve_target_dtype(fastvideo_args) autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast latents = batch.latents num_channels = getattr( @@ -499,9 +578,18 @@ class Kandinsky5DmdDenoisingStage(Kandinsky5DenoisingStage): continue t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype) + attn_metadata = None + if sparse_params is not None: + attn_metadata = NablaAttentionMetadataBuilder().build( + current_timestep=i, + sta_mask=sparse_params["sta_mask"], + P=sparse_params["P"], + visual_shape=sparse_params["visual_shape"], + ) autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled) if device.type == "cuda" else contextlib.nullcontext()) - with autocast_ctx: + with set_forward_context(current_timestep=i, attn_metadata=attn_metadata, + forward_batch=batch), autocast_ctx: pred_velocity = self.transformer( hidden_states=latents.to(dtype=target_dtype), encoder_hidden_states=prompt_embeds, diff --git a/fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py b/fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py new file mode 100644 index 000000000..2d161f4e7 --- /dev/null +++ b/fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression test for exported Kandinsky5 DMD checkpoints resolving to the +wrong pipeline. + +``fastvideo.train.entrypoint.dcp_to_diffusers`` copies the base T2V +checkpoint's ``model_index.json`` unchanged into every export (see +``examples/train/configs/fine_tuning/kandinsky5/README.md``), so +``_class_name`` still says the base T2V pipeline. Without an explicit +override, ``fastvideo.registry.get_model_info`` resolves such a directory to +``Kandinsky5T2VPipeline`` (via the path-based "t2v" fallback detector in +``fastvideo/registry.py``) instead of ``Kandinsky5DMDPipeline`` -- silently +running the full-length multi-step sampler on DMD-distilled weights instead +of the four-step re-noise sampler. ``override_pipeline_cls_name`` bypasses +that resolution entirely and must be paired with ``Kandinsky5DMDConfig`` (not +the default ``Kandinsky5T2VConfig``, whose ``dmd_denoising_steps=None`` makes +``Kandinsky5DmdDenoisingStage`` raise immediately). +""" + +import json + +from fastvideo.configs.pipelines.kandinsky5 import ( + Kandinsky5DMDConfig, + Kandinsky5T2VConfig, +) +from fastvideo.pipelines.basic.kandinsky5.kandinsky5_dmd_pipeline import Kandinsky5DMDPipeline +from fastvideo.pipelines.basic.kandinsky5.kandinsky5_pipeline import Kandinsky5T2VPipeline +from fastvideo.pipelines.pipeline_registry import PipelineType +from fastvideo.registry import get_model_info + + +def _write_fake_export(tmp_path, name: str): + """Minimal on-disk stand-in for a dcp_to_diffusers export: just enough + for verify_model_config_and_directory to accept it.""" + model_dir = tmp_path / name + (model_dir / "transformer").mkdir(parents=True) + (model_dir / "model_index.json").write_text( + json.dumps({ + "_class_name": "Kandinsky5T2VPipeline", + "_diffusers_version": "0.30.0", + })) + return model_dir + + +def test_kandinsky5_dmd_config_sets_denoising_steps(): + assert Kandinsky5T2VConfig().dmd_denoising_steps is None + assert Kandinsky5DMDConfig().dmd_denoising_steps == [1000, 750, 500, 250] + + +def test_exported_dmd_checkpoint_resolves_to_t2v_pipeline_without_override(tmp_path): + """Documents the bug: an unmodified export resolves to the T2V pipeline.""" + model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat") + info = get_model_info(model_path=str(model_dir), pipeline_type=PipelineType.BASIC) + assert info.pipeline_cls is Kandinsky5T2VPipeline + + +def test_override_pipeline_cls_name_resolves_dmd_pipeline(tmp_path): + model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat_override") + info = get_model_info( + model_path=str(model_dir), + pipeline_type=PipelineType.BASIC, + override_pipeline_cls_name="Kandinsky5DMDPipeline", + ) + assert info.pipeline_cls is Kandinsky5DMDPipeline diff --git a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py index fd4048def..40f98972f 100644 --- a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py +++ b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py @@ -49,6 +49,14 @@ KANDINSKY5_OVERFIT_DATA_DIR = DATA_DIR / "kandinsky5_overfit" PREPROCESSED_DIR = DATA_DIR / "kandinsky5_overfit_preprocessed" STAGE1_OUTPUT_DIR = DATA_DIR / "outputs" / "kandinsky5_e2e_stage1" STAGE2_OUTPUT_DIR = DATA_DIR / "outputs" / "kandinsky5_e2e_stage2" +# Kept out of STAGE1_OUTPUT_DIR: that directory is scanned with +# glob("checkpoint-*") below to find the latest raw DCP checkpoint, and a +# "checkpoint--diffusers" export dir sitting alongside "checkpoint-" +# would also match that pattern -- on a rerun that doesn't start from a +# clean data/ dir, it would sort after "checkpoint-" and get picked as +# stage1_ckpt instead, feeding an already-exported diffusers directory back +# into dcp_to_diffusers as if it were a DCP checkpoint. +STAGE1_DIFFUSERS_DIR = DATA_DIR / "outputs" / "kandinsky5_e2e_stage1_diffusers" STAGE1_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "fine_tuning" / "kandinsky5" / "t2v_480p_qat.yaml") @@ -101,6 +109,29 @@ def _run_preprocessing() -> None: subprocess.run(cmd, cwd=str(REPO_ROOT), check=True) +def _export_dcp_to_diffusers(checkpoint_dir: Path, output_dir: Path) -> None: + """Convert a stage-1 DCP checkpoint to a diffusers-style directory. + + Stage 2's ``models.*.init_from`` (like real-training YAML configs, see + ``dmd2_t2v_480p_qat.yaml``) expects a diffusers model directory + (``model_index.json`` + component subfolders), not the raw + ``checkpoint-N`` DCP layout ``CheckpointManager`` writes (``dcp/`` + + metadata/RNG state only). ``--verify`` strictly reloads the exported + transformer immediately so a key-mapping bug fails here, at the export + boundary, rather than deep inside the stage-2 launch below. + """ + cmd = [ + sys.executable, "-m", + "fastvideo.train.entrypoint.dcp_to_diffusers", + "--checkpoint", str(checkpoint_dir), + "--output-dir", str(output_dir), + "--role", "student", + "--overwrite", + "--verify", + ] + subprocess.run(cmd, cwd=str(REPO_ROOT), check=True) + + def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_overrides: list[str], env_overrides: dict[str, str]) -> None: cmd = [ @@ -109,7 +140,7 @@ def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_ov "--nproc_per_node", NUM_GPUS, "-m", "fastvideo.train.entrypoint.train", "--config", str(config), - "--training.data.data_path", str(PREPROCESSED_DIR / "combined_parquet_dataset"), + "--training.data.data_path", str(PREPROCESSED_DIR), "--training.distributed.num_gpus", NUM_GPUS, "--training.distributed.hsdp_shard_dim", NUM_GPUS, "--training.loop.max_train_steps", str(max_train_steps), @@ -145,16 +176,21 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample(): f"no stage-1 checkpoint produced under {STAGE1_OUTPUT_DIR}") stage1_checkpoints = sorted(STAGE1_OUTPUT_DIR.glob("checkpoint-*")) - stage1_ckpt = str(stage1_checkpoints[-1]) + stage1_ckpt = stage1_checkpoints[-1] + + stage1_diffusers_dir = STAGE1_DIFFUSERS_DIR / stage1_ckpt.name + _export_dcp_to_diffusers(stage1_ckpt, stage1_diffusers_dir) + assert (stage1_diffusers_dir / "model_index.json").exists(), ( + f"dcp_to_diffusers export did not produce a diffusers model dir at {stage1_diffusers_dir}") _run_stage( STAGE2_CONFIG, STAGE2_OUTPUT_DIR, max_train_steps=3, extra_overrides=[ - "--models.student.init_from", stage1_ckpt, - "--models.teacher.init_from", stage1_ckpt, - "--models.critic.init_from", stage1_ckpt, + "--models.student.init_from", str(stage1_diffusers_dir), + "--models.teacher.init_from", str(stage1_diffusers_dir), + "--models.critic.init_from", str(stage1_diffusers_dir), ], env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"}, ) diff --git a/fastvideo/tests/stages/test_kandinsky5_target_dtype.py b/fastvideo/tests/stages/test_kandinsky5_target_dtype.py new file mode 100644 index 000000000..fa9ce6dcc --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky5_target_dtype.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression test: Kandinsky5 denoising stages must honor +``pipeline_config.dit_precision`` for a plain (non-FSDP) transformer. + +``Kandinsky5DenoisingStage._resolve_target_dtype`` scans parameters and +defaults to bf16 when every one is fp32 -- a heuristic needed only for the +FSDP2-wrapped transformer ``ValidationCallback`` reuses from training +(``maybe_load_fsdp_model`` hardcodes its ``MixedPrecisionPolicy`` to +``param_dtype=torch.bfloat16`` independent of ``dit_precision``, so the +live training module's true compute dtype can disagree with what +``dit_precision`` claims). Applying that same scan unconditionally broke +plain inference: ``TransformerLoader.load()`` asserts every parameter +matches ``dit_precision`` exactly for a non-FSDP load, so an explicit +fp32 pipeline (all params fp32, loop never finds a non-fp32 param) fell +through to the "safe default" of bf16 -- silently running fp32-requested +inference in bf16 instead. + +This is a pure logic test on a plain ``torch.nn.Module`` stand-in (never an +``FSDPModule``) -- no GPU, no model load, no distributed init needed. +""" +from __future__ import annotations + +import types + +import torch + +from fastvideo.pipelines.stages.kandinsky5 import Kandinsky5DenoisingStage + + +def _make_stage(transformer: torch.nn.Module) -> Kandinsky5DenoisingStage: + """Bypass __init__: only set the field _resolve_target_dtype reads.""" + stage = Kandinsky5DenoisingStage.__new__(Kandinsky5DenoisingStage) + stage.transformer = transformer + return stage + + +def _fastvideo_args(dit_precision: str) -> types.SimpleNamespace: + return types.SimpleNamespace(pipeline_config=types.SimpleNamespace(dit_precision=dit_precision)) + + +def test_resolve_target_dtype_honors_explicit_fp32_for_plain_transformer(): + transformer = torch.nn.Linear(4, 4).to(torch.float32) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("fp32")) + + assert resolved == torch.float32, ( + "an explicit fp32 pipeline_config must not be silently cast to bf16 for a " + "plain (non-FSDP) transformer") + + +def test_resolve_target_dtype_honors_explicit_fp16_for_plain_transformer(): + transformer = torch.nn.Linear(4, 4).to(torch.float16) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("fp16")) + + assert resolved == torch.float16 + + +def test_resolve_target_dtype_honors_explicit_bf16_for_plain_transformer(): + transformer = torch.nn.Linear(4, 4).to(torch.bfloat16) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("bf16")) + + assert resolved == torch.bfloat16 diff --git a/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py b/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py index c6c62dbc3..2681ce8d0 100644 --- a/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py +++ b/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py @@ -16,6 +16,13 @@ Two silent-fallback traps this guards against: the defensive pattern in ``fastvideo/models/loader/component_loader.py``. +This test itself must not become that stale-selection source for whichever +model test runs next in the same pytest process: ``FASTVIDEO_ATTENTION_BACKEND`` +is set via ``monkeypatch`` (auto-restored at teardown, including on +failure) and the cache is cleared again in a ``finally`` block, since +``monkeypatch`` only restores the env var, not the separate process-wide +cache keyed on it. + Unlike ``fastvideo.platforms.cuda``'s other backends, ATTN_QAT_TRAIN has no silent-fallback path at all if the kernel isn't built -- backend selection raises ``ImportError`` loudly. This test skips (does not fail) when the @@ -58,7 +65,7 @@ _FIXTURE = str( @pytest.mark.usefixtures("distributed_setup") -def test_kandinsky5_attn_qat_train_engages_and_backprops(): +def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch): if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -67,82 +74,89 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(): if not is_attn_qat_train_available(): pytest.skip("fastvideo_kernel ATTN_QAT_TRAIN kernel is not built") - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_TRAIN" + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN") _cached_get_attn_backend.cache_clear() + try: + from fastvideo.train.models.kandinsky5 import Kandinsky5Model + from fastvideo.train.utils.config import load_run_config - from fastvideo.train.models.kandinsky5 import Kandinsky5Model - from fastvideo.train.utils.config import load_run_config + cfg = load_run_config(_FIXTURE) + model = Kandinsky5Model( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=True, + ) - cfg = load_run_config(_FIXTURE) - model = Kandinsky5Model( - init_from=cfg.models["student"]["init_from"], - training_config=cfg.training, - trainable=True, - ) + device = torch.device("cuda:0") + dtype = torch.bfloat16 + transformer = model.transformer.to(device=device, dtype=dtype) - device = torch.device("cuda:0") - dtype = torch.bfloat16 - transformer = model.transformer.to(device=device, dtype=dtype) + attn = transformer.visual_transformer_blocks[0].self_attention + assert isinstance(attn.local_attention, LocalAttention) + assert attn.local_attention.backend == AttentionBackendEnum.ATTN_QAT_TRAIN, ( + f"expected ATTN_QAT_TRAIN, got {attn.local_attention.backend} -- " + "backend selection silently fell back") - attn = transformer.visual_transformer_blocks[0].self_attention - assert isinstance(attn.local_attention, LocalAttention) - assert attn.local_attention.backend == AttentionBackendEnum.ATTN_QAT_TRAIN, ( - f"expected ATTN_QAT_TRAIN, got {attn.local_attention.backend} -- " - "backend selection silently fell back") + arch = transformer.config.arch_config + patch_size = arch.patch_size + in_visual_dim = arch.in_visual_dim + in_text_dim = arch.in_text_dim + in_text_dim2 = arch.in_text_dim2 - arch = transformer.config.arch_config - patch_size = arch.patch_size - in_visual_dim = arch.in_visual_dim - in_text_dim = arch.in_text_dim - in_text_dim2 = arch.in_text_dim2 + grid_t, grid_h, grid_w = 2, 4, 4 + latent_t = grid_t * patch_size[0] + latent_h = grid_h * patch_size[1] + latent_w = grid_w * patch_size[2] - grid_t, grid_h, grid_w = 2, 4, 4 - latent_t = grid_t * patch_size[0] - latent_h = grid_h * patch_size[1] - latent_w = grid_w * patch_size[2] + latents = torch.randn( + 1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype, + requires_grad=True) + if bool(getattr(transformer, "visual_cond", False)): + # See Kandinsky5Model._build_distill_input_kwargs / + # Kandinsky5LatentPreparationStage: visual_cond=True checkpoints + # always expect [real | zero_cond | zero_mask] concatenated on the + # channel dim. + cond = torch.zeros_like(latents) + mask = torch.zeros(*latents.shape[:-1], 1, device=device, dtype=dtype) + hidden_states = torch.cat([latents, cond, mask], dim=-1) + else: + hidden_states = latents + encoder_hidden_states = torch.randn(1, 8, in_text_dim, device=device, dtype=dtype) + pooled_projections = torch.randn(1, in_text_dim2, device=device, dtype=dtype) + timestep = torch.tensor([500], device=device, dtype=dtype) + visual_rope_pos = [ + torch.arange(grid_t, device=device), + torch.arange(grid_h, device=device), + torch.arange(grid_w, device=device), + ] + text_rope_pos = torch.arange(encoder_hidden_states.shape[1], device=device) - latents = torch.randn( - 1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype, - requires_grad=True) - if bool(getattr(transformer, "visual_cond", False)): - # See Kandinsky5Model._build_distill_input_kwargs / - # Kandinsky5LatentPreparationStage: visual_cond=True checkpoints - # always expect [real | zero_cond | zero_mask] concatenated on the - # channel dim. - cond = torch.zeros_like(latents) - mask = torch.zeros(*latents.shape[:-1], 1, device=device, dtype=dtype) - hidden_states = torch.cat([latents, cond, mask], dim=-1) - else: - hidden_states = latents - encoder_hidden_states = torch.randn(1, 8, in_text_dim, device=device, dtype=dtype) - pooled_projections = torch.randn(1, in_text_dim2, device=device, dtype=dtype) - timestep = torch.tensor([500], device=device, dtype=dtype) - visual_rope_pos = [ - torch.arange(grid_t, device=device), - torch.arange(grid_h, device=device), - torch.arange(grid_w, device=device), - ] - text_rope_pos = torch.arange(encoder_hidden_states.shape[1], device=device) + with set_forward_context(current_timestep=0, attn_metadata=None): + out = transformer( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + visual_rope_pos=visual_rope_pos, + text_rope_pos=text_rope_pos, + scale_factor=(1.0, 2.0, 2.0), + sparse_params=None, + return_dict=True, + ).sample - with set_forward_context(current_timestep=0, attn_metadata=None): - out = transformer( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - pooled_projections=pooled_projections, - timestep=timestep, - visual_rope_pos=visual_rope_pos, - text_rope_pos=text_rope_pos, - scale_factor=(1.0, 2.0, 2.0), - sparse_params=None, - return_dict=True, - ).sample + assert torch.isfinite(out).all().item(), "output contains NaN/Inf" - assert torch.isfinite(out).all().item(), "output contains NaN/Inf" - - out.sum().backward() - assert latents.grad is not None - assert torch.isfinite(latents.grad).all().item(), "input grad contains NaN/Inf" - assert attn.to_query.weight.grad is not None - assert torch.isfinite(attn.to_query.weight.grad.to_local() - if hasattr(attn.to_query.weight.grad, "to_local") else - attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf" + out.sum().backward() + assert latents.grad is not None + assert torch.isfinite(latents.grad).all().item(), "input grad contains NaN/Inf" + assert attn.to_query.weight.grad is not None + assert torch.isfinite(attn.to_query.weight.grad.to_local() + if hasattr(attn.to_query.weight.grad, "to_local") else + attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf" + finally: + # monkeypatch restores FASTVIDEO_ATTENTION_BACKEND itself, but the + # selector cache is a separate process-wide functools.cache keyed on + # (head_size, dtype, supported_backends) -- not on the env var -- so + # it must be cleared again here or a later model test in this same + # pytest process could reuse this test's ATTN_QAT_TRAIN selection. + _cached_get_attn_backend.cache_clear() diff --git a/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py b/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py new file mode 100644 index 000000000..c8ae97f3f --- /dev/null +++ b/fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression test: preprocess_kandinsky5_overfit.py must not silently +truncate a string caption to its first character. + +The documented ``videos2caption.json`` schema +(``docs/training/data_preprocess.md``) stores ``cap`` as a plain string +(e.g. ``"Ocean waves at sunset..."``), while some producers (this repo's +own e2e test fixtures, mirroring ``preprocess_hunyuan_overfit.py``) store a +non-empty list of caption variants instead. Indexing unconditionally with +``item["cap"][0]`` silently takes the first *character* of a string caption +("O") instead of erroring -- ``get_caption`` accepts and validates both +forms once at ingestion. + +Pure logic test -- no GPU, no model load. +""" +from __future__ import annotations + +import pytest + +from fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit import get_caption + + +def test_get_caption_accepts_plain_string(): + assert get_caption({"path": "a.mp4", "cap": "Ocean waves at sunset..."}) == "Ocean waves at sunset..." + + +def test_get_caption_accepts_nonempty_list(): + assert get_caption({"path": "a.mp4", "cap": ["a synthetic test clip"]}) == "a synthetic test clip" + + +def test_get_caption_uses_first_list_element(): + assert get_caption({"path": "a.mp4", "cap": ["first", "second"]}) == "first" + + +@pytest.mark.parametrize("bad_cap", [ + "", + [], + [""], + None, + 123, + {"nested": "dict"}, +]) +def test_get_caption_rejects_invalid_values(bad_cap): + with pytest.raises(ValueError): + get_caption({"path": "a.mp4", "cap": bad_cap}) diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py index 0c57ec335..da9d8e0b5 100644 --- a/fastvideo/train/entrypoint/dcp_to_diffusers.py +++ b/fastvideo/train/entrypoint/dcp_to_diffusers.py @@ -20,6 +20,10 @@ The checkpoint must contain ``metadata.json`` (written by ``CheckpointManager``). If the checkpoint predates metadata support, pass ``--config`` explicitly to provide the training YAML. + +Pass ``--verify`` to strictly reload the exported transformer immediately +after writing it, so a key-mapping bug fails here instead of deep inside a +later training/inference launch that loads the exported directory. """ from __future__ import annotations @@ -200,6 +204,27 @@ def _save_role_pretrained( return str(dst) +def _strict_reload_verify(*, output_dir: str, training_config: Any) -> None: + """Reload the just-exported transformer from disk and fail loudly on + any key mismatch. + + ``TransformerLoader.load()`` already loads strictly (``strict=True``) + for every non-Cosmos25 model, so this doesn't add leniency -- it moves + the failure to right after export, with a clear "the export is broken" + error, instead of surfacing deep inside a later training/inference + launch that happens to load this directory. + """ + from fastvideo.train.utils.moduleloader import load_module_from_path + + logger.info("Verifying export: strictly reloading transformer from %s", output_dir) + load_module_from_path( + model_path=output_dir, + module_type="transformer", + training_config=training_config, + ) + logger.info("Strict reload verification passed.") + + def convert( *, checkpoint_dir: str, @@ -207,6 +232,7 @@ def convert( config_path: str | None = None, role: str = "student", overwrite: bool = False, + verify: bool = False, ) -> str: """Load a DCP checkpoint and export as a diffusers model. @@ -298,6 +324,10 @@ def convert( model=model, ) logger.info("Export complete: %s", result) + + if verify: + _strict_reload_verify(output_dir=result, training_config=tc) + return result @@ -406,6 +436,13 @@ def main() -> None: action="store_true", help="Overwrite output-dir if it exists.", ) + parser.add_argument( + "--verify", + action="store_true", + help=("After exporting, strictly reload the transformer from " + "the exported directory to catch key-mapping bugs " + "immediately."), + ) args = parser.parse_args(sys.argv[1:]) convert( @@ -414,6 +451,7 @@ def main() -> None: config_path=args.config, role=args.role, overwrite=args.overwrite, + verify=args.verify, )