| """Native Show-o2 visual I/O modules, with the Qwen backbone supplied externally.""" |
| import json, math |
| from pathlib import Path |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from transformers.models.siglip.configuration_siglip import SiglipVisionConfig |
| from native.modeling_siglip import SiglipEncoder |
| from native.modules import (PatchEmbed,RMSNorm,TimestepEmbedder,DiffusionHeadConfig, |
| ModulatedAttentionBlock,FinalLayer) |
|
|
| class VisualFlowComponents(nn.Module): |
| def __init__(self,config,siglip_config): |
| super().__init__() |
| self.config=config |
| p=config['patch_size'];c=config['image_latent_dim'] |
| d=config['hidden_size'];s=config['clip_latent_dim'] |
| self.image_embedder_und=PatchEmbed(patch_size=p,in_chans=c,embed_dim=s) |
| self.image_embedder_gen=PatchEmbed(patch_size=p,in_chans=c,embed_dim=d) |
| vision=SiglipVisionConfig(**siglip_config['vision_config']) |
| vision._attn_implementation='sdpa' |
| vision.num_hidden_layers-=1 |
| self.position_embedding=nn.Embedding((vision.image_size//vision.patch_size)**2,s) |
| self.und_trans=SiglipEncoder(vision) |
| self.fusion_proj=nn.Sequential(RMSNorm(s+d),nn.Linear(s+d,d),nn.GELU(),nn.Linear(d,d)) |
| head=DiffusionHeadConfig() |
| self.time_embed=TimestepEmbedder(head.hidden_size) |
| self.diff_proj=nn.Sequential(nn.Linear(d,head.hidden_size),nn.GELU(),nn.Linear(head.hidden_size,head.hidden_size)) |
| self.time_embed_proj=nn.Linear(head.hidden_size,d) |
| self.diffusion_head_a=nn.ModuleList([ModulatedAttentionBlock(head,i) for i in range(config['num_diffusion_layers'])]) |
| self.diffusion_head_b=FinalLayer(head.hidden_size,p,c) |
|
|
| def position_encoding(self,height,width): |
| |
| weight=self.position_embedding.weight |
| side=int(math.sqrt(weight.shape[0])) |
| value=weight.unsqueeze(0).reshape(1,side,side,weight.shape[1]).permute(0,3,1,2) |
| value=F.interpolate(value,size=(height,width),mode='bicubic',align_corners=False) |
| return value.permute(0,2,3,1).reshape(1,-1,weight.shape[1]) |
|
|
| def encode_latents(self,latents,t): |
| """One image per batch item: normalized VAE BCHW -> time+visual LLM inputs.""" |
| assert latents.ndim==4 and latents.shape[1]==16 |
| p=self.config['patch_size'];h,w=latents.shape[-2:] |
| assert h%p==0 and w%p==0 |
| und=self.image_embedder_und(latents) |
| if (h//p,w//p)==(27,27): |
| pos=self.position_embedding(torch.arange(729,device=latents.device).unsqueeze(0)) |
| else: |
| pos=self.position_encoding(h//p,w//p) |
| und=self.und_trans(und+pos,return_dict=True).last_hidden_state |
| gen=self.image_embedder_gen(latents) |
| fused=self.fusion_proj(torch.cat([und,gen],dim=-1)) |
| time=self.time_embed(t,latents.dtype) |
| return torch.cat([self.time_embed_proj(time).unsqueeze(1),fused],dim=1) |
|
|
| def predict_velocity(self,last_hidden_states,t,modality_positions,attention_mask=None): |
| """External Qwen final states -> native flow prediction in patch-vector order. |
| |
| modality_positions includes each time token. Use the full sequence's |
| native multimodal attention mask for integration with text/interleaving. |
| """ |
| time=self.time_embed(t,last_hidden_states.dtype) |
| x=self.diff_proj(last_hidden_states) |
| ids=torch.arange(x.shape[1],device=x.device).unsqueeze(0) |
| for layer in self.diffusion_head_a: |
| x=layer(hidden_states=x,adaln_input=time,attention_mask=attention_mask, |
| position_ids=ids,modality_positions=modality_positions)[0] |
| return self.diffusion_head_b(x,time,modality_positions) |
|
|
| def velocity_image(self,patch_vectors,height,width): |
| """Exact native single-image output order; input excludes time token.""" |
| p=self.config['patch_size'];c=self.config['image_latent_dim'] |
| b=patch_vectors.shape[0] |
| x=patch_vectors.reshape(b,height//p,width//p,p,p,c) |
| x=x.reshape(b,height*width,c).transpose(1,2) |
| return x.reshape(b,c,height,width) |
|
|
| def load_model(path=None,dtype=torch.float32,device='cpu'): |
| root=Path(path or Path(__file__).parent) |
| with init_empty_weights(): |
| model=VisualFlowComponents(json.loads((root/'source_config.json').read_text()),json.loads((root/'siglip_config.json').read_text())) |
| model.load_state_dict(load_file(root/'model.safetensors'),strict=True,assign=True) |
| return model.to(device=device,dtype=dtype).eval() |
|
|
| def load_vae(path=None,dtype=torch.float32,device='cpu'): |
| from native.wan21_vae import WanVAE |
| root=Path(path or Path(__file__).parent) |
| return WanVAE(z_dim=16,vae_pth=str(root/'vae/source.pth'),dtype=dtype,device=device) |
|
|
