| """LongCat-Next dNaViT/RVQ codec and base pixel decoder. Full AR backbone excluded.""" |
| from pathlib import Path |
| import json,torch |
| from torch import nn |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from native.configuration_longcat_next import LongcatNextConfig |
| from native.visual import LongcatNextVisualTokenizer,VisionTransformerDecoder |
|
|
| def build(config,decoder=False,decoder_state=None): |
| cfg=LongcatNextConfig(**config);m=nn.Module();m.model=nn.Module();m.model.visual_tokenizer=LongcatNextVisualTokenizer(cfg) |
| if decoder: |
| from diffusers import AutoencoderKL |
| m.image_decoder=VisionTransformerDecoder(cfg.visual_config.visual_decoder_config.image_decoder_config) |
| m.image_refiner=nn.Module();m.image_refiner.vae=AutoencoderKL.from_config(config["visual_config"]["visual_decoder_config"]["vae_config"]) |
| w=decoder_state["image_refiner.cond_proj.weight"];m.image_refiner.cond_proj=nn.Linear(w.shape[1],w.shape[0],bias=True) |
| return m |
|
|
| def load_model(path=None,device="cpu",decoder=True): |
| root=Path(path or Path(__file__).parent);c=json.loads((root/"source_config.json").read_text());s=load_file(root/"model.safetensors") |
| aliases=json.loads((root/"aliases.json").read_text()) |
| for alias,canonical in aliases.items():s[alias]=s[canonical] |
| ds=load_file(root/"decoder/model.safetensors") if decoder else None |
| if ds:s.update(ds) |
| with init_empty_weights():m=build(c,decoder,ds) |
| m.load_state_dict(s,strict=True,assign=True);m=m.to(device=device).eval() |
| |
| |
| if decoder:m.image_decoder.float();m.image_refiner.float() |
| return m |
|
|
| def patchify_pixels(pixels): |
| """NCHW RGB in [0,1]; no resize. Spatial dimensions must be divisible by 28.""" |
| if pixels.ndim!=4 or pixels.shape[1]!=3 or any(v%28 for v in pixels.shape[-2:]):raise ValueError("Expected NCHW RGB with dimensions divisible by 28") |
| if pixels.min()<0 or pixels.max()>1:raise ValueError("Expected pixels in [0,1]") |
| mean=pixels.new_tensor([.48145466,.4578275,.40821073])[None,:,None,None];std=pixels.new_tensor([.26862954,.26130258,.27577711])[None,:,None,None] |
| x=(pixels-mean)/std;patches=[];grids=[] |
| for image in x: |
| _,h,w=image.shape;gh,gw=h//14,w//14 |
| p=image.unsqueeze(0).expand(2,-1,-1,-1).reshape(1,2,3,gh//2,2,14,gw//2,2,14).permute(0,3,6,4,7,2,1,5,8).reshape(gh*gw,-1) |
| patches.append(p);grids.append([1,gh,gw]) |
| return torch.cat(patches),torch.tensor(grids,device=pixels.device,dtype=torch.long) |
|
|
| def encode_pixels(model,pixels): |
| patches,grid=patchify_pixels(pixels);ids=model.model.visual_tokenizer.encode(patches,grid) |
| return ids,grid |
|
|
| def code_features(model,ids): |
| if ids.shape[-1]!=8 or ids.min()<0 or ids.max()>=16384:raise ValueError("Expected 8 residual-stage IDs in [0,16383]") |
| books=model.model.visual_tokenizer.visual_bridge_model.quantizer.quantize.codebooks |
| return sum(books[i].embed[ids[...,i].long()] for i in range(8)) |
|
|
| def decode_codes(model,ids,grid_thw): |
| features=code_features(model,ids).reshape(-1,3584).to(next(model.image_decoder.parameters())) |
| decoded=model.image_decoder(features,grid_thw)["images"] |
| return decoded |
|
|
|
|
| def load_refiner(path=None,device="cpu"): |
| """Load the complete native flow Transformer, refiner VAE and conditioning projection.""" |
| from native.image_refiner import ImageRefinerContainer |
| root=Path(path or Path(__file__).parent);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) |
| return m.to(device=device).eval() |
|
|
| def refine_codes(codec,refiner,ids,grid_thw,seed=1019,num_inference_steps=28,output_type="pt"): |
| """Native flow refinement of the base reconstruction. Requires both bundles loaded.""" |
| from native.image_refiner import RefinerPipeline,RefinerImageProcessor,de_transform |
| from native.refiner_modules import FlowMatchEulerDiscreteScheduler |
| device=next(refiner.parameters()).device;cfg=json.loads((Path(__file__).parent/"source_config.json").read_text())["visual_config"]["visual_decoder_config"] |
| scheduler=FlowMatchEulerDiscreteScheduler(**cfg["scheduler_config"]) |
| pipeline=RefinerPipeline(vae=refiner.vae,transformer=refiner.base_transformer,cond_proj=refiner.cond_proj,scheduler=scheduler) |
| raw=decode_codes(codec,ids,grid_thw);images=[] |
| for image in raw: |
| x=de_transform(image,mean=[.48145466,.4578275,.40821073],std=[.26862954,.26130258,.27577711],rescale_factor=1/255) |
| images.append(RefinerImageProcessor.normalize(x).squeeze(0).to(device)) |
| features=code_features(codec,ids) |
| if features.ndim==2:features=features.unsqueeze(0) |
| result=pipeline(encoder_hidden_states=features.to(device=device,dtype=next(refiner.cond_proj.parameters()).dtype),grid_thw_list=[tuple(map(int,g)) for g in grid_thw],image=images,generator=torch.Generator(device=device).manual_seed(seed),num_inference_steps=num_inference_steps,output_type=output_type,return_dict=True,enable_processor_bar=False) |
| return result.images |
|
|
