Add v7 distillation FP8 sparsity DINOv3 multiview symmetry MoE stack

v7/README.md

@@ -0,0 +1,59 @@
+# Companion Forge v7 — Distilled 3D Model Stack
+
+Target: **Hugging Face Jobs `l4x1` / NVIDIA L4 (SM89, 24 GB VRAM)**.
+
+v7 keeps the validated v6.4 ONNX/TensorRT runtime as the production fallback and adds a training/optimization layer inspired by the Kimi-K3 work: progressive few-step distillation, selective FP8, 2:4 structured sparsity, DINOv3 conditioning, multi-view fusion, geometry teachers, symmetry guidance, geometry-aware MoE, and timestep caching.
+
+## Pipeline
+
+```text
+image/text
+  -> DINOv3 bridge + optional multi-view fusion
+  -> SS flow student (10 -> 4 -> 2 -> 1 steps)
+  -> SLat flow student (10 -> 4 -> 2 -> 1 steps)
+  -> full custom ONNX/TensorRT SLat DAE
+  -> MeshFlow/TRELLIS teacher losses during training only
+  -> AniGen skeleton/skin teacher
+  -> symmetry + rig guidance
+  -> rigged GLB
+```
+
+## What is implemented
+
+- `training/distill_core.py`: progressive Euler macro-step KD + endpoint consistency + MeanFlow-style interval consistency.
+- `training/distill_flow.py`: L4-friendly LoRA-then-merge training for SS/SLat flow students, stages `10→4→2→1`.
+- `training/fp8_sparse.py`: selective ModelOpt FP8 hooks and exact magnitude-based 2:4 pruning for eligible Linear weights.
+- `training/dinov3_bridge.py` + `distill_dinov3.py`: DINOv3 ViT-S bridge to AniGen's fixed `1374×1024` conditioning contract, distilled against DINOv2 ViT-L/14-reg.
+- `runtime/multiview.py`: pose-Fourier multi-view token fusion for front/left/back/right conditioning.
+- `runtime/symmetry.py`: flow-time reflection/C2/C4 velocity symmetrization plus rig symmetry loss.
+- `training/teacher_hybrid.py`: cached MeshFlow/TRELLIS geometry teacher + AniGen rig teacher interface.
+- `training/geometry_moe.py`: ModernMOE-inspired shared + top-k routed 3D experts.
+- `runtime/timestep_cache.py`: residual timestep cache for >=4-step students; automatically unnecessary after 1–2-step distillation.
+- `bench/quality_eval.py`: geometry/symmetry/joint/speed comparison harness.
+
+## L4 training order
+
+1. Distill DINOv3 bridge (keep current DINOv2 TRT runtime as fallback).
+2. Build teacher cache using v6.4 outputs plus optional MeshFlow/TRELLIS geometry targets.
+3. Distill SS `10→4`, SLat `10→4`; validate.
+4. Apply selective FP8 PTQ; if quality drops, run short QAT.
+5. Apply 2:4 pruning to MLP/projection weights, freeze masks, recovery distillation.
+6. Distill `4→2`, then `2→1` only if geometry/rig metrics stay within thresholds.
+7. Train geometry-inductive MoE as an optional higher-capacity student; distill it back to a dense/sparse deployment student if its routing overhead is not worthwhile on L4.
+8. Re-export ONNX external-data and compile SM89 TensorRT engines. Output heads/mesh remain FP32; sparse topology ops remain the validated custom plugins.
+
+## Safety/quality gates
+
+A new student is not promoted merely because it is faster. Promotion requires finite tensors, exit code 0, GLB skin/skeleton validation, no lazy PyTorch fallback, and quality thresholds against the current v6.4 teacher. `1+1` is therefore an experimental final stage, not an automatic default.
+
+## External technology mapping
+
+- NVIDIA FastGen concepts: progressive KD / MeanFlow / consistency-style few-step training.
+- NVIDIA Model Optimizer: FP8 PTQ/QAT and export path.
+- TensorRT: 2:4 structured sparse tactics where eligible.
+- Meta DINOv3: stronger dense visual conditioning.
+- Meta MeshFlow / Microsoft TRELLIS.2: geometry/topology teacher targets only; licensing/runtime isolation is preserved.
+- SymTRELLIS idea: velocity symmetrization adapted to AniGen flow and rig constraints.
+- ModernMOE idea: routed/shared lightweight experts specialized for geometry roles.
+
+The current v6.4 runtime remains the rollback target until each v7 stage has its own L4 E2E validation artifact.

v7/bench/quality_eval.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+import argparse,json,time,torch
+
+def chamfer(a,b):
+    d=torch.cdist(a.float()[None],b.float()[None]).squeeze(0);return float(d.min(1).values.mean()+d.min(0).values.mean())
+def symmetry_x(v):
+    a=v.float().clone();a[:,0]*=-1;d=torch.cdist(v.float()[None],a[None]).squeeze(0);return float(d.min(1).values.mean())
+def joint_error(a,b):return float(torch.linalg.vector_norm(a.float()-b.float(),dim=-1).mean()) if a.shape==b.shape else None
+
+def main():
+    ap=argparse.ArgumentParser();ap.add_argument('--teacher',required=True);ap.add_argument('--student',required=True);ap.add_argument('--out',required=True);a=ap.parse_args();t=torch.load(a.teacher,map_location='cpu',weights_only=False);s=torch.load(a.student,map_location='cpu',weights_only=False)
+    r={}
+    if 'vertices' in t and 'vertices' in s:r['chamfer']=chamfer(s['vertices'],t['vertices']);r['symmetry_x']=symmetry_x(s['vertices'])
+    if 'joints' in t and 'joints' in s:r['joint_error']=joint_error(s['joints'],t['joints'])
+    if 'runtime_s' in s:r['runtime_s']=s['runtime_s'];r['speedup']=t.get('runtime_s',0)/s['runtime_s'] if t.get('runtime_s') else None
+    open(a.out,'w').write(json.dumps(r,indent=2));print(json.dumps(r))
+if __name__=='__main__':main()

