VideoX-Fun/examples/z_image_fun/launcher_axmodel.py

#!/usr/bin/env python3
"""使用 AXModel 推理链路 (transformer + VAE decoder)。"""
from __future__ import annotations

import argparse
import json
import os
import random
import sys
from contextlib import contextmanager
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple, Union

current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
    sys.path.insert(0, project_root) if project_root not in sys.path else None

SCRIPT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SCRIPT_DIR.parents[2]
if REPO_ROOT.as_posix() not in sys.path:
    sys.path.insert(0, REPO_ROOT.as_posix())

import numpy as np
import torch
from axengine import InferenceSession as AxInferenceSession
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils.torch_utils import randn_tensor
from omegaconf import OmegaConf
from PIL import Image
from loguru import logger
from tqdm import tqdm

from videox_fun.models import AutoTokenizer, Qwen3ForCausalLM
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from videox_fun.utils.utils import get_image_latent


# -----------------------------------------------------------------------------
# 模型与资源路径
# -----------------------------------------------------------------------------
MODEL_NAME = "models/Diffusion_Transformer/Z-Image-Turbo/"
CONFIG_PATH = REPO_ROOT / "VideoX-Fun" / "config" / "z_image" / "z_image.yaml"
TRANSFORMER_CONFIG_PATH = REPO_ROOT / "VideoX-Fun" / "pulsar2_configs" / "transformers_subgraph.json"
TRANSFORMER_ONNX_PATH = REPO_ROOT / "VideoX-Fun" / "compiled_subgraph_from_onnx" / "frontend" / "optimized_quant_axmodel.onnx"
TRANSFORMER_AXMODEL_DIR = REPO_ROOT / "VideoX-Fun" / "comliled_subgraph_from_all_onnx"  # compiled_slice_quant_onnx
VAE_DECODER_AXMODEL = REPO_ROOT / "VideoX-Fun" / "vae_decoder.axmodel"
SAVE_DIR = REPO_ROOT / "VideoX-Fun" / "samples" / "z-image-t2i-axmodel"

# -----------------------------------------------------------------------------
# 运行配置
# -----------------------------------------------------------------------------
DEFAULT_PROMPTS = [
    "(masterpiece, best quality) solo female on a tropical beach, golden hour rim light, cinematic grading",
    "nighttime cyberpunk boulevard, neon reflections on wet asphalt, volumetric fog, wide shot",
    "sunrise over alpine mountains, low clouds in valleys, god rays, ultra-detailed landscape",
    "modern minimal living room, soft natural light, Scandinavian design, high-resolution interior render",
    "classical oil painting of a renaissance noblewoman, chiaroscuro lighting, rich textures",
    "macro photography of a dewdrop on a leaf, extreme detail, shallow depth of field",
    "futuristic sports car parked under neon lights, glossy paint, cinematic 35mm look",
    "ancient library with towering bookshelves, warm candlelight, dust motes in air",
    "portrait of an astronaut in full suit, visor reflection showing earth, studio lighting",
    "stormy sea with a lone lighthouse, crashing waves, dramatic clouds, long exposure feel",
    "cybernetic samurai standing in rain, backlit silhouette, moody blue-orange palette",
    "lush rainforest waterfall, soft mist, saturated greens, wide-angle composition",
    "product shot of a smartwatch on marble, softbox lighting, crisp shadows, advertisement style",
    "architectural exterior of a glass skyscraper at dusk, warm interior lights, reflections",
    "vintage film photograph of a 1950s diner at night, grain and halation, neon signage",
    "hyperrealistic bowl of ramen, steam rising, glossy broth, detailed toppings",
    "fantasy castle on a floating island, waterfalls falling into clouds, sunset lighting",
    "high-fashion editorial portrait, dramatic chiaroscuro, sharp focus on eyes",
    "aerial view of winding river through autumn forest, golden and crimson leaves",
    "studio shot of running shoes mid-air, motion blur trails, vibrant background gradient",
    "noir city alley in the 1940s, hard shadows, rain-slick pavement, moody atmosphere",
    "desert caravan at twilight, silhouettes of camels, soft purple sky, cinematic scope",
    "close-up of a mechanical watch movement, intricate gears, metallic reflections",
    "bioluminescent underwater reef, glowing corals, schools of fish, deep blue tones",
    "portrait of an elderly man with weathered face, soft window light, fine skin detail",
    "snowy village at night, warm cabin lights, smoke from chimneys, peaceful mood",
    "futuristic data center aisle, cool cyan lighting, depth and symmetry",
    "oil painting of a bowl of fruit in Dutch masters style, rich textures, dramatic lighting",
    "sunlit meadow with wildflowers, shallow depth of field, pastel color palette",
    "sci-fi corridor with volumetric light shafts, pristine white surfaces, wide lens",
    "luxury wristwatch on black velvet, high contrast, advertisement macro shot",
    "medieval marketplace at dawn, merchants setting up, soft warm light, lively details",
    "((masterpiece,best quality))1 young beautiful girl,ultra detailed,official art,unity 8k wallpaper,masterpiece, best quality, official art, extremely detailed CG unity 8k wallpaper, highly detailed, 1 girl, aqua eyes, light smile, ((grey hair)), hair flower, bracelet, choker, ribbon, JK, look at viewer, on the beach, in summer,"
]
prompt_idx = random.randint(0, len(DEFAULT_PROMPTS) - 1)
PROMPT = DEFAULT_PROMPTS[prompt_idx]
NEG_PROMPT = " "
GUIDANCE_SCALE = 0.0
SEED = 42
HEIGHT, WIDTH = 512, 512
NUM_INFERENCE_STEPS = 9
NUM_CHANNELS_LATENTS = 16
VAE_SCALE_FACTOR = 8
PATCH_SIZE = 2
FPATCH_SIZE = 1
MAX_SEQ_LEN = 128
VAE_SCALING_FACTOR = 0.3611
VAE_SHIFT_FACTOR = 0.1159

