anyflow: split co-train sampling cfg, factor time_shift and train_mode

Author: juliusberner

fastgen/configs/data.py

@@ -125,3 +125,78 @@
     shuffle=True,
     sampler_start_idx=None,
 )
+
+# ################################################################################
+# VidProM / Wan 480p (832x480, 81 frames)
+# ################################################################################
+# WDS shards of VidProM prompts rendered with Wan2.1. Two flavors:
+#   * ``vidprom_wan_480p_latents``  — precomputed Wan-VAE latents [16, 21, 60, 104]
+#     plus UMT5 text embeddings [512, 4096]; no VAE/T5 needed at train time
+#     (set ``config.model.enable_preprocessors = False``).
+#   * ``vidprom_wan_res720_ar16_9_f121_clean`` — the raw 720p mp4 source shards.
+# Both live under ``$VIDPROM_ROOT_DIR`` (default ``$DATA_ROOT_DIR/vidprom_wan``).
+VIDPROM_ROOT_DIR = os.getenv("VIDPROM_ROOT_DIR", f"{DATA_ROOT_DIR}/vidprom_wan")
+_VIDPROM_CLEAN_DIR = f"{VIDPROM_ROOT_DIR}/vidprom_wan_res720_ar16_9_f121_clean"
+
+# Curation lists (json: {shard: [fname, ...]}) applied to both the raw and the
+# latent shards so the two stay index-consistent.
+VidProM_Wan_Ignore_Index_Paths = [
+    # captions containing an <img tag
+    f"{_VIDPROM_CLEAN_DIR}/ignore_index/img_tag_23.json",
+    # lowest 5% by avg. rank over HPSv2, ClipScore, VideoAlign VQ/MQ/TA, InternVideo2
+    f"{_VIDPROM_CLEAN_DIR}/ignore_index/lowest_5_perc_weighted_rank_25473.json",
+    # nsfw captions
+    f"{_VIDPROM_CLEAN_DIR}/ignore_index/nsfw_375.json",
+    # split-screen captions
+    f"{_VIDPROM_CLEAN_DIR}/ignore_index/split_screen_606.json",
+    # prompts carrying "ar" / "gs" / "seed" / "neg" commands
+    f"{_VIDPROM_CLEAN_DIR}/ignore_index/prompt_cmd_4507.json",
+    # samples absent from the latent shards (keeps raw and latent in sync)
+    f"{_VIDPROM_CLEAN_DIR}/ignore_index/missing_files_165.json",
+]
+
+# Shared across the raw and latent loaders. Video-only options (sequence_length,
+# img_size) are not listed here: WDSLoader/BaseWDSLoader reject them, only
+# VideoWDSLoader accepts them.
+VidProM_Wan_Loader_Kwargs = dict(
+    batch_size=2,
+    num_workers=2,
+    deterministic=False,
+    sampler_start_idx=None,
+    shard_count_file=f"{_VIDPROM_CLEAN_DIR}/shard_count.json",
+    ignore_index_paths=VidProM_Wan_Ignore_Index_Paths,
+)
+
+_VIDPROM_NEG_PROMPT_EMB = f"{VIDPROM_ROOT_DIR}/vidprom_wan_480p_latents/neg_prompt_emb.npy"
+
+# Precomputed Wan-VAE latents + UMT5 text embeddings.
+VidProM_Wan_Latent_Loader_Config = L(WDSLoader)(
+    datatags=[f"WDS:{VIDPROM_ROOT_DIR}/vidprom_wan_480p_latents"],
+    **VidProM_Wan_Loader_Kwargs,
+    key_map={"real": "latent.pth", "condition": "txt_emb.pth"},
+    files_map={"neg_condition": _VIDPROM_NEG_PROMPT_EMB},
+)
+
+# Raw 720p mp4 shards (needs the VAE + text encoder, i.e. enable_preprocessors=True).
+VidProM_Wan_Loader_Config = L(VideoWDSLoader)(
+    datatags=[f"WDS:{_VIDPROM_CLEAN_DIR}"],
+    **VidProM_Wan_Loader_Kwargs,
+    sequence_length=81,
+    img_size=(832, 480),
+    key_map={"real": "mp4"},
+    presets_map={"neg_condition": "neg_prompt_wan", "condition": "empty_string"},
+)
+
+# Validation shards: 10 hand-crafted prompts in shard 0, 10 random VBench prompts
+# per shard afterwards, with Wan-1.3B (cfg=5, shift=5, skip_layers=[10]) renders
+# as the reference "real" sample.
+VidProM_Wan_Latent_Val_Loader_Config = L(WDSLoader)(
+    datatags=[f"WDS:{VIDPROM_ROOT_DIR}/validation_latent"],
+    batch_size=1,
+    num_workers=1,
+    train=False,
+    sampler_start_idx=None,
+    shards_per_worker=1,
+    key_map={"real": "latent.pth", "condition": "txt_emb.pth"},
+    files_map={"neg_condition": _VIDPROM_NEG_PROMPT_EMB},
+)

fastgen/configs/experiments/WanT2V/config_anyflow.py

