ducvh5/looptrans-a5-b5x4-c5-tf-fm-tm

LoopTrans A5 + B5 x 4 + C5: TF + FM + teacher matching and distill continuation

This repository contains the 13,000-step, 16,500-step, 21,500-step, and 21,750-step transformer checkpoints from looped_a5_b5x4_c5_vpred_tf_fm_logit_normal_fulltrain. Steps 13,000, 16,500, and 21,500 include both EMA and non-EMA (student) weights. Step 21,750 contains EMA weights only, saved on a new best validation VideoReward result outside the regular checkpoint interval.

Also included: stage-two distill-only step 2,000, initialized from the stage-one 21,750 EMA checkpoint, with both EMA and non-EMA (student) weights. The 2,000 count is local to this continuation stage (23,750 updates along the lineage).

Also included: 13,000-step EMA-only checkpoints from both Qwen text experiments: logit-normal(0,1) with shift 2, and independent uniform with shift 5. Both start from the pretrained teacher and use TF + FM + teacher matching. Each training draw chooses 50% VBench-style / 50% detailed Qwen text-only recaptions; held-out validation uses the original captions. These releases contain EMA weights only.

The student has 15 unique transformer blocks arranged as A5 -> (B5 repeated 4 times) -> C5, for 30 transformer-block applications per denoising evaluation. The shared B blocks use trainable per-loop modulation tables. Each variant contains 722,570,560 stored parameters in FP32 safetensors format.

These are transformer-only weights. Load the VAE, text encoder, tokenizer, and scheduler from Wan2.1-T2V-1.3B-Diffusers. The custom LoopedWanTransformer3DModel implementation is included as looped_wan.py.

All eleven weights are preserved without dtype conversion. Each weight file is approximately 2.89 GB. release_manifest.json includes SHA-256 checksums.

Training timestep sampling uses u = sigmoid(z), z ~ Normal(0,1) for logit-normal, or u ~ Uniform(0,1) for uniform. Apply the rational shift once: sigma = s*u / (1 + (s-1)*u), then pass t = 1000*sigma to the transformer. For logit-normal this equals sigma = sigmoid(z + log(s)). Draws are independent (no batch stratification). Both Qwen runs use logit-normal(0,1), shift 2 for validation loss, and the video generation recipe below uses inference shift 8.

VideoReward validation

All student checkpoint scores below use EMA weights only. The teacher uses its original pretrained weights. These are saved validation results on the same 32 held-out captions, with one generated video per caption and seeds 1234–1265. Teacher and original draft results come from the project's eval/outputs/validation32_* runs; this repository's checkpoint results come from training validation.

Model	Training step	Training timestep sampling	Weights	Overall ↑	VQ ↑	MQ ↑	TA ↑
Wan2.1-T2V-1.3B teacher	—	Not recorded here (pretraining)	Pretrained	2.3185	0.2548	0.7337	1.3301
Ours: TF + FM + teacher matching	13,000	Logit-normal(μ=0, σ=1), rational shift 2; independent	EMA	1.8325	0.2506	0.8487	0.7332
Ours: TF + FM + teacher matching, Qwen text	13,000	Logit-normal(μ=0, σ=1), rational shift 2; independent	EMA	1.4372	0.2164	0.7065	0.5143
Ours: TF + FM + teacher matching, Qwen text	13,000	Uniform(0,1), rational shift 5; independent	EMA	1.3272	0.0946	0.6270	0.6056
Ours: TF + FM + teacher matching	16,500	Logit-normal(μ=0, σ=1), rational shift 2; independent	EMA	1.8361	0.2560	0.9099	0.6702
Original draft distillation	20,000	Uniform(0,1), rational shift 5; independent	EMA	1.5650	0.0239	0.7139	0.8272
Ours: TF + FM + teacher matching	21,500	Logit-normal(μ=0, σ=1), rational shift 2; independent	EMA	2.1331	0.4105	0.9838	0.7388
Ours: TF + FM + teacher matching	21,750	Logit-normal(μ=0, σ=1), rational shift 2; independent	EMA	2.1682	0.3977	1.0028	0.7677
Ours: stage 2, distill-only	21,750 EMA + 2,000 distill	Logit-normal(μ=0, σ=1), rational shift 2; independent	EMA	2.2637	0.4233	1.0232	0.8172
Original draft distillation	30,800	Uniform(0,1), rational shift 5; independent	EMA	1.8049	0.0797	0.7684	0.9569