SAMPLER_MAP = {
    "Flow": FlowMatchEulerDiscreteScheduler,
    "Flow_Unipc": FlowUniPCMultistepScheduler,
    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
}
SAMPLER_NAME = "Flow"

# 默认最终输出，如果不存在 auto 子图则回退到最后一个 cfg 输出
DEFAULT_FINAL_OUTPUT = None

# -----------------------------------------------------------------------------
# 工具函数 (复制自原 launcher 并微调)
# -----------------------------------------------------------------------------

def _infer_module_device(module: torch.nn.Module) -> torch.device:
    param = next(module.parameters(), None)
    if param is not None:
        return param.device
    buffer = next(module.buffers(), None)
    if buffer is not None:
        return buffer.device
    return torch.device("cpu")


@contextmanager
def module_to_device(module: torch.nn.Module, target_device: torch.device):
    if module is None:
        yield module
        return
    original_device = _infer_module_device(module)
    target_device = target_device or original_device
    needs_move = original_device != target_device
    moved_to_cuda = needs_move and target_device.type == "cuda"
    if needs_move:
        module.to(target_device)
    try:
        yield module
    finally:
        if needs_move:
            module.to(original_device)
            if moved_to_cuda and torch.cuda.is_available():
                cache_device = target_device.index or torch.cuda.current_device()
                with torch.cuda.device(cache_device):
                    torch.cuda.empty_cache()


def _encode_prompt(
    tokenizer: AutoTokenizer,
    text_encoder: Qwen3ForCausalLM,
    prompt: Union[str, List[str]],
    device: torch.device,
    prompt_embeds: Optional[List[torch.FloatTensor]] = None,
    max_sequence_length: int = 512,
) -> List[torch.FloatTensor]:
    if prompt_embeds is not None:
        return prompt_embeds
    prompts = [prompt] if isinstance(prompt, str) else list(prompt)
    for idx, item in enumerate(prompts):
        messages = [{"role": "user", "content": item}]
        prompts[idx] = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
        )
    text_inputs = tokenizer(
        prompts,
        padding="max_length",
        max_length=max_sequence_length,
        truncation=True,
        return_tensors="pt",
    )
    text_input_ids = text_inputs.input_ids.to(device)
    prompt_masks = text_inputs.attention_mask.to(device).bool()
    with module_to_device(text_encoder, device):
        prompt_embeds = text_encoder(
            input_ids=text_input_ids,
            attention_mask=prompt_masks,
            output_hidden_states=True,
        ).hidden_states[-2]
    return [prompt_embeds[i] for i in range(len(prompt_embeds))]


def encode_prompt(
    tokenizer: AutoTokenizer,
    text_encoder: Qwen3ForCausalLM,
    prompt: Union[str, List[str]],
    device: torch.device,
    do_classifier_free_guidance: bool,
    negative_prompt: Optional[Union[str, List[str]]],
    max_sequence_length: int,
) -> Tuple[List[torch.FloatTensor], List[torch.FloatTensor]]:
    prompt_embeds = _encode_prompt(
        tokenizer, text_encoder, prompt, device, None, max_sequence_length
    )
    negative_embeds: List[torch.FloatTensor] = []
    if do_classifier_free_guidance:
        neg = negative_prompt or ""
        negative_list = [neg] if isinstance(neg, str) else list(neg)
        negative_embeds = _encode_prompt(
            tokenizer, text_encoder, negative_list, device, None, max_sequence_length
        )
    return prompt_embeds, negative_embeds


