| |
| """Inference with the trained BaseNet -> Brainmu generation-expert LoRA.""" |
| from __future__ import annotations |
| import argparse, json, math, os, random, sys |
| from pathlib import Path |
| import cv2 |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from accelerate import init_empty_weights, load_checkpoint_and_dispatch |
| from PIL import Image |
| from safetensors import safe_open |
| from project_config import load_config |
|
|
| class LoRALinear(nn.Module): |
| def __init__(self, base: nn.Linear, rank: int, alpha: float, dropout: float): |
| super().__init__(); self.base=base; self.base.requires_grad_(False); self.rank=rank; self.alpha=alpha; self.scaling=alpha/rank |
| self.dropout=nn.Dropout(dropout) if dropout>0 else nn.Identity() |
| self.lora_A=nn.Parameter(torch.empty(rank,base.in_features,device=base.weight.device,dtype=torch.float32),requires_grad=False) |
| self.lora_B=nn.Parameter(torch.empty(base.out_features,rank,device=base.weight.device,dtype=torch.float32),requires_grad=False) |
| def forward(self,x): |
| return self.base(x)+F.linear(F.linear(self.dropout(x),self.lora_A),self.lora_B)*self.scaling |
|
|
| def resolve_parent(root,dotted): |
| parts=dotted.split('.'); parent=root |
| for part in parts[:-1]: parent=parent[int(part)] if part.isdigit() else getattr(parent,part) |
| return parent,parts[-1] |
|
|
| def inject_lora(model,spec,adapter_path): |
| modules={} |
| for name in spec['target_modules']: |
| parent,leaf=resolve_parent(model,name); base=getattr(parent,leaf) |
| if not isinstance(base,nn.Linear): raise TypeError(f'{name} is {type(base)}, expected nn.Linear') |
| wrapped=LoRALinear(base,int(spec['rank']),float(spec['alpha']),float(spec['dropout'])); setattr(parent,leaf,wrapped); modules[name]=wrapped |
| with safe_open(adapter_path,framework='pt',device='cpu') as f: |
| keys=set(f.keys()); expected={f'{n}.lora_A' for n in modules}|{f'{n}.lora_B' for n in modules} |
| if keys!=expected: raise RuntimeError(f'adapter key mismatch missing={len(expected-keys)} extra={len(keys-expected)}') |
| with torch.no_grad(): |
| for name,m in modules.items(): |
| m.lora_A.copy_(f.get_tensor(f'{name}.lora_A').to(m.lora_A.device)); m.lora_B.copy_(f.get_tensor(f'{name}.lora_B').to(m.lora_B.device)) |
| return modules |
|
|
| def set_seed(seed): |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed) |
|
|
| def ssim_gray(a,b): |
| c1,c2=.01**2,.03**2; ma=cv2.GaussianBlur(a,(11,11),1.5); mb=cv2.GaussianBlur(b,(11,11),1.5) |
| sa=cv2.GaussianBlur(a*a,(11,11),1.5)-ma*ma; sb=cv2.GaussianBlur(b*b,(11,11),1.5)-mb*mb; sab=cv2.GaussianBlur(a*b,(11,11),1.5)-ma*mb |
| return float((((2*ma*mb+c1)*(2*sab+c2))/((ma*ma+mb*mb+c1)*(sa+sb+c2))).mean()) |
|
|
| def metrics(pred,gt): |
| p=np.asarray(pred.convert('L'),np.float32)/255.; t=np.asarray(gt.convert('L').resize(pred.size),np.float32)/255.; mse=float(np.mean((p-t)**2)) |
| return {'psnr_db':float(-10*np.log10(max(mse,1e-12))),'ssim':ssim_gray(p,t)} |
|
|
| def files(path): |
| if path.is_file(): return [path] |
| return sorted(p for p in path.iterdir() if p.suffix.lower() in {'.png','.jpg','.jpeg','.webp'}) |
|
|
| def main(): |
| base=Path(__file__).resolve().parents[1] |
| settings=load_config()['generation']['inference'] |
| ap=argparse.ArgumentParser() |
| ap.add_argument('--input',type=Path,required=True,help='BaseNet PNG or directory of BaseNet PNGs') |
| ap.add_argument('--output-dir',type=Path,required=True) |
| ap.add_argument('--gt',type=Path,help='optional GT PNG or directory with matching filenames') |
| ap.add_argument('--brainmu-repo',type=Path,default=base/'vendor/Brainmu') |
| ap.add_argument('--model-path',type=Path,default=base/'checkpoints/Brainmu') |
| ap.add_argument('--adapter',type=Path,default=base/'checkpoints/lora/adapter.safetensors') |
| ap.add_argument('--adapter-config',type=Path,default=base/'checkpoints/lora/adapter_config.json') |
| ap.add_argument('--prompt',default=settings['prompt']) |
| ap.add_argument('--steps',type=int,default=settings['steps']); ap.add_argument('--seed',type=int,default=settings['seed']); ap.add_argument('--seed-offset',type=int,default=0); ap.add_argument('--seed-mode',choices=['sequential','base-prefix'],default='sequential',help='base-prefix uses the leading numeric clip id so all temporal spans of one clip share a seed') |
| a=ap.parse_args() |
| a.input=a.input.resolve(); a.output_dir=a.output_dir.resolve(); a.brainmu_repo=a.brainmu_repo.resolve(); a.model_path=a.model_path.resolve(); a.adapter=a.adapter.resolve(); a.adapter_config=a.adapter_config.resolve() |
| if a.gt: a.gt=a.gt.resolve() |
| a.output_dir.mkdir(parents=True,exist_ok=True) |
| inputs=files(a.input) |
| if not inputs: raise RuntimeError(f'no images found: {a.input}') |
| sys.path.insert(0,str(a.brainmu_repo.resolve())); os.chdir(a.brainmu_repo) |
| from data.data_utils import add_special_tokens |
| from data.transforms import ImageTransform |
| from inferencer import InterleaveInferencer |
| from modeling.autoencoder import load_ae |
| from modeling.brainmu import Brainmu,BrainmuConfig,Qwen2Config,Qwen2ForCausalLM,SiglipVisionConfig,SiglipVisionModel |
| from modeling.qwen2 import Qwen2Tokenizer |
| model_path=a.model_path.resolve(); device=torch.device('cuda:0'); torch.cuda.set_device(device); torch.backends.cuda.matmul.allow_tf32=True |
| llm=Qwen2Config.from_json_file(str(model_path/'llm_config.json')); llm.qk_norm=True; llm.tie_word_embeddings=False; llm.layer_module='Qwen2MoTDecoderLayer' |
| vit=SiglipVisionConfig.from_json_file(str(model_path/'vit_config.json')); vit.rope=False; vit.num_hidden_layers-=1 |
| vae,vae_cfg=load_ae(local_path=str(model_path/'ae.safetensors')) |
| cfg=BrainmuConfig(visual_gen=True,visual_und=True,llm_config=llm,vit_config=vit,vae_config=vae_cfg,vit_max_num_patch_per_side=70,connector_act='gelu_pytorch_tanh',latent_patch_size=2,max_latent_size=64,timestep_shift=1.0) |
| with init_empty_weights(): |
| lm=Qwen2ForCausalLM(llm); vm=SiglipVisionModel(vit); model=Brainmu(lm,vm,cfg); model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit,meta=True) |
| print('LOAD_BASE_BEGIN',flush=True) |
| model=load_checkpoint_and_dispatch(model,checkpoint=str(model_path/'ema.safetensors'),device_map={'':0},dtype=torch.bfloat16,force_hooks=True) |
| model.requires_grad_(False).eval(); vae=vae.to(device=device,dtype=torch.float32).eval().requires_grad_(False) |
| oe,od=vae.encode,vae.decode; vae.encode=lambda x:oe(x.to(device=device,dtype=torch.float32)); vae.decode=lambda z:od(z.to(device=device,dtype=torch.float32)) |
| spec=json.load(open(a.adapter_config)); adapters=inject_lora(model,spec,a.adapter); model.eval(); print(f'LORA_LOADED modules={len(adapters)} adapter={a.adapter}',flush=True) |
| tok=Qwen2Tokenizer.from_pretrained(str(model_path)); tok,new_ids,_=add_special_tokens(tok) |
| infer=InterleaveInferencer(model,vae,tok,ImageTransform(400,256,16),ImageTransform(392,252,14),new_ids) |
| rows=[] |
| for i,p in enumerate(inputs): |
| seed_index = i if a.seed_mode == 'sequential' else int(p.stem.split('_',1)[0]) |
| seed=a.seed+a.seed_offset+seed_index; set_seed(seed); inp=Image.open(p).convert('RGB') |
| with torch.no_grad(): |
| pred=infer.interleave_inference([inp,a.prompt],think=False,understanding_output=False,cfg_text_scale=settings['cfg_text_scale'],cfg_img_scale=settings['cfg_img_scale'],cfg_interval=settings['cfg_interval'],timestep_shift=settings['timestep_shift'],num_timesteps=a.steps,cfg_renorm_min=settings['cfg_renorm_min'],cfg_renorm_type=settings['cfg_renorm_type'])[-1] |
| out=a.output_dir/(p.stem+'.png'); pred.save(out); row={'input':str(p),'output':str(out),'seed':seed,'prompt':a.prompt} |
| if a.gt: |
| gp=a.gt if a.gt.is_file() else a.gt/(p.stem+'.png'); row.update(metrics(pred,Image.open(gp).convert('RGB'))); row['gt']=str(gp) |
| rows.append(row); print(json.dumps(row),flush=True) |
| result={'count':len(rows),'items':rows} |
| if a.gt: |
| result['mean_psnr_db']=float(np.mean([x['psnr_db'] for x in rows])); result['mean_ssim']=float(np.mean([x['ssim'] for x in rows])) |
| (a.output_dir/'metrics.json').write_text(json.dumps(result,indent=2)+'\n'); print('INFERENCE_COMPLETE '+json.dumps({k:v for k,v in result.items() if k!='items'}),flush=True) |
| if __name__=='__main__': main() |
|
|