@@ -6,19 +6,26 @@
 AnyFlow's pretrain objective is MeanFlow's with a fixed ``beta08`` per-timestep
 weighting, a finite-difference JVP, shifted timestep sampling, and a
 ``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config
-runs``MeanFlowModel`` directly. The values below mirror the reference recipe
+runs ``MeanFlowModel`` directly. The values below mirror the reference recipe
 ``train_wan1b_student_shift5_81f_480p_lr5e-5_6k_b32.yml``.
-Known deviations from the reference: full-rank fine-tuning instead of the paper's
-rank-256 LoRA — the same deviation applies to the on-policy stage;
+
+Known deviations from the reference, both of which also apply to the on-policy stage:
+full-rank fine-tuning instead of the paper's rank-256 LoRA; and under
+``guidance_fuse_scale`` the 1/g rescaling of dF/dt is applied only to the samples that
+kept their condition, where the reference's ``compute_central_difference`` rescales the
+whole batch. A dropped sample's fused prediction is plain ``u_uncond``, so an ungated
+1/g would regress it onto a different fixed point -- ours is gated on ``keep``.
 
 The on-policy stage (paper Stage 2) lives in ``config_anyflow_onpolicy.py``.
 """
 
 import copy
 
+
 import fastgen.configs.methods.config_mean_flow as config_mean_flow
 from fastgen.configs.data import VideoLoaderConfig
 from fastgen.configs.net import Wan_1_3B_Config
+from fastgen.methods import AnyFlowModel
 
 
 def create_config():
@@ -99,8 +106,11 @@ def create_config():
     # ------ inference / validation ------
     config.model.student_sample_type = "ode"
     config.model.student_sample_steps = 4
-    # 4-step shifted schedule shift*s/(1+(shift-1)*s) on s=linspace(1,0,5)
-    config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0]
+    # Shifted schedule under the same map the (t, r) sampling applies. The pretrain
+    # stage has a single scheduler in the reference, so it shares that shift.
+    config.model.sample_t_cfg.t_list = AnyFlowModel.rollout_t_list(
+        config.model.student_sample_steps, config.model.sample_t_cfg.shift, config.model.net.max_t
+    ).tolist()
 
     # ------ data / trainer ------
     config.dataloader_train = copy.deepcopy(VideoLoaderConfig)

fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py

@@ -17,15 +17,23 @@
 
     trainer.checkpointer.pretrained_ckpt_path=<stage1>/checkpoints/0006000.pth
 
-Known deviation from the reference: full-rank fine-tuning instead of the paper's
-rank-256 LoRA.
+Known deviations from the reference: full-rank fine-tuning instead of the paper's
+rank-256 LoRA; the noising times for the DMD gradient and the fake score are drawn
+on [0.001, 0.999] rather than the reference's [0, 1] (its ``dmd_cfg`` sets no
+``min_timestep`` / ``max_timestep``, so it clamps to the full range) -- FastGen's
+convention for the other rectified-flow Wan configs; and, in the co-trained loss, the
+1/g rescaling of dF/dt under ``guidance_fuse_scale`` is gated on the samples that kept
+their condition, where the reference's ``compute_central_difference`` rescales the whole
+batch (see ``config_anyflow.py`` and ``FlowMapLossMixin._compute_mf_loss``).
 """
 
 import copy
 
+
 import fastgen.configs.methods.config_anyflow as config_anyflow_default
 from fastgen.configs.data import VideoLoaderConfig
 from fastgen.configs.net import Wan_1_3B_Config
+from fastgen.methods import AnyFlowModel
 
 
 def create_config():
@@ -77,11 +85,13 @@ def create_config():
     config.model.guidance_scale = 4.0
 
     # DMD gradient noising time: reference `generator_loss` draws torch.rand
-    # then applies the shift -> shifted-uniform.
+    # then applies the shift -> shifted-uniform. The bounds keep the score models off
+    # the degenerate endpoints (see the deviation note above); `fake_score_sample_t_cfg`
+    # inherits them through the deepcopy below.
     config.model.sample_t_cfg.time_dist_type = "shifted"
     config.model.sample_t_cfg.shift = 5.0
-    config.model.sample_t_cfg.min_t = 0.0
-    config.model.sample_t_cfg.max_t = 1.0
+    config.model.sample_t_cfg.min_t = 0.001
+    config.model.sample_t_cfg.max_t = 0.999
 
     # Fake-score noising time: reference `discriminator_loss` draws
     # logit_normal(0, 1) then applies the same shift. This is a DIFFERENT
@@ -92,13 +102,12 @@ def create_config():
     config.model.fake_score_sample_t_cfg.train_p_std = 1.0
 
     # ------ student rollout (reference rollout_cfg) ------
+    # The rollout grid's shift comes from `cotrain_sample_t_cfg` below: the reference
+    # builds its rollout pipeline from the same `scheduler` it draws the co-trained
+    # (t, r) from, separately from the DMD noising time above.
     config.model.student_sample_type = "ode"
     config.model.student_sample_steps_list = [2, 4, 8, 16, 50]
     config.model.student_sample_steps = 4
