VideoX-Fun/examples/z_image_fun/gradio_app.py

import os
import sys
import random
from pathlib import Path
from typing import Optional, Tuple

import gradio as gr
import torch
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from safetensors.torch import load_file

# 将仓库根目录加入 sys.path，方便直接运行本脚本
CURRENT_FILE = Path(__file__).resolve()
PROJECT_ROOTS = [CURRENT_FILE.parent, CURRENT_FILE.parent.parent, CURRENT_FILE.parent.parent.parent]
for root in PROJECT_ROOTS:
    root_str = str(root)
    if root_str not in sys.path:
        sys.path.insert(0, root_str)
REPO_ROOT = PROJECT_ROOTS[-1]

from videox_fun.models import (  # noqa: E402
    AutoencoderKL,
    AutoTokenizer,
    Qwen3ForCausalLM,
    ZImageControlTransformer2DModel,
)
from videox_fun.pipeline import ZImageControlPipeline  # noqa: E402
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler  # noqa: E402
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler  # noqa: E402

# 基础配置，可按需修改
CONFIG_PATH = REPO_ROOT / "config" / "z_image" / "z_image_control.yaml"
MODEL_NAME = REPO_ROOT / "models" / "Diffusion_Transformer" / "Z-Image-Turbo"
TRANSFORMER_CKPT = REPO_ROOT / "models" / "Personalized_Model" / "Z-Image-Turbo-Fun-Controlnet-Union.safetensors"
DEFAULT_POSE_PATH = REPO_ROOT / "asset" / "pose_1024x1024.png"
SAMPLERS = {
    "Flow": FlowMatchEulerDiscreteScheduler,
    "Flow_Unipc": FlowUniPCMultistepScheduler,
    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
}
DEFAULT_SAMPLER = "Flow"

PIPELINE: Optional[ZImageControlPipeline] = None
PIPELINE_DEVICE: Optional[torch.device] = None
PIPELINE_DTYPE: Optional[torch.dtype] = None


def _pick_dtype() -> torch.dtype:
    if torch.cuda.is_available():
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    return torch.float32


def _load_pipeline() -> Tuple[ZImageControlPipeline, torch.device, torch.dtype]:
    global PIPELINE, PIPELINE_DEVICE, PIPELINE_DTYPE
    if PIPELINE is not None:
        return PIPELINE, PIPELINE_DEVICE, PIPELINE_DTYPE

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    dtype = _pick_dtype()

    config = OmegaConf.load(CONFIG_PATH)

    transformer = ZImageControlTransformer2DModel.from_pretrained(
        MODEL_NAME,
        subfolder="transformer",
        low_cpu_mem_usage=True,
        torch_dtype=dtype,
        transformer_additional_kwargs=OmegaConf.to_container(config["transformer_additional_kwargs"]),
    ).to(dtype)

    if TRANSFORMER_CKPT.exists():
        if TRANSFORMER_CKPT.suffix == ".safetensors":
            state_dict = load_file(TRANSFORMER_CKPT)
        else:
            state_dict = torch.load(TRANSFORMER_CKPT, map_location="cpu")
        if "state_dict" in state_dict:
            state_dict = state_dict["state_dict"]
        missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
        print(f"[load] transformer ckpt loaded, missing={len(missing)}, unexpected={len(unexpected)}")
    else:
        print(f"[warn] transformer checkpoint not found at {TRANSFORMER_CKPT}, using base weights")

    vae = AutoencoderKL.from_pretrained(MODEL_NAME, subfolder="vae").to(dtype)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder="tokenizer")
    text_encoder = Qwen3ForCausalLM.from_pretrained(
        MODEL_NAME,
        subfolder="text_encoder",
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
    )

    scheduler_cls = SAMPLERS.get(DEFAULT_SAMPLER, FlowMatchEulerDiscreteScheduler)
    scheduler = scheduler_cls.from_pretrained(MODEL_NAME, subfolder="scheduler")

    pipe = ZImageControlPipeline(
        vae=vae,
        tokenizer=tokenizer,
        text_encoder=text_encoder,
        transformer=transformer,
        scheduler=scheduler,
    )

    if torch.cuda.is_available():
        pipe.to(device=device)
    else:
        # CPU 推理较慢，这里保留在 CPU 上避免超出显存
        pipe.to(device)

    PIPELINE = pipe
    PIPELINE_DEVICE = device
    PIPELINE_DTYPE = dtype
    return pipe, device, dtype