v7/configs/l4_v7.yaml

@@ -0,0 +1,116 @@
+version: 7
+target:
+  gpu: NVIDIA L4
+  compute_capability: 8.9
+  flavor: l4x1
+  vram_gib: 24
+  precision_default: fp16
+  tf32: true
+
+distillation:
+  enabled: true
+  methods: [progressive_kd, meanflow, consistency]
+  stages:
+    - {teacher_steps: 10, student_steps: 4, train_steps: 6000}
+    - {teacher_steps: 4, student_steps: 2, train_steps: 8000}
+    - {teacher_steps: 2, student_steps: 1, train_steps: 12000}
+  components: [ss_flow, slat_flow]
+  train_mode: lora_then_merge
+  lora:
+    rank: 16
+    alpha: 32
+    dropout: 0.0
+    target_regex: '(qkv|to_q|to_k|to_v|proj|fc1|fc2|mlp|linear)'
+  optimizer: {name: adamw8bit, lr: 0.0001, weight_decay: 0.01}
+  gradient_checkpointing: true
+  grad_accum: 8
+  losses:
+    velocity_kd: 1.0
+    endpoint_consistency: 0.5
+    geometry_consistency: 0.25
+    skeleton_consistency: 0.25
+    symmetry: 0.10
+
+quantization:
+  modelopt: true
+  mode: fp8_selective
+  calibration_samples: 256
+  qat_steps: 1500
+  include_regex: '(qkv|to_q|to_k|to_v|proj|fc1|fc2|mlp|linear)'
+  exclude_regex: '(geo_head|skin_head|skl_head|out_layer|mesh|sparse|conv|norm)'
+  output_heads_fp32: true
+  sparse_ops_fp16: true
+
+structured_sparsity:
+  enabled: true
+  pattern: '2:4'
+  include_regex: '(fc1|fc2|mlp|proj)'
+  exclude_regex: '(qkv|geo_head|skin_head|skl_head|out_layer|mesh|sparse|conv)'
+  recovery_steps: 2000
+  fixed_mask: true
+  tensorrt_sparse_weights: true
+
+dinov3:
+  enabled: true
+  student_model: facebook/dinov3-vits16-pretrain-lvd1689m
+  teacher_model: dinov2_vitl14_reg
+  output_dim: 1024
+  target_tokens: 1374
+  input_size: 512
+  distill_steps: 5000
+  loss: {mse: 1.0, cosine: 0.5}
+  export: {onnx: true, tensorrt: true, fp8: true}
+
+multiview:
+  enabled: true
+  views: [front, left, back, right]
+  aggregation: pose_fourier_attention
+  pose_dim: 64
+  fallback_single_view: true
+
+teachers:
+  meshflow:
+    enabled: true
+    repo: facebook/meshflow
+    role: geometry_topology_teacher
+    license_gate: noncommercial_research
+  trellis2:
+    enabled: true
+    repo: microsoft/TRELLIS.2-4B
+    role: pbr_geometry_teacher
+  anigen:
+    enabled: true
+    role: skeleton_skin_teacher
+
+symmetry:
+  enabled: true
+  mode: velocity_symmetrization
+  groups: [reflection_x, c2, c4]
+  auto_detect: true
+  rig_aware:
+    bilateral_limbs: true
+    limb_length_penalty: 0.15
+    joint_mirror_penalty: 0.25
+    center_of_mass_penalty: 0.05
+
+moe:
+  enabled: true
+  mode: geometry_inductive
+  router_top_k: 2
+  shared_experts: 1
+  experts: [coarse_shape, surface_detail, limbs, head_face, skeleton, skin, materials]
+  expert_hidden_ratio: 0.5
+  residual_reuse: true
+  train_after_few_step: true
+
+cache_diffusion:
+  enabled: true
+  mode: residual_timestep_cache
+  apply_when_steps_ge: 4
+  similarity_threshold: 0.985
+  max_reuse: 2
+
+artifacts:
+  runtime_repo: patdev/Companion-Forge-L4-ONNX
+  students_repo: patdev/Companion-Forge-v7-Students
+  teacher_cache_repo: patdev/Companion-Forge-v7-Teacher-Cache