-    # Validation schedule for `student_sample_steps = 4`: `shift * s / (1 + (shift
-    # - 1) * s)` on `s = linspace(1, 0, 5)` -- the same grid the per-NFE rollout
-    # builds. Recompute if `shift`, `max_t` or `student_sample_steps` changes.
-    config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0]
 
     # ------ co-trained Stage-1 flow-map loss (reference cotrain_forward_kl) ------
     # FastGen's VSD loss carries a 0.5 factor the reference's DMD loss does
@@ -115,9 +124,22 @@ def create_config():
     config.model.guidance_fuse_scale = 3.0
     config.model.cond_dropout_prob = 0.1
     config.model.precision_amp_jvp = "float32"
-    config.model.sample_t_cfg.flow_matching_ratio = 0.5
-    config.model.sample_t_cfg.consistency_ratio = 0.25
-    config.model.sample_t_cfg.deterministic_buckets = True
+    # (t, r) sampling for the co-trained loss, drawn from the reference's `scheduler`.
+    # Its shift also drives the student's rollout grid, so it is independent of the DMD
+    # noising-time shift above; the reference recipe sets both to 5.0.
+    config.model.cotrain_sample_t_cfg.time_dist_type = "shifted"
+    config.model.cotrain_sample_t_cfg.shift = 5.0
+    config.model.cotrain_sample_t_cfg.min_t = 0.0
+    config.model.cotrain_sample_t_cfg.max_t = 1.0
+    config.model.cotrain_sample_t_cfg.flow_matching_ratio = 0.5
+    config.model.cotrain_sample_t_cfg.consistency_ratio = 0.25
+    config.model.cotrain_sample_t_cfg.deterministic_buckets = True
+
+    # Validation schedule at `student_sample_steps` -- the same grid the per-NFE rollout
+    # builds, so it follows the co-trained sampling shift.
+    config.model.sample_t_cfg.t_list = AnyFlowModel.rollout_t_list(
+        config.model.student_sample_steps, config.model.cotrain_sample_t_cfg.shift, config.model.net.max_t
+    ).tolist()
 
     # ------ optimization (reference: AdamW lr=2e-6, betas=(0.0, 0.999), wd=0,
     # grad clip 1.0, EMA 0.99) ------

fastgen/configs/experiments/WanT2V/config_mf.py

@@ -74,12 +74,8 @@ def create_config():
     config.model.sample_t_cfg.min_t = 0.001
     config.model.sample_t_cfg.max_t = 0.999
 
-    config.dataloader_train = VideoLatentLoaderConfig
+    config.dataloader_train = copy.deepcopy(VideoLatentLoaderConfig)
     config.dataloader_train.batch_size = 1
 
-    # 480p (832x480) resolution
-    config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8)
-    config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1
-
     config.log_config.group = "wan_mf"
     return config

fastgen/configs/methods/config_anyflow.py

@@ -42,10 +42,11 @@ class ModelConfig(DMD2ModelConfig):
     ``cotrain_forward_kl``).
     """
 
-    # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the
-    # extra fields are ignored by the DMD2 noising-time sampling.
-    sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig)
-    sample_r_cfg: MeanFlowSampleRConfig = attrs.field(factory=MeanFlowSampleRConfig)
+    # MeanFlow-style (t, r) sampling for the co-trained flow-map loss. Separate from
+    # the inherited `sample_t_cfg`, which DMD2 uses for the noising time: the reference
+    # draws these from its `scheduler` and the noising time from its `dmd_scheduler`.
+    cotrain_sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig)
+    cotrain_sample_r_cfg: MeanFlowSampleRConfig = attrs.field(factory=MeanFlowSampleRConfig)
     loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig)
 
     # Weight of the co-trained Stage-1 flow-map loss in the student update.
@@ -70,7 +71,7 @@ class ModelConfig(DMD2ModelConfig):
     cond_dropout_prob: Optional[float] = None
     cond_keys_no_dropout: List[str] = attrs.field(factory=list)
 
-    # Precision for autocast in the co-trained loss JVP (None = training precision).
+    # Precision for autocast in the co-trained loss JVP (None disables autocast there).
     precision_amp_jvp: str | None = None
 
     # MeanFlow's target-side guidance knobs. AnyFlow typically guides on the

fastgen/configs/methods/config_mean_flow.py

@@ -129,7 +129,7 @@ class ModelConfig(BaseModelConfig):
     # guidance t end
     guidance_t_end: float = 1.0
 
-    # precision for autocast in JVP (none defaults to training precision)
+    # precision for autocast in JVP (none disables autocast in the JVP region)
     precision_amp_jvp: str | None = None
 
 

fastgen/configs/methods/config_scm.py

@@ -77,7 +77,7 @@ class ModelConfig(BaseModelConfig):
     # optimizer
     net_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(RAdamOptimizerConfig))
 
-    # precision for autocast in JVP (none defaults to training precision)
+    # precision for autocast in JVP (none disables autocast in the JVP region)
     precision_amp_jvp: str | None = None