Generation uses 50 UniPC steps, shift 8, CFG 6, 81 frames, 480×832, and 16 FPS, with the same negative prompt. Transformer weights are FP32 with BF16 autocast, text embeddings are BF16, and VAE decoding is FP32. All loop students use four loops.

Scoring uses KlingTeam/VideoReward revision 4f26600130683e6f1de9f5d463887f28e8ef995c, BF16/SDPA, batch size 1, 2-fps frame sampling, a 200,704-pixel per-frame budget, and scoring seeds 0–31. Scores are arithmetic means over all 32 videos. VQ, MQ, and TA denote normalized visual quality, motion quality, and text alignment; Overall = VQ + MQ + TA. Higher is better; scores are signed rewards, not percentages. This comparison covers the fixed validation set and is not a full VBench benchmark.

Training objective

The stage-one checkpoint metadata records flow matching, final-output teacher velocity matching, and hidden-state teacher forcing on the B loop blocks only. Their recorded loss coefficients are 1 each. Training uses logit-normal sigma sampling with logit mean 0, standard deviation 1, and rational shift 2. EMA decay is 0.999. Shared initialization is recorded as 1. Full checkpoint-specific metadata is provided beside each checkpoint's model folders.

Stage two: distill-only continuation. Source experiment: looped_a5_b5x4_c5_distill_only_from_tf_fm_step0021750_ema_fulltrain. Initialized from stage-one 21,750 EMA and trained for 2,000 additional updates. Only teacher velocity matching is optimized (teacher_lambda=1, student_all): each of the four B-loop exits is decoded through shared C blocks and the output head, then matched to the full frozen teacher. Raw exit weights are 0.3333,0.3333,0.3333,1, normalized to approximately 1/6,1/6,1/6,1/2. FM, hidden-state TF, and hidden-state AR loss weights are zero in this stage. Training retains logit-normal mean 0, std 1, shift 2, caption dropout 0.1, and EMA decay 0.999. Transformer parameters are trained in full. The validation table reports the stage-two EMA scores only, using the same 32-caption generation/scoring protocol as stage one and the teacher.

The uploaded weights do not contain optimizer state, source videos, cached training latents, or local dataset indices.

Loading

Use the provided custom class rather than the standard Wan transformer class. First download the implementation and the selected model folder, then import looped_wan.py from the downloaded repository:

import sys
import torch
from huggingface_hub import snapshot_download

repo = "ducvh5/looptrans-a5-b5x4-c5-tf-fm-tm"
step = 21750  # or 13000 / 16500 / 21500
variant = "ema"  # "student" is available at 13000 / 16500 / 21500 only
folder = f"step{step:07d}/{variant}"
# For Qwen text EMA-only step 13,000, use either:
# folder = "qwen_text_logit_normal2/step0013000/ema"
# folder = "qwen_text_uniform5/step0013000/ema"
# For stage-two distill step 2,000 instead, use:
# folder = f"distill_from_step0021750/step0002000/{variant}"  # ema or student

snapshot = snapshot_download(
    repo,
    allow_patterns=["looped_wan.py", f"{folder}/*"],
)
sys.path.insert(0, snapshot)
from looped_wan import LoopedWanTransformer3DModel

transformer = LoopedWanTransformer3DModel.from_pretrained(
    snapshot,
    subfolder=folder,
    torch_dtype=torch.float32,
    use_safetensors=True,
).eval()
transformer.inference_loops = 4

The source environment uses Diffusers 0.35.1 and Transformers 4.49.0. The project's validation recipe uses 50 UniPC steps, sampling shift 8, guidance scale 6, 81 frames at 480x832, 16 FPS, and seed 1234 for the first prompt. Training shifts 2/5 and inference shift 8 serve different roles.

The weights are intermediate research checkpoints. Validation scores above cover EMA students and the pretrained teacher on the fixed 32-caption protocol.

Base model and license

This model is derived from the Apache-2.0 licensed Wan2.1-T2V-1.3B model. The recorded Diffusers base revision is 0fad780a534b6463e45facd96134c9f345acfa5b. Credit to the Wan team for the pretrained model and base implementation.