v7/runtime/multiview.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+import math,torch
+import torch.nn as nn
+
+class PoseFourier(nn.Module):
+    def __init__(self,bands=8,out=64):super().__init__();self.bands=bands;self.proj=nn.Linear(bands*4,out)
+    def forward(self,azimuth_deg,elevation_deg):
+        a=torch.deg2rad(azimuth_deg.float());e=torch.deg2rad(elevation_deg.float());freq=2**torch.arange(self.bands,device=a.device,dtype=a.dtype)
+        z=torch.cat([torch.sin(a[:,None]*freq),torch.cos(a[:,None]*freq),torch.sin(e[:,None]*freq),torch.cos(e[:,None]*freq)],-1);return self.proj(z)
+
+class MultiViewAggregator(nn.Module):
+    """Order-aware but permutation-stable view fusion for AniGen condition tokens."""
+    def __init__(self,dim=1024,pose_dim=64,heads=8):
+        super().__init__();self.pose=PoseFourier(out=pose_dim);self.pose_proj=nn.Linear(pose_dim,dim);self.attn=nn.MultiheadAttention(dim,heads,batch_first=True);self.norm=nn.LayerNorm(dim)
+    def forward(self,view_tokens,azimuth,elevation,view_mask=None):
+        # view_tokens [B,V,T,C]
+        b,v,t,c=view_tokens.shape;p=self.pose(azimuth.reshape(-1),elevation.reshape(-1)).view(b,v,-1);x=view_tokens+self.pose_proj(p)[:,:,None,:]
+        # fuse each token position across views; output remains [B,T,C]
+        q=x.mean(1);kv=x.transpose(1,2).reshape(b*t,v,c);qq=q.reshape(b*t,1,c);key_padding=None if view_mask is None else ~view_mask[:,None,:].expand(b,t,v).reshape(b*t,v)
+        y,_=self.attn(qq,kv,kv,key_padding_mask=key_padding,need_weights=False);return self.norm(q+y.reshape(b,t,c))
+
+DEFAULT_VIEWS={'front':(0.,0.),'left':(-90.,0.),'back':(180.,0.),'right':(90.,0.)}

v7/runtime/symmetry.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+import torch
+
+# Flow-time symmetry projection inspired by velocity symmetrization: transform state,
+# evaluate equivariant velocity externally, inverse-transform velocity, then average.
+
+def dense_reflect_x(x): return torch.flip(x,[-1])
+def dense_rot_z_90(x,k=1): return torch.rot90(x,k,(-2,-1))
+
+def sparse_reflect_x(s,res=64):
+    c=s.coords.clone();c[:,3]=(res-1)-c[:,3];return s.replace(coords=c)
+def sparse_rot_z_90(s,res=64,k=1):
+    c=s.coords.clone();x,y=c[:,3].clone(),c[:,2].clone();k%=4
+    if k==1:c[:,3],c[:,2]=(res-1)-y,x
+    elif k==2:c[:,3],c[:,2]=(res-1)-x,(res-1)-y
+    elif k==3:c[:,3],c[:,2]=y,(res-1)-x
+    return s.replace(coords=c)
+
+def transform_state(x,kind,res=64,inverse=False):
+    sparse=hasattr(x,'coords') and hasattr(x,'replace')
+    if kind=='identity':return x
+    if kind=='reflection_x':return sparse_reflect_x(x,res) if sparse else dense_reflect_x(x)
+    if kind.startswith('rot_z_'):
+        k=int(kind.rsplit('_',1)[1]);k=(-k)%4 if inverse else k
+        return sparse_rot_z_90(x,res,k) if sparse else dense_rot_z_90(x,k)
+    raise ValueError(kind)
+
+def group_ops(group):
+    if group=='reflection_x':return ['identity','reflection_x']
+    if group=='c2':return ['identity','rot_z_2']
+    if group=='c4':return ['identity','rot_z_1','rot_z_2','rot_z_3']
+    return ['identity']
+
+@torch.no_grad()
+def symmetrized_velocity(model,x,xs,t,cond,group='reflection_x',res=64,**kwargs):
+    vs=[];vss=[]
+    for op in group_ops(group):
+        tx=transform_state(x,op,res);ts=transform_state(xs,op,res);v,sv=model(tx,ts,t,cond,**kwargs)
+        inv=op
+        if op.startswith('rot_z_'):inv='rot_z_'+str((-int(op.rsplit('_',1)[1]))%4)
+        vs.append(transform_state(v,inv,res));vss.append(transform_state(sv,inv,res))
+    def avg(items):
+        if hasattr(items[0],'feats'):return items[0].replace(feats=torch.stack([z.feats for z in items]).mean(0))
+        return torch.stack(items).mean(0)
+    return avg(vs),avg(vss)
+
+def rig_symmetry_loss(joints,parents=None):
+    """Soft bilateral prior; assumes x axis is left/right in normalized object coordinates."""
+    if joints.numel()==0:return joints.sum()*0
+    x=joints[...,0];center=x.mean(-1,keepdim=True);sorted_x=torch.sort(x-center,dim=-1).values
+    mirror=-torch.flip(sorted_x,[-1]);return torch.nn.functional.smooth_l1_loss(sorted_x,mirror)

v7/runtime/timestep_cache.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+class ResidualTimestepCache:
+    """Training-free residual cache for multi-step flow inference.
+    Stores model residual/velocity and reuses it only when normalized state similarity is high.
+    Intended for >=4-step students; disabled automatically for 1-2 step students.
+    """
+    def __init__(self,threshold=.985,max_reuse=2):self.threshold=threshold;self.max_reuse=max_reuse;self.reset()
+    def reset(self):self.prev_state=None;self.prev_output=None;self.reuse_count=0
+    @staticmethod
+    def _feat(x):return x.feats if hasattr(x,'feats') else x
+    def similarity(self,x):
+        if self.prev_state is None:return -1.
+        a=self._feat(x).float().flatten(1);b=self._feat(self.prev_state).float().flatten(1)
+        if a.shape!=b.shape:return -1.
+        return float(F.cosine_similarity(a,b,dim=-1).mean())
+    def get(self,x):
+        sim=self.similarity(x)
+        if sim>=self.threshold and self.reuse_count<self.max_reuse:
+            self.reuse_count+=1;return self.prev_output,sim
+        return None,sim
+    def put(self,x,out):self.prev_state=x;self.prev_output=out;self.reuse_count=0
+
+class CachedFlowModel:
+    def __init__(self,model,threshold=.985,max_reuse=2):self.model=model;self.cache=ResidualTimestepCache(threshold,max_reuse);self.hits=0;self.calls=0
+    def __call__(self,x,xs,t,cond,*args,**kwargs):
+        self.calls+=1;cached,sim=self.cache.get(x)
+        if cached is not None:self.hits+=1;return cached
+        out=self.model(x,xs,t,cond,*args,**kwargs);self.cache.put(x,out);return out
+    @property
+    def hit_rate(self):return self.hits/max(1,self.calls)

