"""UniSpace inference helpers. Copied verbatim from the official UniSpace release (https://github.com/yjb6/UniSpace): * unispace/eval/gen/gen_images_qwen3_unified_mot.py -> decode_latent, generate_image * unispace/eval/gen/gen_images_imgedit_qwen3_unified_mot.py -> encode_ref_image, prepare_unified_vit_input, compute_target_size_from_ref, editing_image Only the argparse/torchrun driver code was dropped; the inference logic, op ordering and defaults are unchanged. Copyright 2025 Bytedance Ltd. and/or its affiliates. SPDX-License-Identifier: Apache-2.0 """ import os import torch from PIL import Image from modeling.unimm.unimm_mot import UnimmMoT from modeling.unimm.qwen3 import NaiveCache from data.data_utils import get_flattened_position_ids_extrapolate import torch._dynamo torch._dynamo.config.disable = True def move_generation_input_to_device(generation_input, device): for k, v in generation_input.items(): if isinstance(v, torch.Tensor): generation_input[k] = v.to(device) return generation_input def decode_latent(latent, h, w, latent_ch, vae_model, use_merge_gen, device_type): if use_merge_gen: from modeling.unimm.unimm_mot import UnimmMoT latent = UnimmMoT.spatial_unshuffle(latent, h, w) latent = latent.reshape(1, h, w, latent_ch).permute(0, 3, 1, 2) with torch.amp.autocast(device_type, dtype=torch.bfloat16): image = vae_model.decode(latent) return ((image * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy() def generate_image( prompt, num_timesteps=50, cfg_scale=4.0, cfg_interval=None, cfg_renorm_min=0.9, timestep_shift=1.0, max_t=1.0, num_images=1, resolution=256, inference_mode='flash', device=None, device_type='cuda', gen_model=None, tokenizer=None, new_token_ids=None, vae_model=None, seed=None, intermediates_dir=None, ): if cfg_interval is None: cfg_interval = [0, 1.0] generator = torch.Generator("cpu").manual_seed(seed) if seed is not None else None _llm_cfg = gen_model.config.llm_config num_layers = (_llm_cfg.text_config.num_hidden_layers if hasattr(_llm_cfg, 'text_config') and _llm_cfg.text_config is not _llm_cfg else _llm_cfg.num_hidden_layers) past_key_values = NaiveCache(num_layers) newlens = [0] * num_images new_rope = [0] * num_images # 1. encode prompt prefix into KV cache generation_input_text, newlens, new_rope = gen_model.prepare_prompts( curr_kvlens=newlens, curr_rope=new_rope, prompts=[prompt] * num_images, tokenizer=tokenizer, new_token_ids=new_token_ids, ) generation_input_text = move_generation_input_to_device(generation_input_text, device) with torch.no_grad(), torch.amp.autocast(device_type, dtype=torch.bfloat16): past_key_values = gen_model.forward_cache_update_text(past_key_values, **generation_input_text) # 2. prepare VAE latent slots generation_input = gen_model.prepare_vae_latent( curr_kvlens=newlens, curr_rope=new_rope, image_sizes=[(resolution, resolution)] * num_images, new_token_ids=new_token_ids, generator=generator, ) generation_input = move_generation_input_to_device(generation_input, device) # 3. CFG prefix (text-free) cfg_past_key_values = NaiveCache(num_layers) cfg_newlens = [0] * num_images cfg_new_rope = [0] * num_images cfg_prompt_input, cfg_newlens, cfg_new_rope = gen_model.prepare_cfg_prompts( curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope, num_images=num_images, tokenizer=tokenizer, new_token_ids=new_token_ids, ) cfg_prompt_input = move_generation_input_to_device(cfg_prompt_input, device) with torch.no_grad(), torch.amp.autocast(device_type, dtype=torch.bfloat16): cfg_past_key_values = gen_model.forward_cache_update_text(cfg_past_key_values, **cfg_prompt_input) generation_input_cfg = gen_model.prepare_vae_latent_cfg( curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope, image_sizes=[(resolution, resolution)] * num_images, ) generation_input_cfg = move_generation_input_to_device(generation_input_cfg, device) # 4. flow matching denoising with torch.no_grad(), torch.amp.autocast(device_type, dtype=torch.bfloat16): result = gen_model.generate_image( past_key_values=past_key_values, num_timesteps=num_timesteps, max_t=max_t, cfg_text_scale=cfg_scale, cfg_renorm_min=cfg_renorm_min, cfg_interval=cfg_interval, timestep_shift=timestep_shift, cfg_text_past_key_values=cfg_past_key_values, cfg_text_packed_position_ids=generation_input_cfg['cfg_packed_position_ids'], cfg_text_key_values_lens=generation_input_cfg['cfg_key_values_lens'], cfg_text_packed_query_indexes=generation_input_cfg['cfg_packed_query_indexes'], cfg_text_packed_key_value_indexes=generation_input_cfg['cfg_packed_key_value_indexes'], generation_input_text=generation_input_text, inference_mode=inference_mode, return_intermediates=(intermediates_dir is not None), **generation_input, ) if intermediates_dir is not None: unpacked_latent, intermediate_latents = result else: unpacked_latent = result intermediate_latents = [] # 5. VAE decode latent_ch = gen_model.config.vae_config.z_channels # 1280 latent_ds = gen_model.config.vae_config.downsample # 16 h = w = resolution // latent_ds # 256//16 = 16 use_merge_gen = gen_model.config.use_spatial_merge_gen # save intermediate steps if intermediates_dir is not None and intermediate_latents: os.makedirs(intermediates_dir, exist_ok=True) # only decode first image in batch for speed split_sizes = generation_input['packed_seqlens'].tolist() split_sizes = [s - 2 for s in split_sizes] # strip bos/eos for step_idx, t_val, packed_xt in intermediate_latents: per_img = packed_xt.split(split_sizes) latent = per_img[0].float() arr = decode_latent(latent, h, w, latent_ch, vae_model, use_merge_gen, device_type) Image.fromarray(arr).save( os.path.join(intermediates_dir, f"step_{step_idx:03d}_t{t_val:.3f}.png")) image_list = [] for latent in unpacked_latent: arr = decode_latent(latent.float(), h, w, latent_ch, vae_model, use_merge_gen, device_type) image_list.append(Image.fromarray(arr)) return image_list def move_to_device(d, device): for k, v in d.items(): if isinstance(v, torch.Tensor): d[k] = v.to(device) return d # ═══════════════════════════════════════════════════════ # 模型加载 # ═══════════════════════════════════════════════════════ def encode_ref_image(pil_image, vae_model, mot_model, image_size, device, device_type, ref_max_size=None, ref_min_size=256, ref_max_pixels=1_048_576, ref_use_mar=False, ref_mar_resolution=1024): """PIL → RAE und 支路编码 → (merged_tokens, pos_ids) 三种模式（互斥，优先级：ref_use_mar > ref_max_size > 默认）： 默认（ref_max_size=None, ref_use_mar=False）: 强制 resize 到正方形 (image_size, image_size)。 ref_max_size=N: 保持宽高比，MaxLongEdgeMinShortEdgeResize(max_size=N, min_size=ref_min_size, stride=32, max_pixels=ref_max_pixels)，与训练 ref_vit_image_transform_args 一致。 ref_use_mar=True: snap 到最近的 mar_{ref_mar_resolution} bucket（与训练 editref config 一致）， ref 和 target 使用同一套 bucket，分辨率完全对齐。 """ import torchvision.transforms.functional as TF from data.transforms import MaxLongEdgeMinShortEdgeResize img = pil_image.convert('RGB') if ref_use_mar: # mar bucket 模式：snap 到最近的 mar_{ref_mar_resolution} bucket target_h, target_w = compute_target_size_from_ref(pil_image, ref_mar_resolution) img = img.resize((target_w, target_h), Image.BICUBIC) # PIL: (width, height) h_img, w_img = target_h, target_w elif ref_max_size is not None: # 可变模式：保持宽高比，与训练 ref_vit_image_transform_args 一致 resizer = MaxLongEdgeMinShortEdgeResize( max_size=ref_max_size, min_size=ref_min_size, stride=32, max_pixels=ref_max_pixels, ) img = resizer(img) w_img, h_img = img.size # PIL: (width, height) else: # 原有行为：强制正方形 img = img.resize((image_size, image_size), Image.BICUBIC) h_img = w_img = image_size img_tensor = TF.to_tensor(img) * 2.0 - 1.0 # [3, H, W] in [-1,1] img_tensor = img_tensor.unsqueeze(0).to(device) with torch.no_grad(), torch.amp.autocast(device_type, dtype=torch.bfloat16): z = vae_model.encode(img_tensor) # [1, z_ch, h, w] h_grid = h_img // 16 w_grid = w_img // 16 z_tokens = z[0].permute(1, 2, 0).reshape(h_grid * w_grid, -1).to(dtype=torch.bfloat16) if mot_model.config.use_spatial_merge_und: merged_tokens = UnimmMoT.spatial_merge(z_tokens, h_grid, w_grid) patch_size = 32 # 16 * 2 else: merged_tokens = z_tokens patch_size = 16 pos_ids = get_flattened_position_ids_extrapolate( h_img, w_img, patch_size=patch_size, max_num_patches_per_side=mot_model.max_latent_size, ).to(device) return merged_tokens, pos_ids def prepare_unified_vit_input(curr_kvlens, curr_rope, vit_tokens, pos_ids, new_token_ids, device): """打包 forward_cache_update_unified_vit 所需输入（参考 vlmeval_adapter_qwen3_unified_mot.py）""" packed_text_ids, packed_text_indexes = [], [] packed_vit_token_indexes = [] packed_position_ids_list, packed_seqlens, packed_indexes = [], [], [] packed_key_value_indexes = [] _curr = curr = 0 newlens, new_rope_out = [], [] for curr_kvlen, curr_position_id in zip(curr_kvlens, curr_rope): packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) curr += curr_kvlen packed_text_ids.append(new_token_ids['start_of_image']) packed_text_indexes.append(_curr) packed_indexes.append(curr) curr += 1 _curr += 1 num_img_tokens = vit_tokens.shape[0] packed_vit_token_indexes.extend(range(_curr, _curr + num_img_tokens)) packed_indexes.extend(range(curr, curr + num_img_tokens)) curr += num_img_tokens _curr += num_img_tokens packed_text_ids.append(new_token_ids['end_of_image']) packed_text_indexes.append(_curr) packed_indexes.append(curr) curr += 1 _curr += 1 pos_1d = torch.full((num_img_tokens + 2,), curr_position_id, dtype=torch.long) packed_position_ids_list.append(pos_1d) packed_seqlens.append(num_img_tokens + 2) newlens.append(curr_kvlen + num_img_tokens + 2) new_rope_out.append(curr_position_id + num_img_tokens + 2) return { 'packed_text_ids': torch.tensor(packed_text_ids, dtype=torch.long).to(device), 'packed_text_indexes': torch.tensor(packed_text_indexes, dtype=torch.long).to(device), 'packed_vit_tokens': vit_tokens, 'packed_vit_token_indexes': torch.tensor(packed_vit_token_indexes, dtype=torch.long).to(device), 'packed_vit_position_ids': pos_ids, 'packed_position_ids': torch.cat(packed_position_ids_list).to(device), 'packed_seqlens': torch.tensor(packed_seqlens, dtype=torch.int).to(device), 'packed_indexes': torch.tensor(packed_indexes, dtype=torch.long).to(device), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long).to(device), 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int).to(device), }, newlens, new_rope_out # ═══════════════════════════════════════════════════════ # VAE decode # ═══════════════════════════════════════════════════════ def compute_target_size_from_ref(ref_image, resolution, stride=16): """按 ref 图宽高比 snap 到最近的 aspect-ratio bucket，与训练分布对齐。 resolution=1024 → mar_1024，resolution=512 → mar_512，其余按比例计算。 bucket key = H/W，value = [H, W]。 """ from data.data_utils import MULTI_RESOLUTION_MAP bucket_key = f'mar_{resolution}' if bucket_key in MULTI_RESOLUTION_MAP: bucket = MULTI_RESOLUTION_MAP[bucket_key] pil_w, pil_h = ref_image.size # PIL: (width, height) ref_ratio = pil_h / pil_w # H/W best_key = min(bucket.keys(), key=lambda k: abs(float(k) - ref_ratio)) H, W = bucket[best_key] return (int(H), int(W)) else: # 没有对应 bucket，按长边对齐 resolution pil_w, pil_h = ref_image.size if pil_w >= pil_h: target_w = resolution target_h = resolution * pil_h / pil_w else: target_h = resolution target_w = resolution * pil_w / pil_h target_h = max(stride, round(target_h / stride) * stride) target_w = max(stride, round(target_w / stride) * stride) return (int(target_h), int(target_w)) # ═══════════════════════════════════════════════════════ # 编辑推理核心 # ═══════════════════════════════════════════════════════ def editing_image( gen_model, tokenizer, new_token_ids, vae_model, ref_image, prompt, ref_image_size=448, ref_max_size=None, ref_min_size=256, ref_use_mar=False, ref_mar_resolution=1024, target_size=(256, 256), num_timesteps=50, cfg_text_scale=4.0, cfg_interval=None, cfg_renorm_min=0.0, timestep_shift=1.0, device=None, device_type='cuda', seed=None, ): if cfg_interval is None: cfg_interval = [0, 1.0] generator = torch.Generator("cpu").manual_seed(seed) if seed is not None else None _llm_cfg = gen_model.config.llm_config num_layers = (_llm_cfg.text_config.num_hidden_layers if hasattr(_llm_cfg, 'text_config') and _llm_cfg.text_config is not _llm_cfg else _llm_cfg.num_hidden_layers) h, w = target_size # --- ref 图 und 编码（条件路径）--- vit_tokens, pos_ids = encode_ref_image( ref_image, vae_model, gen_model, ref_image_size, device, device_type, ref_max_size=ref_max_size, ref_min_size=ref_min_size, ref_use_mar=ref_use_mar, ref_mar_resolution=ref_mar_resolution) # ========== conditional 路径 ========== past_key_values = NaiveCache(num_layers) newlens = [0] new_rope = [0] vit_inp, newlens, new_rope = prepare_unified_vit_input( newlens, new_rope, vit_tokens, pos_ids, new_token_ids, device) with torch.no_grad(), torch.amp.autocast(device_type, enabled=True, dtype=torch.bfloat16): past_key_values = gen_model.forward_cache_update_unified_vit(past_key_values, **vit_inp) # user\n{instruction} inp_text, newlens, new_rope = gen_model.prepare_prompts( curr_kvlens=newlens, curr_rope=new_rope, prompts=[f"user\n{prompt}"], tokenizer=tokenizer, new_token_ids=new_token_ids, raw=True, ) inp_text = move_to_device(inp_text, device) with torch.no_grad(), torch.amp.autocast(device_type, enabled=True, dtype=torch.bfloat16): past_key_values = gen_model.forward_cache_update_text(past_key_values, **inp_text) # assistant\n inp_asst, newlens, new_rope = gen_model.prepare_prompts( curr_kvlens=newlens, curr_rope=new_rope, prompts=["assistant\n"], tokenizer=tokenizer, new_token_ids=new_token_ids, raw=True, ) inp_asst = move_to_device(inp_asst, device) with torch.no_grad(), torch.amp.autocast(device_type, enabled=True, dtype=torch.bfloat16): past_key_values = gen_model.forward_cache_update_text(past_key_values, **inp_asst) # 目标图 latent slot generation_input = gen_model.prepare_vae_latent( curr_kvlens=newlens, curr_rope=new_rope, image_sizes=[(h, w)], new_token_ids=new_token_ids, generator=generator, ) generation_input = move_to_device(generation_input, device) # ========== unconditional 路径（只保留 ref 图，drop 指令）========== cfg_past_key_values = NaiveCache(num_layers) cfg_newlens = [0] cfg_new_rope = [0] vit_inp_cfg, cfg_newlens, cfg_new_rope = prepare_unified_vit_input( cfg_newlens, cfg_new_rope, vit_tokens, pos_ids, new_token_ids, device) with torch.no_grad(), torch.amp.autocast(device_type, enabled=True, dtype=torch.bfloat16): cfg_past_key_values = gen_model.forward_cache_update_unified_vit(cfg_past_key_values, **vit_inp_cfg) # 只有 assistant\n，跳过 user 指令 inp_asst_cfg, cfg_newlens, cfg_new_rope = gen_model.prepare_prompts( curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope, prompts=["assistant\n"], tokenizer=tokenizer, new_token_ids=new_token_ids, raw=True, ) inp_asst_cfg = move_to_device(inp_asst_cfg, device) with torch.no_grad(), torch.amp.autocast(device_type, enabled=True, dtype=torch.bfloat16): cfg_past_key_values = gen_model.forward_cache_update_text(cfg_past_key_values, **inp_asst_cfg) generation_input_cfg = gen_model.prepare_vae_latent_cfg( curr_kvlens=cfg_newlens, curr_rope=cfg_new_rope, image_sizes=[(h, w)], ) generation_input_cfg = move_to_device(generation_input_cfg, device) # ========== flow matching 生成 ========== with torch.no_grad(), torch.amp.autocast(device_type, enabled=True, dtype=torch.bfloat16): unpacked_latent = gen_model.generate_image( past_key_values=past_key_values, num_timesteps=num_timesteps, cfg_text_scale=cfg_text_scale, cfg_interval=cfg_interval, cfg_renorm_min=cfg_renorm_min, timestep_shift=timestep_shift, cfg_text_past_key_values=cfg_past_key_values, cfg_text_packed_position_ids=generation_input_cfg["cfg_packed_position_ids"], cfg_text_key_values_lens=generation_input_cfg["cfg_key_values_lens"], cfg_text_packed_query_indexes=generation_input_cfg["cfg_packed_query_indexes"], cfg_text_packed_key_value_indexes=generation_input_cfg["cfg_packed_key_value_indexes"], inference_mode='flash', **generation_input, ) # ========== VAE decode ========== latent_ch = gen_model.config.vae_config.z_channels # 1280 latent_ds = gen_model.config.vae_config.downsample # 16 h_lat = h // latent_ds w_lat = w // latent_ds use_merge_gen = getattr(gen_model.config, 'use_spatial_merge_gen', False) arr = decode_latent( unpacked_latent[0].float(), h_lat, w_lat, latent_ch, vae_model, use_merge_gen, device_type) return Image.fromarray(arr)