| """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 |
| |
| |
| |
| 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) |
|
|