v7/training/dinov3_bridge.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel
+
+class DINOv3ConditionBridge(nn.Module):
+    """DINOv3 -> AniGen conditioning contract: 1374 tokens x 1024 dims.
+    1369 spatial tokens are obtained by 32x32 -> 37x37 interpolation; 5 special tokens
+    are learned from global features so the downstream AniGen cross-attention shape is unchanged.
+    """
+    def __init__(self,model_id='facebook/dinov3-vits16-pretrain-lvd1689m',out_dim=1024,target_grid=37,special_tokens=5):
+        super().__init__();self.encoder=AutoModel.from_pretrained(model_id);h=self.encoder.config.hidden_size
+        self.proj=nn.Linear(h,out_dim,bias=False);self.special=nn.Sequential(nn.Linear(h,out_dim),nn.GELU(),nn.Linear(out_dim,special_tokens*out_dim));self.target_grid=target_grid;self.special_tokens=special_tokens;self.out_dim=out_dim
+    def forward(self,pixel_values):
+        o=self.encoder(pixel_values=pixel_values,return_dict=True);x=o.last_hidden_state
+        # DINOv3 HF exposes sequence tokens. Keep CLS/global separate and infer square patch tail.
+        g=x[:,0]; patches=x[:,1:];n=patches.shape[1];side=int(math.sqrt(n))
+        if side*side!=n:
+            # Some checkpoints expose register tokens; peel them from the front until the tail is square.
+            for r in range(min(16,n)):
+                q=n-r;s=int(math.sqrt(q))
+                if s*s==q: patches=patches[:,1+r:];side=s;break
+        b,c=patches.shape[0],patches.shape[-1];p=patches.transpose(1,2).reshape(b,c,side,side)
+        p=F.interpolate(p,size=(self.target_grid,self.target_grid),mode='bicubic',align_corners=False).flatten(2).transpose(1,2)
+        p=self.proj(p);sp=self.special(g).view(b,self.special_tokens,self.out_dim);return torch.cat([sp,p],dim=1)
+
+class FeatureDistillLoss(nn.Module):
+    def __init__(self,mse=1.0,cosine=0.5):super().__init__();self.mse=mse;self.cosine=cosine
+    def forward(self,student,teacher):
+        a=F.mse_loss(student.float(),teacher.float());b=1-F.cosine_similarity(student.float(),teacher.float(),dim=-1).mean();return self.mse*a+self.cosine*b,{'mse':a.detach(),'cosine_loss':b.detach()}

v7/training/distill_core.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, Callable
+import torch
+import torch.nn.functional as F
+
+@dataclass
+class FlowBatch:
+    x: Any
+    x_skl: Any
+    cond: Any
+    neg_cond: Any|None=None
+    kwargs: dict|None=None
+
+class StateAdapter:
+    """Math on dense tensors or AniGen SparseTensor-like objects (.feats/.replace)."""
+    @staticmethod
+    def feats(x): return x.feats if hasattr(x,'feats') else x
+    @staticmethod
+    def replace(x, feats): return x.replace(feats=feats) if hasattr(x,'replace') else feats
+    @classmethod
+    def add(cls,a,b,alpha=1.0): return cls.replace(a,cls.feats(a)+alpha*cls.feats(b))
+    @classmethod
+    def sub(cls,a,b): return cls.replace(a,cls.feats(a)-cls.feats(b))
+    @classmethod
+    def scale(cls,a,s): return cls.replace(a,cls.feats(a)*s)
+    @classmethod
+    def mse(cls,a,b): return F.mse_loss(cls.feats(a).float(),cls.feats(b).float())
+
+@torch.no_grad()
+def teacher_rollout(model, batch:FlowBatch, t_hi:float, t_lo:float, steps:int):
+    """Euler teacher trajectory from noise-side t_hi toward data-side t_lo."""
+    assert t_hi>t_lo and steps>=1
+    x,xs=batch.x,batch.x_skl; dt=(t_hi-t_lo)/steps; kw=batch.kwargs or {}
+    t=t_hi
+    for _ in range(steps):
+        tt=torch.full((StateAdapter.feats(x).shape[0],),t*1000,device=StateAdapter.feats(x).device,dtype=torch.float32)
+        v,vs=model(x,xs,tt,batch.cond,neg_cond=batch.neg_cond,**kw)
+        x=StateAdapter.add(x,v,alpha=-dt); xs=StateAdapter.add(xs,vs,alpha=-dt); t-=dt
+    return x,xs
+
+def macro_velocity(x_start,x_end,delta_t:float): return StateAdapter.scale(StateAdapter.sub(x_start,x_end),1.0/delta_t)
+
+def student_loss(student,batch:FlowBatch,t_hi:float,t_lo:float,teacher_end,weights:dict):
+    dt=t_hi-t_lo; x_end,xs_end=teacher_end; kw=batch.kwargs or {}
+    tt=torch.full((StateAdapter.feats(batch.x).shape[0],),t_hi*1000,device=StateAdapter.feats(batch.x).device,dtype=torch.float32)
+    pv,pvs=student(batch.x,batch.x_skl,tt,batch.cond,neg_cond=batch.neg_cond,**kw)
+    tv=macro_velocity(batch.x,x_end,dt); tvs=macro_velocity(batch.x_skl,xs_end,dt)
+    vel=StateAdapter.mse(pv,tv)+StateAdapter.mse(pvs,tvs)
+    pred_end=StateAdapter.add(batch.x,pv,alpha=-dt); pred_skl_end=StateAdapter.add(batch.x_skl,pvs,alpha=-dt)
+    endpoint=StateAdapter.mse(pred_end,x_end)+StateAdapter.mse(pred_skl_end,xs_end)
+    return weights.get('velocity_kd',1.0)*vel+weights.get('endpoint_consistency',0.5)*endpoint, {'velocity':float(vel.detach()),'endpoint':float(endpoint.detach())}
+
+def meanflow_pair_loss(student,batch:FlowBatch,t0:float,t1:float,teacher_end,weight:float=0.25):
+    """MeanFlow-style interval consistency: same macro displacement from both endpoints."""
+    if weight<=0:return torch.zeros((),device=StateAdapter.feats(batch.x).device)
+    x1,s1=teacher_end; mid=(t0+t1)/2; kw=batch.kwargs or {}; bsz=StateAdapter.feats(batch.x).shape[0]
+    tm=torch.full((bsz,),mid*1000,device=StateAdapter.feats(batch.x).device,dtype=torch.float32)
+    # Approximate midpoint state from teacher endpoints; stable and cheap for LoRA distillation.
+    xm=StateAdapter.scale(StateAdapter.add(batch.x,x1),0.5); sm=StateAdapter.scale(StateAdapter.add(batch.x_skl,s1),0.5)
+    vm,vsm=student(xm,sm,tm,batch.cond,neg_cond=batch.neg_cond,**kw)
+    target=macro_velocity(batch.x,x1,t0-t1); target_s=macro_velocity(batch.x_skl,s1,t0-t1)
+    return weight*(StateAdapter.mse(vm,target)+StateAdapter.mse(vsm,target_s))

