| from pathlib import Path |
| import json,sys,time,torch,shutil |
| root=Path(__file__).resolve().parent |
| sys.path.insert(0,str(root.resolve())) |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from native.configuration_longcat_next import LongcatNextConfig |
| from native.image_refiner import ImageRefinerContainer |
| from native.refiner_modules import RotaryPosEmbed |
| torch.set_num_threads(4);checks=[];t=time.time() |
| def check(name,ok,**details): |
| assert ok,(name,details) |
| checks.append(dict(name=name,passed=True,**details));print(name,details,flush=True) |
| cfg=LongcatNextConfig(**json.loads((root/'source_config.json').read_text())).visual_config.visual_decoder_config |
| with init_empty_weights():m=ImageRefinerContainer(cfg) |
| state={k.removeprefix('image_refiner.'):v for k,v in load_file(root/'decoder/model.safetensors').items() if k.startswith('image_refiner.')} |
| state.update({k.removeprefix('image_refiner.'):v for k,v in load_file(root/'refiner/model.safetensors').items()}) |
| m.load_state_dict(state,strict=True,assign=True);m.eval();del state |
| check('strict_full_flow_refiner',True,tensors=len(m.state_dict()),transformer_tensors=len(m.base_transformer.state_dict())) |
| dtype=next(m.base_transformer.parameters()).dtype;check('native_bfloat16_weights',dtype==torch.bfloat16) |
| torch.manual_seed(1019);latent=torch.randn(1,16,4,4,dtype=dtype);reference=[[torch.randn(16,4,4,dtype=dtype)]];features=torch.randn(1,4,3584,dtype=dtype);mask=torch.ones(1,4,dtype=torch.bool);tc=cfg.transformer_config |
| freqs=RotaryPosEmbed.get_freqs_cis(tc.axes_dim_rope,tc.axes_lens,theta=10000) |
| with torch.inference_mode(): |
| cond=m.cond_proj(features);check('native_refiner_conditioning',cond.shape==(1,4,2048) and torch.isfinite(cond).all().item()) |
| result=m.base_transformer(latent,torch.tensor([.5],dtype=dtype),cond,freqs,mask,ref_image_hidden_states=reference,return_dict=False) |
| pred=result if torch.is_tensor(result) else (result[0] if isinstance(result,(tuple,list)) else result.sample) |
| if isinstance(pred,list):pred=pred[0].unsqueeze(0) |
| check('conditioned_flow_forward',pred.shape==latent.shape and torch.isfinite(pred).all().item(),shape=list(pred.shape),dtype=str(pred.dtype)) |
| update=latent.float()-.1*pred.float();check('single_euler_update',torch.isfinite(update).all().item() and not torch.equal(update,latent.float()),mean_abs_update=(update-latent.float()).abs().mean().item()) |
| altered=m.base_transformer(latent,torch.tensor([.5],dtype=dtype),cond*0,freqs,mask,ref_image_hidden_states=reference,return_dict=False) |
| other=altered if torch.is_tensor(altered) else (altered[0] if isinstance(altered,(tuple,list)) else altered.sample) |
| if isinstance(other,list):other=other[0].unsqueeze(0) |
| check('condition_affects_flow',not torch.equal(pred,other),max_abs_difference=(pred-other).abs().max().item()) |
| ex=json.loads((root/'refiner/extraction.json').read_text());check('refiner_source_tensor_bytes_exact',ex['source_tensor_bytes_equal']) |
| report=dict(strict_load=True,source_tensor_bytes_equal=True,checks=checks,smoke=dict(passed=True,device='cpu',dtype=str(dtype),latent_shape=[1,16,4,4],reference_images=1),elapsed_seconds=time.time()-t,limitations=['Tiny-latent native flow forward and one Euler update, not full-resolution image refinement or perceptual-quality evaluation.','No LongCat language/audio backbone.']) |
| (root/'refiner_validation.json').write_text(json.dumps(report,indent=2)+'\n') |
| base=json.loads((root/'validation.json').read_text());base['checks'].extend([{'name':'refiner_'+c['name'],**{k:v for k,v in c.items() if k!='name'}} for c in checks]);base['full_flow_refiner_validation']='refiner_validation.json';base['limitations']=[x for x in base['limitations'] if 'separate validation report' not in x];base['limitations'].append('Full flow-refiner weights and tiny-latent conditioned forward are validated separately; no full-resolution quality benchmark.');(root/'validation.json').write_text(json.dumps(base,indent=2)+'\n') |
| print('LONGCAT FULL REFINER VALIDATED',report['elapsed_seconds'],flush=True) |
|
|
