o 9j@sxdZddlmZddlZddlZddlZddlmZddlm Z ddl Z ddl Z ddl mmZddlmZddlmZddlmZmZdd lmZmZmZd d lmZd d lmZd d l m!Z!zddl"Z"dZ#Wn e$ytdZ#Ynwee%&j'dZ(e(ddZ)e(ddZ*e(dddZ+d3ddZ,d4d!d"Z-d3d#d$Z.d5d&d'Z/ (d6d7d.d/Z0d0d1Z1e2d2kre1dSdS)8uBV3 face training — distill teacher's .npz targets into a causal TCN. Usage: PYTHONPATH=. python3 -m models.v3_face.train # full train PYTHONPATH=. python3 -m models.v3_face.train --smoke # 10 scenarios × 5 epochs PYTHONPATH=. python3 -m models.v3_face.train --device cuda:1 --epochs 80 ) annotationsN)Path)Dict)AdamW)LambdaLR) DataLoaderWeightedRandomSampler) LIPSYNC_ONLYEXPRESSION_ONLYSHARED_CHANNELS) V3FaceConfig)BlendshapeDataset) V3FaceModelTFdata v3_trainingemotionmodelsv3_face checkpointscfgr return torch.TensorcCsbtj|jf|jtjd}tD]}|j||<qtD] }d|j|j||<q|j|d<|j|d<|S)zPer-channel L1 weight vector, shape (52,). Lipsync channels get cfg.lipsync_weight (audio-sync matters most). Eye-blink (ch 8, 9) get cfg.eye_blink_weight (sparse, hard to learn). Everything else gets cfg.expression_weight. dtype? ) torchfull output_dimexpression_weightfloat32r lipsync_weightr eye_blink_weightrwchr)L/dataset/kemix-engine/package/face/animasync-face-v3/models/v3_face/train.pymake_channel_weights+s   r+predtarget valid_length ch_weightsc Cs|j\}}}tj||jdd}||dk}|||ddddf} | jdd} | |} |j dd} | | S) uL1 loss masked to valid frames, weighted per channel. pred, target: (B, T, C) valid_length: (B,) int — number of valid frames per sample ch_weights: (C,) devicerr Ndim?min) shaperaranger1 unsqueezefloatabsmeansumclamp) r,r-r.r/BTC frame_idxmaskdiffdenomr)r)r* masked_l1?s   rGcCstj|jf|jtjd}tD]}|j||<qtD]}|j||<q|j |d<|j |d<|j dur9|j |d<|j |d<|j durHdD]}|j ||<q@|S)aPer-channel velocity penalty weights, shape (52,). Lipsync channels get heavy smoothing; brows + cheek get light smoothing (preserve V2 prosody motion); eye-blink channels get near-zero so the sharp 5-frame blink kernel survives training. rrrN)rr r) rr r!velocity_expression_weightr#r velocity_lipsync_weightr velocity_shared_weightvelocity_eye_blink_weightvelocity_eye_squint_weightvelocity_brow_weightr&r)r)r*make_velocity_weightsRs          rR vel_weightsc Cs|j\}}}|ddddf|ddddf}|ddddf|ddddf}|||ddddf} | jdd} tj|d|jdd} | |ddk} | | } | j dd} | | S) zEPer-channel-weighted L1 on per-frame difference (smoothness penalty).Nr r2r3r0rr5r6) r8r<r=rr9r1r:r;r>r?) r,r-r.rSr@rA_pred_vtarget_vrErCrDrFr)r)r*masked_velocityrs ((  rWr5trainboolvelocity_scaler;Dict[str, float]c Cs||ddddd} |rtnt} | |D]} | dj|dd} | dj|dd}| dj|dd}| d j|dd}|| |}t||||}t||||| }||}|ry|jdd |tj j | |j ||| d t|7<| d t|7<| d |jdkrt|nd7<| dd7<qWdn1swYtd| d}| d || d || d |dS)Nr)lossl1velnaudioT) non_blockingcondr-r.) set_to_noner]r^r_r`r )r]r^r_)rXr enable_gradno_gradtorGrW zero_gradbackwardnnutilsclip_grad_norm_ parameters grad_clipstepr;itemvelocity_weightmax)modeldl optimizer schedulerr/rSrr1rXrZsumsgrad_ctxbatchrarcr-r.r,loss_l1loss_velr]r`r)r)r* run_epochs:   &$r|c$s6 t}|jdttd|jdttd|jdttd|jdtdd|jdtdd|jdtdd|jdd d |jd td d|jd ddd|jdgdddd|jdtddd|jdddd|jdddd|jdddd |jd!dd"d |jd#dd$d |jd%td&d'd|jd(td&d)d|jd*tdd+d|jd,tdd-d|jd.tdd/d|jd0td1d2d|jd3td1d4d|jd5td1d6d|jd7td1d8d|jd9tdd:d|jd;tddd|jd?dd@d|jdAtd1dBd|jdCtdDdEd|jdFtddGd|jdHtddId|jdJt ddKd| }t |j r7|j _ |jr?|j_|jrG|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j_|j _ |j!_!|j"j#dLdLdMt$%|j%}jd&ksjd&ksjdurʈjdurjnj}t&dNjdOdPjdOdQ|dOdRjdurt&dSjdOdTj'dOdUjdurt&dVjdOdTj'dOdUjd1krt&dWjdOdXjd1kr"t&dYjdOdZjd[j(d\d]j(d^jd1kr@t&d_jdOdZjd[j(d\d]j(d^jd1kr^t&d`jdOdZjd[j(d\d]j(d^jdurst&dajdOdbjdOdcjdurt&ddjdOdejdOdfjdurt&dgjdOdbjdurjnjdOdcjdurt&dhjdOdijdurt&djjdOdkjrt&dlj dOdmj!d1krt&dnj!dodZj!d[j(d\d]j(d^t)|j*|j+dpfidqj,drjdsjdtjdujdvjdwjdxjdyjdzjd{jd|jd}jd~jdj dj!}t)|j*|j+dfj,jjjjjjjjjjjjj j!d}|j-rg}dD]\}fdd|j.Dd|}|/|qo|s|j.dd}||_.|j.dd|_.d_ d_0d _t&dt1|dt1|dj djt&dt1|d|2t&dt1|d|2|3j4} t5| t1|dLd} t6|j| |j7dLdd} t6|jd|j7dLd} t89|} t&d| j:ddOd| j;dod| j;d dod|jd|j<t$j?|j<|dd}| @|dt&d|j<d|Addd|AdtdddU|jBr| B}t&d|ddOd| jCddOdddlDmE}|D]}q}|jFotG}|jFrtGst&d|r|jHpdjIdt1jJt1jKdjd|j-rdnd}tFjL|jM|jN|ijO| j:| j;| j;d t1|t1||2|2t ||j-d dtFjP| dddt&dtFjQRdd| SD}tT|jdjUd}tVdt1| j dfdd }tW||}tX9|}tY9|}ddlDmE}mZ}|j[dkrR|D] }d1||<d1||<q;t&dt1|dn+|j[dks\|jBr}|D] }d1||<d1||<q^t&d|jBrrdnddt1|dƝtdǃ}t\j D]}t]]}j^dkrt_d&|j^nd&}t`| | |||||dL|dɍ }t`| | |||||d|dɍ } t]]|}!t&d|d˛d|ddd|ddd| ddd| ddd|addқd|dOd|!dod՝| d|k}"|"r| d}| bjO|| d| d|d|dd֜}|j[dkrdnd|j[}#|jcr+|#d|jc}#|"rGt$d||j"d|#dٝt&d|#d|ddUt$d||j"d|#dٝ|rytFje||d|d|d| d| d| d||ad|!dޜ |dߍqt&d|dt&d|j"|r|tFjfd<tFgdSdS)Nz --npz_dir)typedefaultz --emotion_dirz --out_dirz--epochsz --batch_sizez--lrz--devicezcuda:0)r~z --num_workersrKz--smoke store_truez)Smoke test: 10 train scenarios, 5 epochs.)actionhelpz--focus)alllipsync expressionrzChannel focus. 'all' = train all 52 channels (default). 'lipsync' = only LIPSYNC + SHARED channels have loss (expression branch sees zero gradient). 'expression' = only EXPRESSION_ONLY channels have loss (lipsync branch sees zero gradient).)choicesr~rz--resumezvLoad model weights from a checkpoint .pt before training. Use with --freeze_lipsync for phase-2 expression retraining.)r}r~rz--freeze_lipsynczFreeze shared backbone + lipsync branch + lipsync head. Only the expression branch + head will train. Lipsync output stays bit-for-bit identical to what was loaded via --resume.z--wandbz5Log to Weights & Biases. Requires `wandb login` once.z--wandb_projectzanimasync-v3-facezW&B project name.)r~rz--wandb_run_namez)W&B run name. Defaults to auto-generated.z--wandb_entityz4W&B entity (team or user). Defaults to your default.z--lipsync-target-gainr5zMultiply PURE_LIPSYNC target channels (jaw, mouth mechanics, tongue, cheekPuff) by this factor at load time, then clamp [0,1]. Default 1.0 = no change.z--expression-target-gainaRMultiply EMOTIONAL target channels (brows incl. innerUp, cheekSquint, eyeSquint, eyeWide, mouth Dimple/Frown/Smile, noseSneer) by this factor at load time, then clamp [0,1]. Default 1.0 = no change. If --emotional-mouth-target-gain is also set, THIS knob covers only the pure-expression subset (brows + eyeSquint + eyeWide + cheekSquint).z--emotional-mouth-target-gainuOptional separate gain for emotional-mouth target channels (mouthDimple, mouthFrown, mouthSmile, noseSneer). When set, decouples from --expression-target-gain so brows/eyes can go higher than mouth. Used by v18b to avoid pushing shared mouth channels past the point where crisp_mouth normalization destabilizes lipsync. None → same as --expression-target-gain (backward-compat with v14/v18).z--velocity-eye-squint-weightuPer-channel velocity penalty for eyeSquint L/R (ch 18, 19). Higher than the default expression velocity weight suppresses jitter at high gain. None → use velocity_expression_weight (backward-compat). Suggested 0.8 for v18b.z--velocity-brow-weightuPer-channel velocity penalty for the 5 brow channels (ch 0-4). Same mechanism as the eyeSquint weight — suppresses brow jitter at high expression gain. None → use velocity_expression_weight (backward-compat). Suggested 0.7 for v18c.z--plosive-damp-targetr\a-Bake the runtime plosive damper into training targets. When mouthClose > 0.4 on a frame, mouthPress/Roll/Shrug (ch 35,36,39,40,41,42) get multiplied by (1 - this_value * smoothstep). 0 = off (default). 0.30 = matches the production main-viewer setting that prevents 'lips swallowed' on m/b/p plosives.z--smooth-target-sigma-browuGaussian σ (frames @ 30 fps) for pre-smoothing the 5 brow target channels BEFORE the gain. The proper fix for brow flicker at high gain — smooths the input the model is asked to fit so jitter never enters the training signal. 0 = off (default). Suggested 2.0 (~67ms) for v18e.z--smooth-target-sigma-eye-wideuGaussian σ (frames @ 30 fps) on the eyeWide target channels (20, 21) before gain. Pair with eyeSquint smoothing to fix surprise → other-emotion transitions where eyeWide↓ and eyeSquint↑ cross discontinuously. Typical: same value as eyeSquint.z --smooth-target-sigma-eye-squintuGaussian σ (frames @ 30 fps) for pre-smoothing the eyeSquint L/R target channels BEFORE the gain. Real orbicularis oculi is slow + sustained, so we can smooth heavier than brows without losing useful motion. 0 = off (default). Suggested 3.0 (~100ms) for v18e.z--brow-innerup-happy-gainuqOverride the gain on browInnerUp (ch 2) for happy emotions (joy / laughter / excitement / gratitude). Set to 1.0 with expression-target-gain=2.2 to keep browInnerUp un-amplified on happy frames so the avatar doesn't look concerned/apologetic when saying happy things, while leaving the 2.2× boost everywhere else. None → no override (backward compat with v14..v18f).z--brow-happy-gainuPer-emotion gain override for ALL brow channels (ch 0–4) on happy frames (joy / laughter / excitement / gratitude). Use when brows are pushed very high via --brow-target-gain but should stay calm (~v14 level 1.4) on smiling frames. None → no override (backward compat).z--brow-target-gainu.Override target gain for ALL brow channels (ch 0–4) independently of --expression-target-gain. Use to keep brows at a v14-style mild gain (e.g. 1.4) while pushing the rest of expression (cheekSquint / eyeSquint / eyeWide) much higher. None → brows ride the global expression gain (backward compat).z --soft-clipaReplace the dataset's hard upper-clamp at 1.0 with a smooth knee-then-asymptote curve. Preserves emotional cross-fade tails that would otherwise be chopped off when high target gains push values well past 1.0. Linear below --soft-clip-knee, asymptotic above. Lower bound stays hard 0.z--smooth-cond-sigma-emotionuGaussian σ (in frames @ 30 fps) for smoothing the emotion one-hot in cond[:, :16] along time. Smooths turn-boundary step changes so the model learns gradual emotion cross-fades. Must also be applied at inference time. Typical: 5–10 (~170–330 ms). 0 = no smoothing (backward compat).z--soft-clip-kneegffffff?zKnee point for --soft-clip. Below this value the curve is exactly linear (no distortion). Above, it bends toward 1.0. Default 0.7 = first 70%% of the dynamic range untouched. Range (0, 1).z--eye-wide-maxuHard cap on eyeWide (ch 20, 21) AFTER gain, before the [0,1] clamp. Prevents bug-eyed saturation at high expression gain. Typical: 0.5–0.7. None → no cap (backward compat).z--brow-surprise-gainu$Per-frame brow scale-down on surprise/fluster proxy (high eyeWide). Mirrors the viewer 'brow cap' slider: brows (ch 0–4) get multiplied by (1 - wide_ramp * value), where wide_ramp ramps in over [0.10, 0.30] of post-gain eyeWide. Typical: 0.3–0.5. None → no scale-down (backward compat).z --variant-taguOptional suffix appended to checkpoint filenames after --focus, so different gain runs don't clobber each other. E.g. 'v14' → best_lipsync_v14.pt / best_expression_v14.pt.T)parentsexist_oku[target-gain] lipsync×z.2fu expression×u emotional-mouth×z6 (applied per-channel in dataset, then clamped [0,1])u [velocity] eyeSquint(ch 18,19)×z (vs default expression )u[velocity] brows(ch 0-4)×z[plosive-damper] target damp=z, (applied in dataset when mouthClose > 0.4)u [target-smooth] brow(ch 0-4) σ=z frames (~iz.0fz ms @ z fps)u'[target-smooth] eyeSquint(ch 18,19) σ=u%[target-smooth] eyeWide(ch 20,21) σ=z:[per-emotion] browInnerUp(ch 2) on happy frames uses gain z (vs z elsewhere)u[target-gain] brows(ch 0-4)×u (overrides expression×z for brow group only)z5[per-emotion] brows(ch 0-4) on happy frames use gain u6[per-frame] brows(ch 0-4) scaled by (1 - wide_ramp × z?) on surprise/fluster frames (eyeWide proxy, ramp [0.10, 0.30])z"[cap] eyeWide(ch 20,21) capped at z after gainz[soft-clip] saturation knee at z! (linear below, asymptotic above)u![cond-smooth] emotion one-hot σ=z.1fzseed_train_final.jsonl crop_frameslipsync_target_gainexpression_target_gainemotional_mouth_target_gainplosive_damp_targetsmooth_target_sigma_browsmooth_target_sigma_eye_squintsmooth_target_sigma_eye_widebrow_innerup_happy_gainbrow_happy_gainbrow_target_gainbrow_surprise_gain eye_wide_max soft_clipsoft_clip_kneesmooth_cond_sigma_emotionzseed_val.jsonl)rrrrrrrrrrrrrrr))long_rK)solo_rJ)z daily_-splitrJcsg|] }|dkr|qS)r r)).0e)want_catr)r* szmain.. z[smoke] train=z val=z epochs=z batch=ztrain: u scenarios — counts: zval: ) num_samples replacementF) batch_sizesampler num_workers pin_memory drop_last)rshufflerrzV3FaceModel (split-branch): g.Az M params, ~z MB fp32 (~z MB int8)z--resume checkpoint not found: ) map_location weights_onlyrsz[resume] loaded weights from z (prev epoch=epoch?z , val_l1=val_l1nanz.4fz[freeze_lipsync] froze zGM params (shared backbone + lipsync branch + lipsync head). Trainable: zM (expression branch + head).r )LIPSYNC_BRANCH_CHANNELSuY[wandb] requested but `wandb` not installed — skipping. Install with: pip install wandbv3face_h_b_lrz.0e_smoke) n_params size_mb_fp32 size_mb_int8train_scenarios val_scenariostrain_category_countsval_category_countsr1smoke)projectentitynameconfig gradientsd)loglog_freqz[wandb] logging to cSsg|]}|jr|qSr)) requires_grad)rpr)r)r*rs)g?gffffff?)lrbetas weight_decayrointrr;csR|jkr |tdjS|jtdj}ddttjtd|S)Nr rr5) warmup_stepsrrmathcospir7)roprog)r total_stepsr)r* lr_lambdas zmain..lr_lambda)rEXPRESSION_BRANCH_CHANNELSrz[focus=lipsync] masked z expression channelsrz[focus=expressionz/freeze_lipsyncz ] masked z lipsync channelsinfr)rXrZzepoch 3dz train l1=r^z vel=r_z val l1=z lr=z.2ez vw=z s)rsrrrval_veltrain_l1 train_velrTbestz.ptu → saved bestz (val l1=latestr]) rztrain/l1ztrain/velocityz train/losszval/l1z val/velocityzval/lossz val/best_l1r epoch_seconds)roz Done. best val l1: z checkpoints: best_val_l1)rorrr;)hargparseArgumentParser add_argumentrDEFAULT_NPZ_DIRDEFAULT_EMOTION_DIRDEFAULT_OUT_DIRrr;str parse_argsr epochsn_epochsrr learning_raterrrrPrQrrrrrrrrrrrrout_dirmkdirrr1printrLfpsrnpz_dir emotion_dirrrentriesextendrlencategory_countsget_sample_weightslong_oversample_weightrrrrrgrsize_mbresumeexists SystemExitloadload_state_dictgetfreeze_lipsync n_trainablersrwandb_WANDB_AVAILABLEwandb_run_name hidden_dimshared_dilationsbranch_dilationsinit wandb_project wandb_entity__dict__watchrunget_urlrmrrrrrr+rRrfocusrangetimevelocity_warmup_epochsr7r| get_last_lr state_dict variant_tagsaversummaryfinish)$apargsr1mouth_gtrain_dsval_ds smoke_picksr`pickssample_weightsrtrain_dlval_dlrsckptfrozen_nrr( use_wandbrun_nametrainable_paramsrurrvr/rSrbest_valrt0rZtrvadtis_bestsuffixr))rrrr*mainsB                                                                           r*__main__)rr rr) r,rr-rr.rr/rrr) r,rr-rr.rrSrrr)r5)rXrYrZr;rr[)3__doc__ __future__rrrr pathlibrtypingrnumpynprtorch.nn.functionalrj functionalF torch.optimrtorch.optim.lr_schedulerrtorch.utils.datarrscripts.compiler.constantsr r r rr datasetrrsrrr ImportError__file__resolver PROJECT_ROOTrrrr+rGrRrWr|r*__name__r)r)r)r*sP                %[