v7/training/distill_dinov3.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+import argparse,os,sys,json
+from pathlib import Path
+import numpy as np,torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+from .dinov3_bridge import DINOv3ConditionBridge,FeatureDistillLoss
+
+MEAN=(.485,.456,.406);STD=(.229,.224,.225)
+def image_tensor(p,size,device):
+    im=Image.open(p).convert('RGB').resize((size,size),Image.LANCZOS);x=torch.from_numpy(np.asarray(im).copy()).permute(2,0,1).float()/255.;x=transforms.Normalize(MEAN,STD)(x);return x[None].to(device)
+
+def load_teacher(root='/home/user/app'):
+    sys.path.insert(0,root);os.chdir(root);m=torch.hub.load('./ckpts/dinov2','dinov2_vitl14_reg',pretrained=True,source='local').eval().cuda();
+    for p in m.parameters():p.requires_grad=False
+    return m
+
+@torch.no_grad()
+def teacher_features(m,x):
+    f=m(x,is_training=True)['x_prenorm'];return F.layer_norm(f,f.shape[-1:])
+
+def main():
+    ap=argparse.ArgumentParser();ap.add_argument('--images',required=True);ap.add_argument('--out',required=True);ap.add_argument('--model',default='facebook/dinov3-vits16-pretrain-lvd1689m');ap.add_argument('--steps',type=int,default=5000);ap.add_argument('--lr',type=float,default=2e-4);a=ap.parse_args()
+    files=[p for p in Path(a.images).rglob('*') if p.suffix.lower() in {'.png','.jpg','.jpeg','.webp'}];assert files
+    teacher=load_teacher();student=DINOv3ConditionBridge(a.model).cuda().train()
+    for p in student.encoder.parameters():p.requires_grad=False
+    opt=torch.optim.AdamW([*student.proj.parameters(),*student.special.parameters()],lr=a.lr,weight_decay=.01);crit=FeatureDistillLoss(1,.5)
+    for step in range(1,a.steps+1):
+        p=files[(step-1)%len(files)];xt=image_tensor(p,518,'cuda');xs=image_tensor(p,512,'cuda')
+        with torch.inference_mode(),torch.autocast('cuda',dtype=torch.float16):yt=teacher_features(teacher,xt)
+        with torch.autocast('cuda',dtype=torch.float16):ys=student(xs);loss,parts=crit(ys,yt);loss.backward()
+        opt.step();opt.zero_grad(set_to_none=True)
+        if step==1 or step%100==0:print(json.dumps({'step':step,'loss':float(loss.detach()),'mse':float(parts['mse']),'cosine_loss':float(parts['cosine_loss'])}),flush=True)
+    out=Path(a.out);out.mkdir(parents=True,exist_ok=True);torch.save({'model':student.state_dict(),'model_id':a.model,'target_tokens':1374,'output_dim':1024},out/'dinov3_bridge.pt')
+if __name__=='__main__':main()

