File size: 8,409 Bytes
4719196 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | #!/usr/bin/env python3
"""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()
|