def _ensure_pose_image(pose_image: Optional[Image.Image]) -> Image.Image:
    if pose_image is None:
        return Image.open(DEFAULT_POSE_PATH).convert("RGB")
    if pose_image.mode != "RGB":
        pose_image = pose_image.convert("RGB")
    return pose_image


def _align_size(value: int) -> int:
    # pipeline 要求可被 16 整除
    return max(256, (value // 16) * 16)


def infer(
    prompt: str,
    negative_prompt: str,
    pose_image: Optional[Image.Image],
    height: int,
    width: int,
    steps: int,
    guidance_scale: float,
    control_strength: float,
    seed: int,
):
    pipe, device, _ = _load_pipeline()

    if not prompt.strip():
        raise gr.Error("提示词不能为空")

    pose_image = _ensure_pose_image(pose_image)
    height = _align_size(height)
    width = _align_size(width)

    if seed is None or seed < 0:
        seed = random.randint(1, 2**31 - 1)
    generator = torch.Generator(device=device).manual_seed(seed)

    with torch.inference_mode():
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            generator=generator,
            control_image=pose_image,
            control_context_scale=float(control_strength),
            max_sequence_length=128,
        ).images[0]

    return result, seed


def build_ui():
    css = """
    .compact-slider {padding-top: 4px; padding-bottom: 4px;}
    """

    with gr.Blocks(title="Z-Image Turbo 文生图 (Pose)", css=css) as demo:
        gr.Markdown("## Z-Image Turbo 文生图 (含 Pose 控制)")
        gr.Markdown(
            "上传姿态图，输入提示词即可生成图像。右侧为缩略图预览，可放大/下载原分辨率。",
        )

        with gr.Row():
            with gr.Column(scale=1, min_width=320):
                prompt = gr.Textbox(
                    label="提示词",
                    placeholder="描述你想生成的画面",
                    lines=4,
                    value="1 girl, on the beach, summer, full body, highly detailed",
                )
                negative_prompt = gr.Textbox(
                    label="反向提示词",
                    placeholder="不希望出现的元素，例如 '低质量, 模糊'",
                    lines=3,
                    value="lowres, blurry, text, watermark",
                )
                steps = gr.Slider(minimum=4, maximum=30, step=1, value=9, label="采样步数", elem_classes=["compact-slider"])
                guidance_scale = gr.Slider(minimum=0.0, maximum=6.0, step=0.1, value=0.0, label="CFG 指数 (>=1 生效)", elem_classes=["compact-slider"])
                control_strength = gr.Slider(minimum=0.0, maximum=2.0, step=0.05, value=0.75, label="Pose 强度", elem_classes=["compact-slider"])
                height = gr.Slider(minimum=512, maximum=1792, step=16, value=1024, label="高度 (16 的倍数)", elem_classes=["compact-slider"])
                width = gr.Slider(minimum=512, maximum=1792, step=16, value=1024, label="宽度 (16 的倍数)", elem_classes=["compact-slider"])
                seed = gr.Number(value=-1, label="随机种子 (-1 表示随机)", precision=0)
                run_btn = gr.Button("生成", variant="primary")

            with gr.Column(scale=2.4):
                with gr.Row():
                    with gr.Column(scale=0.8, min_width=200):
                        pose_image = gr.Image(
                            label="姿态图上传 (RGB)",
                            type="pil",
                            height=320,
                            width=240,
                            show_download_button=True,
                        )
                    with gr.Column(scale=2.6):
                        result_img = gr.Image(
                            label="生成结果 (缩略图)",
                            type="pil",
                            height=520,
                            show_download_button=True,
                            show_fullscreen_button=True,
                        )
                        used_seed = gr.Number(label="实际种子", precision=0)

        run_btn.click(
            infer,
            inputs=[prompt, negative_prompt, pose_image, height, width, steps, guidance_scale, control_strength, seed],
            outputs=[result_img, used_seed],
        )

    return demo


def main():
    _load_pipeline()
    demo = build_ui()
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, inbrowser=False, share=False)


if __name__ == "__main__":
    main()