v7/training/distill_flow.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+import argparse,copy,json,os,random,sys,time
+from pathlib import Path
+import torch,yaml
+from torch.utils.data import Dataset,DataLoader
+from .lora import inject_lora,trainable_parameters,merge_lora
+from .distill_core import FlowBatch,teacher_rollout,student_loss,meanflow_pair_loss,StateAdapter
+
+class CacheDataset(Dataset):
+    def __init__(self,root): self.files=sorted(Path(root).glob('*.pt')); assert self.files,f'No .pt cache files in {root}'
+    def __len__(self): return len(self.files)
+    def __getitem__(self,i): return torch.load(self.files[i],map_location='cpu',weights_only=False)
+
+def move(v,dev):
+    if torch.is_tensor(v): return v.to(dev,non_blocking=True)
+    if isinstance(v,dict): return {k:move(x,dev) for k,x in v.items()}
+    if isinstance(v,(list,tuple)): return type(v)(move(x,dev) for x in v)
+    return v
+
+def make_batch(d,dev):
+    return FlowBatch(move(d['x'],dev),move(d['x_skl'],dev),move(d['cond'],dev),move(d.get('neg_cond'),dev),move(d.get('kwargs',{}),dev))
+
+def load_anigen_model(component,root='/home/user/app'):
+    sys.path.insert(0,root); os.chdir(root)
+    from anigen.utils.model_utils import load_model_from_path
+    path='ckpts/anigen/ss_flow_solo' if component=='ss_flow' else 'ckpts/anigen/slat_flow_auto'
+    model,cfg=load_model_from_path(path,model_name_in_config='denoiser',device='cuda',use_ema=False)
+    return model,cfg
+
+def save_state(model,out,meta):
+    out=Path(out);out.mkdir(parents=True,exist_ok=True)
+    torch.save(model.state_dict(),out/'student_merged.pt');(out/'meta.json').write_text(json.dumps(meta,indent=2))
+
+def train_stage(component,cfg,stage,cache,out,resume=None):
+    torch.backends.cuda.matmul.allow_tf32=True; torch.backends.cudnn.allow_tf32=True
+    teacher,_=load_anigen_model(component); teacher.eval();
+    for p in teacher.parameters(): p.requires_grad=False
+    student=copy.deepcopy(teacher).train(); replaced=inject_lora(student,cfg['distillation']['lora']['target_regex'],cfg['distillation']['lora']['rank'],cfg['distillation']['lora']['alpha'],cfg['distillation']['lora'].get('dropout',0))
+    if resume: student.load_state_dict(torch.load(resume,map_location='cpu',weights_only=False),strict=False)
+    params=trainable_parameters(student); opt=torch.optim.AdamW(params,lr=cfg['distillation']['optimizer']['lr'],weight_decay=cfg['distillation']['optimizer']['weight_decay'])
+    ds=CacheDataset(cache); dl=DataLoader(ds,batch_size=1,shuffle=True,num_workers=2,pin_memory=True,collate_fn=lambda x:x[0]);it=iter(dl)
+    accum=int(cfg['distillation'].get('grad_accum',8)); total=int(stage['train_steps']); t0=time.time(); opt.zero_grad(set_to_none=True)
+    for step in range(1,total+1):
+        try:d=next(it)
+        except StopIteration:it=iter(dl);d=next(it)
+        b=make_batch(d,'cuda'); # cached states may specify interval; otherwise sample macro interval.
+        t_hi=float(d.get('t_hi',random.uniform(0.25,1.0))); width=1.0/int(stage['student_steps']); t_lo=max(0.0,t_hi-width)
+        # Teacher uses the previous-stage budget over this macro interval.
+        local_teacher_steps=max(1,round(stage['teacher_steps']/stage['student_steps']))
+        with torch.inference_mode(),torch.autocast('cuda',dtype=torch.float16): te=teacher_rollout(teacher,b,t_hi,t_lo,local_teacher_steps)
+        with torch.autocast('cuda',dtype=torch.float16):
+            loss,parts=student_loss(student,b,t_hi,t_lo,te,cfg['distillation']['losses'])
+            loss=loss+meanflow_pair_loss(student,b,t_hi,t_lo,te,0.25 if 'meanflow' in cfg['distillation']['methods'] else 0.0)
+            (loss/accum).backward()
+        if step%accum==0:
+            torch.nn.utils.clip_grad_norm_(params,1.0);opt.step();opt.zero_grad(set_to_none=True)
+        if step==1 or step%100==0: print(json.dumps({'step':step,'loss':float(loss.detach()),**parts,'elapsed_s':round(time.time()-t0,1)}),flush=True)
+        if step%1000==0: torch.save(student.state_dict(),Path(out)/f'adapter_step_{step}.pt')
+    merged=merge_lora(student); meta={'component':component,'stage':stage,'lora_modules':replaced,'merged_modules':merged,'method':cfg['distillation']['methods'],'base':'VAST-AI/AniGen'};save_state(student,out,meta)
+
+def main():
+    ap=argparse.ArgumentParser();ap.add_argument('--config',default='v7/configs/l4_v7.yaml');ap.add_argument('--component',choices=['ss_flow','slat_flow'],required=True);ap.add_argument('--stage',type=int,default=0);ap.add_argument('--cache',required=True);ap.add_argument('--out',required=True);ap.add_argument('--resume');ap.add_argument('--steps',type=int);a=ap.parse_args()
+    cfg=yaml.safe_load(open(a.config)); stage=dict(cfg['distillation']['stages'][a.stage]);
+    if a.steps:stage['train_steps']=a.steps
+    train_stage(a.component,cfg,stage,a.cache,a.out,a.resume)
+if __name__=='__main__':main()

