File size: 3,279 Bytes
a130d59 | 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 | """Load the complete InternVL-U image-generation decoder with native weights."""
import json
import sys
from pathlib import Path
import torch
from safetensors.torch import load_file
class ImageDiT:
def __init__(self, native, root):
self.native = native
self.root = root
@torch.inference_mode()
def predict(self, latents, timestep, condition, *, condition_mask=None,
condition_image_mask=None, condition_grid=None, reference_latents=None):
if latents.ndim != 4 or latents.shape[1] != 16 or any(x % 2 for x in latents.shape[-2:]):
raise ValueError('Expected even spatial dimensions and [B,16,H,W] native VAE latents')
if condition.ndim != 3 or condition.shape[0] != latents.shape[0] or condition.shape[2] != 4096:
raise ValueError('Expected original InternVL-U [B,L,4096] conditioning features')
if condition_mask is None:
condition_mask = torch.ones(condition.shape[:2], dtype=torch.bool, device=condition.device)
if condition_image_mask is None:
condition_image_mask = torch.zeros_like(condition_mask, dtype=torch.bool)
if condition_grid is None:
if condition_image_mask.any():
raise ValueError('Visual conditioning tokens require their native image grid')
condition_grid = [torch.empty((0, 3), dtype=torch.long, device=condition.device) for _ in range(len(latents))]
timestep = torch.as_tensor(timestep, device=latents.device, dtype=latents.dtype)
if timestep.ndim == 0:
timestep = timestep.expand(len(latents))
return self.native.decoder(
hidden_states=latents, encoder_hidden_states=condition, timestep=timestep,
encoder_attention_mask=condition_mask, encoder_image_token_mask=condition_image_mask,
image_fhw_cond=condition_grid, conditional_input=reference_latents, return_dict=False,
)[0]
def make_scheduler(self):
from diffusers import DPMSolverMultistepScheduler
return DPMSolverMultistepScheduler.from_config(json.loads((self.root / 'scheduler_config.json').read_text()))
def load_component(root=None, *, dtype=torch.float32, device='cpu'):
root = Path(root or Path(__file__).resolve().parent)
sys.path.insert(0, str(root))
from intern_native.configuration_internvlu_generation_decoder import InternVLUGenerationDecoderConfig
from intern_native.modeling_internvlu_generation_decoder import InternVLUGenerationDecoder
config = InternVLUGenerationDecoderConfig(**json.loads((root / 'config.json').read_text()))
with torch.device('meta'):
native = InternVLUGenerationDecoder(config)
state = load_file(root / 'model.safetensors')
native.load_state_dict(state, strict=True, assign=True)
del state
# This native module stores deterministic RoPE frequencies as a plain
# tensor attribute. They are absent from state_dict and remain on meta
# after assign=True, so reconstruct them with the upstream constructor.
rope = native.decoder.pos_embed
native.decoder.pos_embed = type(rope)(theta=rope.theta, axes_dim=rope.axes_dim, scale_rope=rope.scale_rope)
native.to(device=device, dtype=dtype).eval()
return ImageDiT(native, root)
|