def _stack_prompt_embeddings(prompt_embeds_input):
    if isinstance(prompt_embeds_input, list):
        return torch.stack(prompt_embeds_input, dim=0)
    return prompt_embeds_input


def prepare_latents(
    batch_size: int,
    num_channels_latents: int,
    height: int,
    width: int,
    dtype: torch.dtype,
    device: torch.device,
    generator: torch.Generator,
) -> torch.FloatTensor:
    height = 2 * (int(height) // (VAE_SCALE_FACTOR * 2))
    width = 2 * (int(width) // (VAE_SCALE_FACTOR * 2))
    shape = (batch_size, num_channels_latents, height, width)
    latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
    return latents


def calculate_shift(
    image_seq_len: int,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
) -> float:
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    return image_seq_len * m + b


def retrieve_timesteps(
    scheduler,
    num_inference_steps: int,
    device: torch.device,
    **kwargs,
):
    scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
    return scheduler.timesteps


# -----------------------------------------------------------------------------
# 参数
# -----------------------------------------------------------------------------


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="AXModel 推理 (transformer + VAE)")
    parser.add_argument("--prompt", type=str, default=None, help="正向提示词，不填则使用预置随机样本")
    parser.add_argument("--negative-prompt", type=str, default=NEG_PROMPT, help="反向提示词")
    parser.add_argument("--steps", type=int, default=NUM_INFERENCE_STEPS, help="迭代步数")
    parser.add_argument("--height", type=int, default=HEIGHT, help="生成高度，需被 16 整除")
    parser.add_argument("--width", type=int, default=WIDTH, help="生成宽度，需被 16 整除")
    parser.add_argument("--seed", type=int, default=SEED, help="随机种子")
    parser.add_argument("--sampler", type=str, choices=list(SAMPLER_MAP.keys()), default=SAMPLER_NAME, help="采样器")
    parser.add_argument("--max-seq-len", type=int, default=MAX_SEQ_LEN, help="最大文本长度")
    parser.add_argument("--save-dir", type=str, default=str(SAVE_DIR), help="结果输出目录")
    parser.add_argument("--transformer-config", type=str, required=True, help="子图配置 json")
    parser.add_argument("--transformer-onnx", type=str, default=None, help="原始 transformer onnx（可选，sub_configs 已覆盖可不填）")
    parser.add_argument("--transformer-subgraph-dir", type=str, required=True, help="子图 axmodel 目录")
    parser.add_argument("--vae-axmodel", type=str, required=True, help="VAE decoder axmodel 路径")
    parser.add_argument("--final-output-name", type=str, default=None, help="指定最终输出 tensor 名称，默认自动推断")
    parser.add_argument("--save-decoder-input", action="store_true", help="是否保存 decoder 输入 npy")
    parser.add_argument("--no-progress", action="store_true", help="关闭进度条输出")
    return parser.parse_args()


# -----------------------------------------------------------------------------
# AX transformer 子图执行器
# -----------------------------------------------------------------------------

from scripts.split_quant_onnx_by_subconfigs import SubGraphSpec, sanitize


class AxSplitTransformer:
    def __init__(self, config_path: Path, onnx_path: Optional[Path], model_dir: Path):
        self.config_path = config_path
        self.onnx_path = onnx_path
        self.model_dir = model_dir
        self._session_cache: Dict[str, AxInferenceSession] = {}
        config_specs = self._load_specs()
        auto_specs = self._load_auto_specs()
        self.specs = config_specs + auto_specs
        last_group = auto_specs if auto_specs else config_specs
        self.final_outputs = list(last_group[-1].end)
        self.default_final_output = DEFAULT_FINAL_OUTPUT or self.final_outputs[0]

    def _get_session(self, spec: SubGraphSpec) -> AxInferenceSession:
        if spec.label not in self._session_cache:
            path = self._expected_path(spec)
            self._session_cache[spec.label] = AxInferenceSession(path.as_posix())
            logger.info(f"加载子图 session: {spec.label} from {path.name}")
        return self._session_cache[spec.label]

    def close(self) -> None:
        # 显式释放缓存的 session
        for key, sess in list(self._session_cache.items()):
            try:
                del sess
            finally:
                self._session_cache.pop(key, None)

    def _load_specs(self) -> List[SubGraphSpec]:
        with self.config_path.open("r", encoding="utf-8") as f:
            config = json.load(f)
        sub_configs = config.get("compiler", {}).get("sub_configs", [])
        if not sub_configs:
            raise ValueError("配置文件缺少 compiler.sub_configs")
        specs: List[SubGraphSpec] = []
        for idx, entry in enumerate(sub_configs):
            start = [name for name in entry.get("start_tensor_names", []) if name]
            end = [name for name in entry.get("end_tensor_names", []) if name]
            if not start or not end:
                raise ValueError(f"sub_config[{idx}] 定义不完整")
            spec = SubGraphSpec(
                label=f"cfg_{idx:02d}",
                start=start,
                end=end,
                node_names=set(),
                source="config",
            )
            specs.append(spec)
        return specs

    def _load_auto_specs(self) -> List[SubGraphSpec]:
        specs: List[SubGraphSpec] = []
        for path in sorted(self.model_dir.glob("auto_*.axmodel")):
            try:
                session = AxInferenceSession(path.as_posix())
                inputs = [info.name for info in session.get_inputs() if getattr(info, "name", None)]
                outputs = [info.name for info in session.get_outputs() if getattr(info, "name", None)]
                # 缓存 session，避免重复打开
                self._session_cache[path.stem] = session
            except Exception as exc:  # pragma: no cover - defensive
                logger.warning(f"跳过 {path.name}，加载/解析 IO 失败: {exc}")
                continue

            if not inputs or not outputs:
                logger.warning(f"跳过 {path.name}，未找到有效的输入/输出定义")
                continue

            specs.append(
                SubGraphSpec(
                    label=path.stem,
                    start=inputs,
                    end=outputs,
                    node_names=set(),
                    source="auto",
                    output_path=path,
                )
            )
        return specs

    def _expected_path(self, spec: SubGraphSpec) -> Path:
        if spec.output_path is not None:
            path = spec.output_path
        else:
            head = sanitize(spec.start[0]) if spec.start else "const"
            tail = sanitize(spec.end[0]) if spec.end else "out"
            filename = f"{spec.label}_{head}_to_{tail}_{spec.source}.axmodel"
            path = self.model_dir / filename
        if not path.exists():
            raise FileNotFoundError(f"缺少 AXModel: {path}")
        return path

    def __call__(
        self,
        latent_np: np.ndarray,
        prompt_np: np.ndarray,
        timestep_np: np.ndarray,
        final_output_name: Optional[str] = None,
    ) -> np.ndarray:
        tensor_store: Dict[str, np.ndarray] = {
            "latent_model_input": latent_np,
            "prompt_embeds": prompt_np,
            "timestep": timestep_np,
        }
        executed = set()
        target = final_output_name or self.default_final_output

        # 就绪驱动执行，单个子图跑完立刻释放 session
        while target not in tensor_store:
            progressed = False
            for spec in self.specs:
                if spec.label in executed:
                    continue
                if not all(name in tensor_store for name in spec.start):
                    continue
                session = self._get_session(spec)
                inputs = {name: tensor_store[name] for name in spec.start}
                outputs = session.run(spec.end, inputs)
                for out_name, value in zip(spec.end, outputs):
                    tensor_store[out_name] = value
                executed.add(spec.label)
                progressed = True
            if not progressed:
                missing = [
                    (spec.label, [name for name in spec.start if name not in tensor_store])
                    for spec in self.specs
                    if spec.label not in executed
                ]
                raise RuntimeError(
                    f"子图调度中断，缺少输入: {missing}; 当前可用: {list(tensor_store.keys())}"
                )

        return tensor_store[target]


# -----------------------------------------------------------------------------
# 主流程
# -----------------------------------------------------------------------------

def main() -> None:
    args = parse_args()
    prompt_text = args.prompt if args.prompt is not None else PROMPT
    logger.info(f"使用的 prompt: {prompt_text}")

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    torch.set_grad_enabled(False)

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder="tokenizer")
    text_encoder = Qwen3ForCausalLM.from_pretrained(
        MODEL_NAME,
        subfolder="text_encoder",
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
    text_encoder.eval()

    scheduler_cls = SAMPLER_MAP[args.sampler]
    scheduler = scheduler_cls.from_pretrained(MODEL_NAME, subfolder="scheduler")

    image_processor = VaeImageProcessor(vae_scale_factor=VAE_SCALE_FACTOR * 2)

    prompt_embeds, _ = encode_prompt(
        tokenizer,
        text_encoder,
        prompt_text,
        device,
        do_classifier_free_guidance=False,
        negative_prompt=args.negative_prompt,
        max_sequence_length=args.max_seq_len,
    )

    latents = prepare_latents(
        batch_size=1,
        num_channels_latents=NUM_CHANNELS_LATENTS,
        height=args.height,
        width=args.width,
        dtype=torch.float32,
        device=device,
        generator=torch.Generator(device=device).manual_seed(args.seed),
    )

    image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
    mu = calculate_shift(
        image_seq_len,
        scheduler.config.get("base_image_seq_len", 256),
        scheduler.config.get("max_image_seq_len", 4096),
        scheduler.config.get("base_shift", 0.5),
        scheduler.config.get("max_shift", 1.15),
    )
    timesteps = retrieve_timesteps(scheduler, args.steps, device=device, mu=mu)

    onnx_path = Path(args.transformer_onnx) if args.transformer_onnx else None
    transformer_runner = AxSplitTransformer(
        Path(args.transformer_config),
        onnx_path,
        Path(args.transformer_subgraph_dir),
    )
    # 优先使用 auto_* 子图里的 sample 输出，避免误用中间特征导致 shape 对不上
    available_outputs = [name for spec in transformer_runner.specs for name in getattr(spec, "end", [])]
    preferred_output = "sample" if "sample" in available_outputs else transformer_runner.default_final_output
    final_output_name = args.final_output_name or preferred_output
    if final_output_name not in available_outputs:
        raise ValueError(f"指定的输出 {final_output_name} 不存在，可选: {available_outputs}")

    prompt_embeds_tensor = _stack_prompt_embeddings(prompt_embeds)

    iterator = timesteps if args.no_progress else tqdm(timesteps, desc="AX Denoising", dynamic_ncols=True)
    for t in iterator:
        timestep = t.expand(latents.shape[0])
        timestep_model_input = (1000 - timestep) / 1000

        latent_model_input = latents.to(torch.float32)
        latent_np = latent_model_input.unsqueeze(2).to(dtype=torch.float32).cpu().numpy()
        prompt_np = prompt_embeds_tensor.to(dtype=torch.float32).cpu().numpy()
        timestep_np = timestep_model_input.to(dtype=torch.float32).cpu().numpy()

        model_out = transformer_runner(latent_np, prompt_np, timestep_np, final_output_name)
        if model_out.ndim == 5 and model_out.shape[2] == 1:
            model_out = np.squeeze(model_out, axis=2)
        model_out_tensor = torch.from_numpy(model_out).to(device=device, dtype=torch.float32)
        if model_out_tensor.dim() == 5 and model_out_tensor.size(2) == 1:
            model_out_tensor = model_out_tensor.squeeze(2)
        noise_pred = -model_out_tensor

        latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]

    # 释放 transformer 缓存的 session
    transformer_runner.close()

    latents = (latents / VAE_SCALING_FACTOR) + VAE_SHIFT_FACTOR
    decoder_input = latents.to(dtype=torch.float32).cpu().numpy()

    if args.save_decoder_input:
        save_dir_path = Path(args.save_dir)
        save_dir_path.mkdir(parents=True, exist_ok=True)
        np.save(save_dir_path / "decoder_input.npy", decoder_input)
        logger.info("已保存 decoder 输入为 npy")

    del transformer_runner

    vae_decoder_session = AxInferenceSession(Path(args.vae_axmodel).as_posix())

    if decoder_input.ndim == 5 and decoder_input.shape[2] == 1:
        decoder_input = np.squeeze(decoder_input, axis=2)
    image = vae_decoder_session.run(None, {"latents": decoder_input})[0]

    del vae_decoder_session
    image = torch.from_numpy(image).to(device=device, dtype=torch.float32)
    image = image_processor.postprocess(image, output_type="pil")

    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    target_path = save_dir / f"z_image_axmodel_{prompt_idx}.png"
    image[0].save(target_path)
    logger.info(f"AXModel 推理完成，结果保存到 {target_path}")


if __name__ == "__main__":
    """
    # 512x512 生成示例命令：
    python3 examples/z_image_fun/launcher_axmodel.py \
        --transformer-config pulsar2_configs/transformers_subgraph.json \
        --transformer-subgraph-dir comliled_subgraph_from_all_onnx \
        --vae-axmodel vae_decoder.axmodel

    # 1728x992 生成示例命令:
    python3 examples/z_image_fun/launcher_axmodel.py \
        --transformer-config pulsar2_configs/transformers_subgraph_1728x992.json \
        --transformer-subgraph-dir transformers_body_only_1728_992_split_onnx \
        --vae-axmodel onnx-models-1728x992/vae_decoder_simp_slim.axmodel \
        --max-seq-len 256 \
        --height 1728 --width 992
    """

    main()