v7/training/fp8_sparse.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+import argparse,json,re,torch
+import torch.nn as nn
+
+@torch.no_grad()
+def prune_2to4_(model:nn.Module,include:str,exclude:str):
+    inc,rex=re.compile(include),re.compile(exclude); touched=[]
+    for name,m in model.named_modules():
+        if not isinstance(m,nn.Linear) or not inc.search(name) or rex.search(name): continue
+        w=m.weight.data; k=w.shape[-1]
+        if k%4: continue
+        flat=w.reshape(-1,k); g=flat.reshape(-1,k//4,4); idx=g.abs().argsort(dim=-1)[...,:2]
+        mask=torch.ones_like(g,dtype=torch.bool); mask.scatter_(-1,idx,False); g.mul_(mask); touched.append(name)
+    return touched
+
+def enforce_2to4_masks(model,masks):
+    with torch.no_grad():
+        for n,p in model.named_parameters():
+            if n in masks:p.mul_(masks[n])
+
+def capture_masks(model): return {n:(p!=0).to(p.dtype) for n,p in model.named_parameters() if p.ndim==2 and (p==0).any()}
+
+def modelopt_fp8(model,forward_loop,include,exclude):
+    import modelopt.torch.quantization as mtq
+    from modelopt.torch.quantization.config import FP8_DEFAULT_CFG
+    cfg=dict(FP8_DEFAULT_CFG); q=list(cfg.get('quant_cfg',[])) if isinstance(cfg.get('quant_cfg'),list) else cfg.get('quant_cfg',{})
+    # ModelOpt supports name filters globally; keep sensitive geometry/output paths unquantized.
+    if isinstance(q,dict):
+        q[f'*{exclude}*']={'enable':False}
+    def calibrate(m): forward_loop(m)
+    mtq.quantize(model=model,config=cfg,forward_func=calibrate)
+    # Explicitly disable excluded module quantizers when available.
+    for name,module in model.named_modules():
+        if re.search(exclude,name):
+            for attr in ('input_quantizer','weight_quantizer','output_quantizer'):
+                qz=getattr(module,attr,None)
+                if qz is not None and hasattr(qz,'disable'): qz.disable()
+    return model
+
+def selective_report(model,include,exclude):
+    rows=[]
+    for n,m in model.named_modules():
+        if isinstance(m,nn.Linear): rows.append({'name':n,'shape':list(m.weight.shape),'fp8_candidate':bool(re.search(include,n) and not re.search(exclude,n))})
+    return rows
+
+def main():
+    ap=argparse.ArgumentParser();ap.add_argument('--model',required=True);ap.add_argument('--out',required=True);ap.add_argument('--include',default='(qkv|to_q|to_k|to_v|proj|fc1|fc2|mlp|linear)');ap.add_argument('--exclude',default='(geo_head|skin_head|skl_head|out_layer|mesh|sparse|conv|norm)');ap.add_argument('--sparsity',action='store_true');a=ap.parse_args()
+    model=torch.load(a.model,map_location='cpu',weights_only=False)
+    touched=prune_2to4_(model,a.include,a.exclude) if a.sparsity else []
+    torch.save(model,a.out);open(a.out+'.report.json','w').write(json.dumps({'sparse_modules':touched,'layers':selective_report(model,a.include,a.exclude)},indent=2))
+if __name__=='__main__':main()

v7/training/geometry_moe.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+EXPERT_NAMES=('coarse_shape','surface_detail','limbs','head_face','skeleton','skin','materials')
+
+class GeometryRouter(nn.Module):
+    def __init__(self,dim,num_experts=len(EXPERT_NAMES),top_k=2):super().__init__();self.top_k=top_k;self.norm=nn.LayerNorm(dim);self.gate=nn.Linear(dim,num_experts,bias=False)
+    def forward(self,x,context=None):
+        h=self.norm(x if context is None else x+context);logits=self.gate(h.float());v,i=torch.topk(logits,self.top_k,dim=-1);return i,F.softmax(v,dim=-1).to(x.dtype)
+
+class ExpertMLP(nn.Module):
+    def __init__(self,dim,hidden,ratio=0.5):super().__init__();h=max(64,int(hidden*ratio));self.net=nn.Sequential(nn.Linear(dim,h),nn.SiLU(),nn.Linear(h,dim))
+    def forward(self,x):return self.net(x)
+
+class GeometryInductiveMoE(nn.Module):
+    """ModernMOE-inspired shared+routed experts for 3D flow transformer FFNs.
+    Expert semantics are encouraged by auxiliary labels/losses; routing itself remains learned.
+    """
+    def __init__(self,dim,hidden,num_experts=len(EXPERT_NAMES),top_k=2,expert_ratio=.5):
+        super().__init__();self.names=EXPERT_NAMES[:num_experts];self.router=GeometryRouter(dim,num_experts,top_k);self.shared=ExpertMLP(dim,hidden,expert_ratio);self.experts=nn.ModuleList([ExpertMLP(dim,hidden,expert_ratio) for _ in range(num_experts)]);self.out_norm=nn.LayerNorm(dim)
+    def forward(self,x,context=None):
+        idx,w=self.router(x,context);out=self.shared(x)
+        # Top-k dispatch; optimized Triton/TRT plugin can replace this without changing weights.
+        for slot in range(idx.shape[-1]):
+            ids=idx[...,slot];ws=w[...,slot]
+            for eid,expert in enumerate(self.experts):
+                mask=ids.eq(eid)
+                if not mask.any():continue
+                y=expert(x[mask]);out[mask]=out[mask]+y*ws[mask,None]
+        return self.out_norm(out),{'expert_idx':idx,'expert_weight':w}
+
+def load_balance_loss(router_info,num_experts):
+    idx=router_info['expert_idx'];w=router_info['expert_weight'];freq=torch.stack([(idx==i).float().mean() for i in range(num_experts)]);prob=torch.stack([torch.where(idx==i,w,torch.zeros_like(w)).sum(-1).mean() for i in range(num_experts)]);return num_experts*(freq*prob).sum()
+
+def expert_hint_loss(router_info,hints):
+    """hints: integer expert ids per token from geometry heuristics/teacher labels."""
+    idx=router_info['expert_idx'][...,0];return (idx!=hints).float().mean()

v7/training/lora.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+import math,re
+import torch
+import torch.nn as nn
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, rank: int=16, alpha: float=32.0, dropout: float=0.0):
+        super().__init__(); self.base=base; self.rank=rank; self.scale=alpha/rank
+        self.lora_a=nn.Linear(base.in_features,rank,bias=False,dtype=base.weight.dtype,device=base.weight.device)
+        self.lora_b=nn.Linear(rank,base.out_features,bias=False,dtype=base.weight.dtype,device=base.weight.device)
+        self.drop=nn.Dropout(dropout) if dropout else nn.Identity()
+        nn.init.kaiming_uniform_(self.lora_a.weight,a=math.sqrt(5)); nn.init.zeros_(self.lora_b.weight)
+        for p in self.base.parameters(): p.requires_grad=False
+    def forward(self,x): return self.base(x)+self.lora_b(self.lora_a(self.drop(x)))*self.scale
+    @torch.no_grad()
+    def merge(self):
+        delta=(self.lora_b.weight.float()@self.lora_a.weight.float())*self.scale
+        self.base.weight.add_(delta.to(self.base.weight.dtype)); return self.base
+
+def inject_lora(model: nn.Module, pattern: str, rank=16, alpha=32.0, dropout=0.0):
+    rx=re.compile(pattern); replaced=[]
+    for full_name,module in list(model.named_modules()):
+        if not isinstance(module,nn.Linear) or not rx.search(full_name): continue
+        if '.' in full_name: parent_name,leaf=full_name.rsplit('.',1); parent=model.get_submodule(parent_name)
+        else: parent,leaf=model,full_name
+        setattr(parent,leaf,LoRALinear(module,rank,alpha,dropout)); replaced.append(full_name)
+    if not replaced: raise RuntimeError(f'No Linear matched LoRA regex: {pattern}')
+    return replaced
+
+def trainable_parameters(model): return [p for p in model.parameters() if p.requires_grad]
+
+def merge_lora(model: nn.Module):
+    merged=[]
+    for full_name,module in list(model.named_modules()):
+        if not isinstance(module,LoRALinear): continue
+        base=module.merge()
+        if '.' in full_name: parent_name,leaf=full_name.rsplit('.',1); parent=model.get_submodule(parent_name)
+        else: parent,leaf=model,full_name
+        setattr(parent,leaf,base); merged.append(full_name)
+    return merged

