| import fcntl, json, random, sys, time | |
| from pathlib import Path | |
| import torch | |
| import torch.utils.checkpoint | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| import sparse_support as s | |
| C, ROOT, RUN, OUT, DATA = s.C, s.ROOT, s.RUN, s.OUT, s.DATA | |
| def schedule(steps, rescale): | |
| return [rescale * (1 - i / steps) / (1 + (rescale - 1) * (1 - i / steps)) for i in range(steps + 1)] | |
| def times(): | |
| _, params = s.sparse_sampler() | |
| assert params['steps'] == 12, params | |
| return schedule(12, params['rescale_t']) | |
| def path(row, seed_index=0): | |
| return DATA / f"{row['id']}_seed{seed_index}.pt" | |
| def trajectory(model, noise, cond, coarse=False, which='base'): | |
| ts = times()[::2] if coarse else times() | |
| x = noise.cuda().float() | |
| states = [x.cpu()] | |
| with s.mode(model, which): | |
| for a, b in zip(ts, ts[1:]): | |
| x = x - (a - b) * s.velocity(model, x, a, cond) | |
| if not torch.isfinite(x).all(): | |
| raise RuntimeError('Non-finite rollout: stopped before caching/training.') | |
| states.append(x.cpu()) | |
| return torch.stack(states) | |
| def decoder_model(): | |
| d = s.load_model(s.DECODER_KEY) | |
| d.float() | |
| if hasattr(d, 'convert_to_fp32'): | |
| d.convert_to_fp32() | |
| return d.eval().requires_grad_(False).cuda() | |
| def mask_metrics(pred, ref): | |
| pred, ref = (pred.bool(), ref.bool()) | |
| inter = int((pred & ref).sum()) | |
| return dict(iou=inter / max(1, int((pred | ref).sum())), precision=inter / max(1, int(pred.sum())), recall=inter / max(1, int(ref.sum())), active=int(pred.sum())) | |
| """Native-conditioned small-data 12-to-6 Euler distillation. Embedded in notebook.""" | |
| import ast | |
| import hashlib | |
| import os | |
| import subprocess | |
| import numpy as np | |
| from types import SimpleNamespace | |
| METHOD = 'native_sparse_small_data_v1' | |
| def bridge(action, payload): | |
| job = ROOT/'bucket_job.json' | |
| s.dump(payload, job) | |
| subprocess.run([os.environ['T2_BUCKET_PY'], str(ROOT/'bucket_bridge.py'), action, str(job)], check=True) | |
| return json.loads(job.with_suffix('.result.json').read_text()) | |
| def upload(paths): | |
| pairs = [[str(p), C['output_prefix']+'/'+str(Path(p).relative_to(RUN))] | |
| for p in paths if Path(p).is_file()] | |
| if pairs: bridge('upload', dict(bucket=C['bucket'], files=pairs)) | |
| s.upload = upload | |
| def rows(): | |
| return json.loads((RUN/'selection.json').read_text())['rows'] | |
| def condition_path(r): | |
| return RUN/'conditions'/f"{r['id']}.pt" | |
| def signature(): | |
| return s.fingerprint(dict(config=json.loads((RUN/'identity.json').read_text()), | |
| selection=rows(), manifest=json.loads((ROOT/'manifest.json').read_text()))) | |
| def fetch(): | |
| # Restore only this new experiment. Never read old conditions or structure latents. | |
| bridge('restore', dict(bucket=C['bucket'], prefix=C['output_prefix'], root=str(RUN))) | |
| identity = {k:v for k,v in C.items() if k not in ('train_steps','session_minutes','save_every','eval_every')} | |
| identity['implementation_sha']=hashlib.sha256((ROOT/'distill.py').read_bytes()+(ROOT/'sparse_support.py').read_bytes()).hexdigest() | |
| ip = RUN/'identity.json' | |
| if ip.exists() and json.loads(ip.read_text()) != identity: | |
| raise RuntimeError('Run settings changed. Use a new RUN_NAME; no old cache was overwritten.') | |
| s.dump(identity, ip) | |
| from huggingface_hub import HfApi, hf_hub_download | |
| mp = RUN/'manifest.json' | |
| if mp.exists(): | |
| manifest = json.loads(mp.read_text()) | |
| else: | |
| # Resolve once; all future restores use immutable revisions. | |
| dino_sha = HfApi().model_info(C['dino_repo']).sha | |
| hf_hub_download(C['dino_repo'], 'config.json', revision=dino_sha) | |
| manifest = dict(source_commit=C['source_commit'], dino_revision=dino_sha, | |
| models={s.KEY:dict(repo='microsoft/TRELLIS.2-4B',revision='af44b45f2e35a493886929c6d786e563ec68364d',prefix='ckpts/ss_flow_img_dit_1_3B_64_bf16'), | |
| s.DECODER_KEY:dict(repo='microsoft/TRELLIS-image-large',revision='25e0d31ffbebe4b5a97464dd851910efc3002d96',prefix='ckpts/ss_dec_conv3d_16l8_fp16')}, | |
| pipeline={'args':{'sparse_structure_sampler':dict(name='FlowEulerGuidanceIntervalSampler',args={'sigma_min':1e-5}, | |
| params=dict(steps=12,rescale_t=5.,guidance_strength=7.5,guidance_rescale=.7,guidance_interval=[.6,1.]))}}) | |
| s.dump(manifest, mp) | |
| s.dump(manifest, ROOT/'manifest.json') | |
| if not (RUN/'selection.json').exists(): | |
| raw = RUN/'images'; raw.mkdir(exist_ok=True) | |
| candidates = bridge('list', dict(bucket=C['bucket'],prefix=C['source_prefix']+'/cache/object_images')) | |
| pngs = sorted(p for p in candidates if p.lower().endswith('.png')) | |
| random.Random(C['seed']).shuffle(pngs) | |
| chosen=[]; seen=set() | |
| if pngs: | |
| for remote in pngs: | |
| # Old collector names each PNG by object ID, not camera view. | |
| p=raw/Path(remote).name | |
| bridge('download',dict(bucket=C['bucket'],files=[[remote,str(p)]])) | |
| with Image.open(p) as im: | |
| im=im.convert('RGBA'); a=np.asarray(im) | |
| digest=hashlib.sha256(im.tobytes()+str(im.size).encode()).hexdigest() | |
| if digest in seen or not np.any(a[:,:,3]>204) or np.all(a[:,:,3]==255): continue | |
| seen.add(digest); chosen.append(dict(id=p.stem,image='images/'+p.name,source=remote,pixel_sha=digest)) | |
| if len(chosen)==50: break | |
| if len(chosen)<50: | |
| # No source latent repair. Stream RAW public images only; don't download the 48-GB repository. | |
| from datasets import load_dataset | |
| dataset_sha=HfApi().dataset_info('LEGO-Eval/object_images').sha | |
| stream=load_dataset('LEGO-Eval/object_images',revision=dataset_sha,split='train',streaming=True) | |
| objects=set() | |
| for item in stream: | |
| key=str(item.get('__key__','')); leaf=key.rsplit('/',1)[-1] | |
| if '_' not in leaf: raise RuntimeError('Dataset key schema changed; cannot establish object-disjoint split.') | |
| obj=leaf.rsplit('_',1)[0] | |
| if obj in objects: continue | |
| objects.add(obj) | |
| im=item.get('png') | |
| if im is None: continue | |
| im=im.convert('RGBA'); a=np.asarray(im) | |
| if not np.any(a[:,:,3]>204) or np.all(a[:,:,3]==255): continue | |
| digest=hashlib.sha256(im.tobytes()+str(im.size).encode()).hexdigest() | |
| oid=hashlib.sha256(obj.encode()).hexdigest() | |
| if digest in seen or any(r['id']==oid for r in chosen): continue | |
| p=raw/(oid+'.png'); im.save(p); seen.add(digest) | |
| chosen.append(dict(id=oid,image='images/'+p.name,source=key,pixel_sha=digest,dataset_revision=dataset_sha)) | |
| if len(chosen)==50: break | |
| if len(chosen)!=50: raise RuntimeError(f'Need 50 distinct alpha images; got {len(chosen)}.') | |
| random.Random(C['seed']).shuffle(chosen) | |
| for i,r in enumerate(chosen): r['split']='train' if i<32 else ('val' if i<40 else 'test') | |
| s.dump(dict(rows=chosen),RUN/'selection.json') | |
| # Keep runnable source with the experiment, but never upload tokens or the whole environment. | |
| source_paths=[] | |
| for name in ('distill.py','sparse_support.py','bucket_bridge.py','config.json'): | |
| dst=RUN/'code'/name; dst.parent.mkdir(exist_ok=True); dst.write_text((ROOT/name).read_text()); source_paths.append(dst) | |
| upload([ip,mp,RUN/'selection.json']+source_paths+[RUN/r['image'] for r in rows()]) | |
| print('50 raw images: 32 train / 8 validation / 10 sealed test. No old latent or DINO cache used.',flush=True) | |
| def native_preprocess(image): | |
| # Execute the exact alpha branch from pinned native source, without importing mesh dependencies. | |
| a=np.asarray(image.convert('RGBA')) | |
| if np.all(a[:,:,3]==255) or not np.any(a[:,:,3]>204): | |
| raise RuntimeError('Input needs a nonempty transparent background; do not silently bypass native preprocessing.') | |
| source=ROOT/'TRELLIS.2/trellis2/pipelines/trellis2_image_to_3d.py' | |
| tree=ast.parse(source.read_text()) | |
| cls=next(n for n in tree.body if isinstance(n,ast.ClassDef) and n.name=='Trellis2ImageTo3DPipeline') | |
| fn=next(n for n in cls.body if isinstance(n,ast.FunctionDef) and n.name=='preprocess_image') | |
| ns=dict(np=np,Image=Image) | |
| exec(compile(ast.Module(body=[fn],type_ignores=[]),str(source),'exec'),ns) | |
| return ns['preprocess_image'](SimpleNamespace(low_vram=True),image.convert('RGBA')) | |
| def conditions(test=False): | |
| from huggingface_hub import snapshot_download | |
| from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor | |
| manifest=json.loads((ROOT/'manifest.json').read_text()) | |
| selected=[r for r in rows() if (r['split']=='test')==test] | |
| pending=[r for r in selected if not condition_path(r).exists()] | |
| if not pending: return | |
| local=snapshot_download(C['dino_repo'],revision=manifest['dino_revision'],allow_patterns=['*.json','*.safetensors']) | |
| extractor=DinoV3FeatureExtractor(local,image_size=512); extractor.cuda() | |
| for r in pending: | |
| with Image.open(RUN/r['image']) as im: prepared=native_preprocess(im) | |
| cond=extractor([prepared]).float().cpu() | |
| if not s.valid_condition(cond): raise RuntimeError('Native DINO condition invalid') | |
| s.save(dict(signature=signature(),cond=cond),condition_path(r)) | |
| s.log('condition_ready',image=r['id'],shape=list(cond.shape)) | |
| del extractor; s.clear() | |
| upload([condition_path(r) for r in selected]) | |
| def read_condition(r): | |
| d=s.read(condition_path(r)) | |
| if d['signature']!=signature(): raise RuntimeError('Condition provenance mismatch') | |
| return d['cond'] | |
| def flow_model(adapters=True): | |
| m=s.load_model(s.KEY,torch.bfloat16).cuda() | |
| if adapters: m=s.install(m) | |
| s.checkpoint_policy(m,0) | |
| return m | |
| def collect_subset(selected, seeds): | |
| m=flow_model(False) # Teacher collection never attaches LoRA. | |
| pending=[] | |
| for r in selected: | |
| cond=read_condition(r) | |
| for k in range(seeds): | |
| p=path(r,k) | |
| if p.exists(): | |
| d=s.read(p) | |
| if d['signature']!=signature() or not torch.isfinite(d['states']).all(): | |
| raise RuntimeError('Trajectory provenance or finite-value check failed') | |
| continue | |
| seed=s.stable_seed(r['id']+f'/native/{k}') | |
| noise=torch.randn((1,m.in_channels,m.resolution,m.resolution,m.resolution),generator=torch.Generator().manual_seed(seed)).cuda() | |
| cg=cond.cuda(); states=trajectory(m,noise,cg) | |
| parity=OUT/'native_parity.json' | |
| if not parity.exists(): | |
| sampler,params=s.sparse_sampler() | |
| native=sampler.sample(m,noise,cond=cg,neg_cond=torch.zeros_like(cg),**params,verbose=False).samples.cpu() | |
| err=float((native-states[-1]).abs().max()) | |
| if not torch.allclose(native,states[-1],atol=2e-4,rtol=2e-4): raise RuntimeError(f'Native sampler parity failed: {err}') | |
| s.dump(dict(signature=signature(),max_abs_error=err),parity); upload([parity]) | |
| base=trajectory(m,noise,cg,True)[-1] | |
| s.save(dict(signature=signature(),method=METHOD,id=r['id'],split=r['split'],seed=seed,states=states,cond=cond,base6=base),p) | |
| pending.append(p); s.log('trajectory_ready',image=r['id'],seed_index=k) | |
| if len(pending)>=5 or s.STOP: upload(pending); pending=[] | |
| if s.STOP: raise RuntimeError('Paused safely. Rerun collection, not setup.') | |
| upload(pending); del m; s.clear() | |
| def preview(records,name): | |
| # Density projections carry information hidden by binary silhouettes. | |
| from PIL import ImageDraw | |
| canvas=Image.new('RGB',(1000,len(records)*150),'white'); draw=ImageDraw.Draw(canvas) | |
| for j,rec in enumerate(records): | |
| r=next(r for r in rows() if r['id']==rec['id']) | |
| with Image.open(RUN/r['image']) as im: | |
| tile=native_preprocess(im); tile.thumbnail((115,115)); canvas.paste(tile,(0,j*150+25)) | |
| draw.text((0,j*150+3),r['id'][:12],fill='black') | |
| for col,label in enumerate(('teacher12','base6','student6')): | |
| if label not in rec: continue | |
| cube=rec[label].reshape(32,32,32).float() | |
| x=125+col*290 | |
| draw.text((x,j*150+3),f'{label}: {int(cube.sum())} voxels',fill='black') | |
| for axis in range(3): | |
| a=(255-255*cube.mean(axis).sqrt()).byte().numpy() | |
| tile=Image.fromarray(a).convert('RGB').resize((90,110)) | |
| canvas.paste(tile,(x+axis*95,j*150+25)) | |
| canvas.save(OUT/name) | |
| def audit_subset(selected,seeds,label): | |
| decoder=decoder_model(); records=[]; pics=[]; paths=[] | |
| for r in selected: | |
| for k in range(seeds): | |
| p=path(r,k); d=s.read(p) | |
| mask=s.decode_mask(decoder,d['states'][-1]).cpu() | |
| base=s.decode_mask(decoder,d['base6']).cpu() | |
| count=int(mask.sum()) | |
| d.update(teacher_mask=mask,base_mask=base); s.save(d,p); paths.append(p) | |
| records.append(dict(id=r['id'],seed_index=k,active=count,passed=64<=count<.90*32768)) | |
| if k==0: pics.append(dict(id=r['id'],teacher12=mask,base6=base)) | |
| del decoder; s.clear() | |
| report=dict(signature=signature(),passed=all(r['passed'] for r in records),records=records, | |
| warning='Heuristic numerical/degeneracy gate, not proof of semantic quality. No failed record is replaced.') | |
| s.dump(report,OUT/f'{label}.json'); preview(pics,f'{label}.png') | |
| points=[dict(id=p['id'],teacher12=p['teacher12'].reshape(32,32,32).nonzero().tolist(), | |
| base6=p['base6'].reshape(32,32,32).nonzero().tolist()) for p in pics] | |
| s.dump(points,OUT/f'{label}_3d.json') | |
| upload(paths+[OUT/f'{label}.json',OUT/f'{label}.png',OUT/f'{label}_3d.json']) | |
| if not report['passed']: raise RuntimeError(f'{label} FAILED. Inspect teacher; no training, no silent replacements.') | |
| def smoke(): | |
| conditions() | |
| selected=[r for r in rows() if r['split']=='train'][:6] | |
| collect_subset(selected,1); audit_subset(selected,1,'smoke_audit') | |
| def approval(label): | |
| p=OUT/f'{label}_approved.json' | |
| if not p.exists() or json.loads(p.read_text()).get('signature')!=signature(): | |
| raise RuntimeError(f'Inspect preview and run the {label} approval cell first.') | |
| report=json.loads((OUT/f'{label}.json').read_text()) | |
| if not report['passed'] or report['signature']!=signature(): | |
| raise RuntimeError('Audit no longer passes for this run.') | |
| def collect(): | |
| approval('smoke_audit') | |
| selected=[r for r in rows() if r['split']!='test'] | |
| collect_subset(selected,2); audit_subset(selected,2,'teacher_audit') | |
| def sample_loss(m,d,interval,on_policy=False,anchor=False): | |
| ts=times(); i=2*interval; cond=d['cond'].cuda().float() | |
| with torch.no_grad(): | |
| x=d['states'][i].cuda().float() | |
| target=d['states'][i+2].cuda().float() | |
| if on_policy or anchor: | |
| # Fresh noise for local on-policy targets; cached noise for exact endpoint anchors. | |
| x=d['states'][0].cuda().float() if anchor else torch.randn_like(x) | |
| with s.mode(m,'student'): | |
| for j in range(interval): x=x-(ts[2*j]-ts[2*j+2])*s.velocity(m,x,ts[2*j],cond) | |
| if not anchor: | |
| with s.mode(m,'base'): | |
| mid=x-(ts[i]-ts[i+1])*s.velocity(m,x,ts[i],cond) | |
| target=mid-(ts[i+1]-ts[i+2])*s.velocity(m,mid,ts[i+1],cond) | |
| target_v=(x-target)/(ts[i]-ts[i+2]) | |
| pred=s.velocity(m,x,ts[i],cond) | |
| loss=F.mse_loss(pred,target_v) | |
| if not torch.isfinite(loss): raise RuntimeError('Nonfinite objective. No update applied.') | |
| return loss | |
| def validation(m,step,split='val'): | |
| rng=s.seed_state(); decoded=[] | |
| for r in rows(): | |
| if r['split']!=split: continue | |
| for k in range(2): | |
| d=s.read(path(r,k)) | |
| student=trajectory(m,d['states'][0],d['cond'].cuda(),True,'student')[-1] | |
| decoded.append((r,k,d,student)) | |
| m.cpu(); s.clear(); decoder=decoder_model(); records=[]; pictures=[] | |
| for r,k,d,student in decoded: | |
| mask=s.decode_mask(decoder,student).cpu() | |
| records.append(dict(id=r['id'],seed_index=k,base6=mask_metrics(d['base_mask'],d['teacher_mask']), | |
| student6=mask_metrics(mask,d['teacher_mask']),base6_mse=float(F.mse_loss(d['base6'],d['states'][-1])), | |
| student6_mse=float(F.mse_loss(student,d['states'][-1])))) | |
| if k==0: pictures.append(dict(id=r['id'],teacher12=d['teacher_mask'],base6=d['base_mask'],student6=mask)) | |
| report=dict(step=step,split=split,records=records) | |
| for model in ('base6','student6'): | |
| for metric in ('iou','precision','recall'): | |
| report[f'{model}_{metric}']=float(np.mean([r[model][metric] for r in records])) | |
| report[f'{model}_mse']=float(np.mean([r[f'{model}_mse'] for r in records])) | |
| # Image-level bootstrap; two seeds are not two independent objects. | |
| deltas=[np.mean([r['student6']['iou']-r['base6']['iou'] for r in records if r['id']==oid]) for oid in sorted({r['id'] for r in records})] | |
| boots=np.random.default_rng(2026).choice(deltas,size=(2000,len(deltas)),replace=True).mean(1) | |
| report['iou_delta_95pct_bootstrap']=np.quantile(boots,[.025,.975]).tolist() | |
| report['image_win_fraction']=float(np.mean(np.asarray(deltas)>0)) | |
| report['qualifies']=(step>0 and report['student6_iou']>report['base6_iou']+.01 and | |
| report['student6_precision']>=report['base6_precision']-.02 and report['student6_recall']>=report['base6_recall']-.02) | |
| # Selection is geometry-led. MSE is reported, never sufficient by itself. | |
| s.dump(report,OUT/f'{split}_{step:05d}.json'); preview(pictures,f'{split}_comparison.png') | |
| upload([OUT/f'{split}_{step:05d}.json',OUT/f'{split}_comparison.png']) | |
| s.log('validation',**{k:v for k,v in report.items() if k!='records'}) | |
| del decoder,decoded; s.clear(); m.cuda(); s.restore_rng(rng) | |
| return report | |
| def preflight(): | |
| approval('teacher_audit') | |
| m=flow_model(); s.checkpoint_policy(m,2) | |
| params=[p for p in m.parameters() if p.requires_grad] | |
| opt=torch.optim.AdamW(params,lr=C['learning_rate']) | |
| d=s.read(path(next(r for r in rows() if r['split']=='train'))) | |
| before=s.adapter_state(m); torch.cuda.reset_peak_memory_stats() | |
| for interval in (0,5): | |
| opt.zero_grad(set_to_none=True) | |
| with s.mode(m,'student'): | |
| loss=sample_loss(m,d,interval,on_policy=(interval==5)); loss.backward() | |
| torch.nn.utils.clip_grad_norm_(params,1.,error_if_nonfinite=True); opt.step() | |
| after=s.adapter_state(m) | |
| delta=sum(float((after[k]-before[k]).square().sum()) for k in after) | |
| if not delta>0: raise RuntimeError('Adapter did not update!') | |
| report=dict(signature=signature(),passed=True,adapter_delta_l2=delta**.5, | |
| peak_allocated_gib=torch.cuda.max_memory_allocated()/2**30, | |
| peak_reserved_gib=torch.cuda.max_memory_reserved()/2**30) | |
| s.dump(report,OUT/'preflight.json'); upload([OUT/'preflight.json']) | |
| print('Forward/backward/update passed; temporary updates discarded.',report,flush=True) | |
| def train(): | |
| approval('teacher_audit') | |
| if json.loads((OUT/'preflight.json').read_text())['signature']!=signature(): raise RuntimeError('Run preflight again') | |
| random.seed(C['seed']); torch.manual_seed(C['seed']); torch.cuda.manual_seed_all(C['seed']) | |
| if (OUT/'test_seal.json').exists(): | |
| raise RuntimeError('Final test already opened. Do not train further against the same test set.') | |
| m=flow_model(); s.checkpoint_policy(m,2) | |
| params=[p for p in m.parameters() if p.requires_grad] | |
| opt=torch.optim.AdamW(params,lr=C['learning_rate'],weight_decay=0.) | |
| step=0; best=-1.; latest=OUT/'latest.pt' | |
| if latest.exists(): | |
| ck=s.read(latest) | |
| if ck['signature']!=signature(): raise RuntimeError('Checkpoint provenance mismatch') | |
| s.apply_adapter(m,ck['adapter']); opt.load_state_dict(ck['optimizer']) | |
| step,best=ck['step'],ck['best_iou']; s.restore_rng(ck['rng']) | |
| else: | |
| baseline=validation(m,0) | |
| if abs(baseline['student6_mse']-baseline['base6_mse'])>1e-5: | |
| raise RuntimeError('Zero-adapter student != base6. Stop before training.') | |
| items=[path(r,k) for r in rows() if r['split']=='train' for k in range(2)] | |
| def persist(): | |
| s.save(dict(method=METHOD,signature=signature(),step=step,best_iou=best,adapter=s.adapter_state(m), | |
| optimizer=opt.state_dict(),rng=s.seed_state()),latest) | |
| upload([latest,OUT/'events.jsonl']) | |
| started=time.monotonic() | |
| while step<C['train_steps']: | |
| if s.STOP or (C['session_minutes'] and time.monotonic()-started>C['session_minutes']*60): break | |
| lr=C['learning_rate']*min(1.,(step+1)/50) | |
| for g in opt.param_groups: g['lr']=lr | |
| opt.zero_grad(set_to_none=True); value=0. | |
| for micro in range(C['micro_batches']): | |
| d=s.read(random.choice(items)); kind=random.random() if step>=100 else 0. | |
| anchor=kind>=.75; online=.5<=kind<.75 | |
| interval=5 if anchor else random.randrange(6) | |
| with s.mode(m,'student'): | |
| loss=sample_loss(m,d,interval,online,anchor)/C['micro_batches']; loss.backward() | |
| value+=float(loss.detach()) | |
| norm=torch.nn.utils.clip_grad_norm_(params,1.,error_if_nonfinite=True) | |
| opt.step(); step+=1; opt.zero_grad(set_to_none=True) | |
| s.log('train',step=step,total=C['train_steps'],loss=value,grad_norm=float(norm),lr=lr, | |
| elapsed_seconds=time.monotonic()-started,peak_vram_gib=torch.cuda.max_memory_allocated()/2**30) | |
| if step%C['save_every']==0: persist() | |
| if step%C['eval_every']==0 or step==C['train_steps']: | |
| s.checkpoint_policy(m,0); report=validation(m,step); s.checkpoint_policy(m,2) | |
| if report['qualifies'] and report['student6_iou']>best: | |
| best=report['student6_iou'] | |
| s.save(dict(signature=signature(),method=METHOD,step=step,adapter=s.adapter_state(m),report=report),OUT/'best.pt') | |
| upload([OUT/'best.pt']) | |
| persist() | |
| if step>=400 and best<0: | |
| s.log('no_improvement_stop',step=step,message='No geometry improvement. Inspect the run, do not blindly extend.') | |
| break | |
| complete=step>=C['train_steps'] or (step>=400 and best<0) | |
| s.dump(dict(step=step,ready_for_test=complete),OUT/'training_status.json') | |
| persist(); upload([OUT/'training_status.json']) | |
| s.log('session_saved',step=step,finished=step>=C['train_steps'],ready_for_test=complete,best_iou=best) | |
| def evaluate(): | |
| # Test is not touched during checkpoint selection. Freeze chosen checkpoint before test generation. | |
| best=OUT/'best.pt'; latest=OUT/'latest.pt' | |
| if not latest.exists(): raise RuntimeError('No trained checkpoint') | |
| status=OUT/'training_status.json' | |
| if not status.exists() or not json.loads(status.read_text())['ready_for_test']: | |
| raise RuntimeError('Training session paused, not completed. Rerun the training cell before opening the final test.') | |
| ck=s.read(best if best.exists() else latest) | |
| if ck['signature']!=signature() or ck['step']==0: raise RuntimeError('Invalid or step-zero checkpoint') | |
| adapter_sha=s.fingerprint({k:hashlib.sha256(v.numpy().tobytes()).hexdigest() for k,v in ck['adapter'].items()}) | |
| seal=OUT/'test_seal.json' | |
| if seal.exists() and json.loads(seal.read_text())['adapter_sha']!=adapter_sha: | |
| raise RuntimeError('Test already used by a different checkpoint. Choose new independent test images in a new run.') | |
| s.dump(dict(adapter_sha=adapter_sha,step=ck['step']),seal); upload([seal]) | |
| conditions(test=True); selected=[r for r in rows() if r['split']=='test'] | |
| collect_subset(selected,2); audit_subset(selected,2,'test_teacher_audit') | |
| m=flow_model(); s.apply_adapter(m,ck['adapter']); report=validation(m,ck['step'],'test') | |
| report.update(checkpoint='best.pt' if best.exists() else 'latest.pt',adapter_sha=adapter_sha, | |
| status='PILOT_IMPROVED' if best.exists() and report['qualifies'] else 'NOT_ESTABLISHED', | |
| limitation='Ten held-out objects, not broad deployment evidence. Geometry and combined shape need visual inspection.') | |
| s.dump(report,OUT/'evaluation.json') | |
| # Save trained candidate even on failure, but never label it successful or replace prior shape files. | |
| s.save_file({k:v.contiguous() for k,v in ck['adapter'].items()},str(OUT/'sparse6_candidate.safetensors')) | |
| sampler=json.loads((ROOT/'manifest.json').read_text())['pipeline']['args']['sparse_structure_sampler'] | |
| ts=times()[::2] | |
| nfe=sum(2 if .6<=t<=1 else 1 for t in ts[:-1]) | |
| s.dump(dict(method=METHOD,signature=signature(),rank=C['rank'],lora_scale=1.,steps=6,timesteps=ts, | |
| flow_forward_calls=nfe,base_model=json.loads((ROOT/'manifest.json').read_text())['models'][s.KEY], | |
| source_commit=C['source_commit'],sampler=sampler,status=report['status'], | |
| note='Set sampler steps to 6. Apply exact .a/.b attention LoRA to unquantized base. Not a Q4-ready replacement.'),OUT/'adapter_config.json') | |
| helper=ROOT/'apply_sparse6.py' | |
| (OUT/'apply_sparse6.py').write_text(helper.read_text()) | |
| upload([OUT/'evaluation.json',OUT/'sparse6_candidate.safetensors',OUT/'adapter_config.json',OUT/'apply_sparse6.py']) | |
| print(json.dumps({k:v for k,v in report.items() if k!='records'},indent=2)) | |
| if __name__=='__main__': | |
| lock=(ROOT/'native_worker.lock').open('w') | |
| try: fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) | |
| except BlockingIOError: raise SystemExit('Another worker is active. Use status; do not start a duplicate.') | |
| action=sys.argv[1] | |
| if action!='fetch': | |
| s.dump(json.loads((RUN/'manifest.json').read_text()),ROOT/'manifest.json') | |
| s.log('gpu',name=torch.cuda.get_device_name()) | |
| dict(fetch=fetch,smoke=smoke,collect=collect,preflight=preflight,train=train,evaluate=evaluate)[action]() | |
Xet Storage Details
- Size:
- 26.2 kB
- Xet hash:
- cf43daaae0ecd9fe01fbebb1ff13bdd68f69304fccbd3f10da888be82f314d1c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.