v7/training/teacher_hybrid.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+import json,torch
+
+@dataclass
+class TeacherTarget:
+    vertices: torch.Tensor|None=None
+    faces: torch.Tensor|None=None
+    normals: torch.Tensor|None=None
+    joints: torch.Tensor|None=None
+    parents: torch.Tensor|None=None
+    skin: torch.Tensor|None=None
+    materials: dict|None=None
+    source: str=''
+
+class TeacherEnsemble:
+    """Geometry teacher = MeshFlow/TRELLIS.2; rig teacher = AniGen.
+    The adapters are intentionally lazy because MeshFlow has a non-commercial research license
+    and TRELLIS.2 has a distinct CUDA stack. Cache outputs once, then train students from tensors.
+    """
+    def __init__(self,meshflow=None,trellis=None,anigen=None):self.meshflow=meshflow;self.trellis=trellis;self.anigen=anigen
+    @torch.no_grad()
+    def generate(self,image,prefer='meshflow'):
+        geom=None
+        if prefer=='meshflow' and self.meshflow is not None:geom=self.meshflow(image)
+        elif self.trellis is not None:geom=self.trellis(image)
+        rig=self.anigen(image) if self.anigen is not None else None
+        return {'geometry':geom,'rig':rig}
+
+def chamfer_like(a,b,max_points=8192):
+    if a is None or b is None:return torch.tensor(0.,device=(a.device if a is not None else b.device))
+    if a.shape[0]>max_points:a=a[torch.randperm(a.shape[0],device=a.device)[:max_points]]
+    if b.shape[0]>max_points:b=b[torch.randperm(b.shape[0],device=b.device)[:max_points]]
+    d=torch.cdist(a.float()[None],b.float()[None]).squeeze(0);return d.min(1).values.mean()+d.min(0).values.mean()
+
+def geometry_teacher_loss(student_vertices,target:TeacherTarget):
+    return chamfer_like(student_vertices,target.vertices) if target.vertices is not None else student_vertices.sum()*0
+
+def rig_teacher_loss(student_joints,student_skin,target:TeacherTarget):
+    loss=student_joints.sum()*0
+    if target.joints is not None and student_joints.shape==target.joints.shape:loss=loss+torch.nn.functional.smooth_l1_loss(student_joints.float(),target.joints.float())
+    if target.skin is not None and student_skin.shape==target.skin.shape:loss=loss+torch.nn.functional.kl_div(student_skin.float().log_softmax(-1),target.skin.float().softmax(-1),reduction='batchmean')
+    return loss
+
+def save_teacher_cache(path,record):
+    path=Path(path);path.parent.mkdir(parents=True,exist_ok=True);torch.save(record,path)