initial open-source release
This commit is contained in:
minorQAQ committed
commit
01c4605498
224 files changed
+55124
No files matched your search
@@ -0,0 +1,207 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Z-Image-Turbo 图像生成服务 —— FastAPI 封装。
|
||||
|
||||
使用 Tongyi-MAI/Z-Image-Turbo 模型（基于 Diffusers），默认 1024×1024 输出。
|
||||
|
||||
启动方式:
|
||||
python image_gen_server.py --port 8904
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import base64
|
||||
import time
|
||||
import argparse
|
||||
import warnings
|
||||
import threading
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from io import BytesIO
|
||||
|
||||
import torch
|
||||
|
||||
# ======================== 配置区域 ========================
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
DEFAULT_PORT = 8904
|
||||
DEFAULT_MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
|
||||
|
||||
# 本地模型路径：优先读环境变量 LOCAL_MODEL_DIR，未设置则回退到远程 ID
|
||||
LOCAL_MODEL_DIR = os.environ.get("LOCAL_MODEL_DIR", "")
|
||||
MODEL_PATH = LOCAL_MODEL_DIR if LOCAL_MODEL_DIR and os.path.isdir(LOCAL_MODEL_DIR) else DEFAULT_MODEL_ID
|
||||
|
||||
# GPU 选择：CUDA_VISIBLE_DEVICES 或环境变量 SPECIFIC_GPU
|
||||
SPECIFIC_GPU = os.environ.get("SPECIFIC_GPU", "")
|
||||
if SPECIFIC_GPU:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = SPECIFIC_GPU
|
||||
print(f"[Z-Image] 已设置 CUDA_VISIBLE_DEVICES={SPECIFIC_GPU}")
|
||||
|
||||
# 兼容性：将前导 warnings 移到 torch 导入后
|
||||
warnings.filterwarnings("ignore", category=FutureWarning)
|
||||
|
||||
|
||||
# ======================== FastAPI ========================
|
||||
from fastapi import FastAPI, Request, HTTPException
|
||||
from fastapi.responses import JSONResponse, Response
|
||||
|
||||
# Lazy-loaded pipeline singleton
|
||||
_pipe = None
|
||||
_pipe_lock = threading.Lock()
|
||||
_device: str | None = None
|
||||
|
||||
|
||||
def _get_device() -> str:
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
return "cpu"
|
||||
|
||||
|
||||
def _get_pipe():
|
||||
"""延迟加载 Z-Image pipeline（线程安全）。
|
||||
使用 Diffusers DiffusionPipeline + trust_remote_code 加载模型仓库中的自定义 ZImagePipeline。
|
||||
"""
|
||||
global _pipe, _device
|
||||
if _pipe is not None:
|
||||
return _pipe, _device
|
||||
|
||||
with _pipe_lock:
|
||||
if _pipe is not None:
|
||||
return _pipe, _device
|
||||
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
_device = _get_device()
|
||||
load_kwargs = {
|
||||
"torch_dtype": torch.bfloat16 if _device == "cuda" else torch.float32,
|
||||
"low_cpu_mem_usage": False,
|
||||
"trust_remote_code": True,
|
||||
}
|
||||
# 本地路径时跳过在线检查
|
||||
if LOCAL_MODEL_DIR and os.path.isdir(LOCAL_MODEL_DIR):
|
||||
load_kwargs["local_files_only"] = True
|
||||
|
||||
print(f"[Z-Image] 加载模型: {MODEL_PATH}，设备: {_device}")
|
||||
_pipe = DiffusionPipeline.from_pretrained(MODEL_PATH, **load_kwargs)
|
||||
_pipe.to(_device)
|
||||
print(f"[Z-Image] 模型加载成功")
|
||||
return _pipe, _device
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
print(f"[Z-Image] 服务启动中，端口: {_model_args.get('port', DEFAULT_PORT)}")
|
||||
# 启动时预加载模型
|
||||
try:
|
||||
_get_pipe()
|
||||
except Exception as e:
|
||||
print(f"[Z-Image] 模型预加载失败（将在首次请求时重试）: {e}")
|
||||
yield
|
||||
print("[Z-Image] 服务关闭")
|
||||
|
||||
|
||||
app = FastAPI(title="Z-Image-Turbo Gen", lifespan=lifespan)
|
||||
|
||||
|
||||
@app.post("/generate")
|
||||
async def generate_image(request: Request):
|
||||
"""生成图像。
|
||||
|
||||
输入 JSON:
|
||||
prompt: str — 必填，正向提示词
|
||||
negative_prompt: str — 可选，负向提示词
|
||||
height: int — 可选，输出高度（默认 1024，须为 64 的倍数）
|
||||
width: int — 可选，输出宽度（默认 1024，须为 64 的倍数）
|
||||
num_inference_steps: int — 可选，推理步数（Turbo 模型建议 9，即 8 次 DiT forward）
|
||||
guidance_scale: float — 可选，CFG（Turbo 模型建议 0.0）
|
||||
seed: int — 可选，随机种子（-1 表示随机）
|
||||
return_format: str — 可选，"base64"（默认）或 "url"
|
||||
|
||||
输出 JSON:
|
||||
success: bool
|
||||
image: str — base64 编码的 PNG 或文件路径
|
||||
format: str — "png"
|
||||
width: int
|
||||
height: int
|
||||
"""
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
raise HTTPException(status_code=400, detail="请求体必须为 JSON")
|
||||
|
||||
prompt = (body.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
raise HTTPException(status_code=400, detail="prompt 不能为空")
|
||||
|
||||
height = int(body.get("height", 1024))
|
||||
width = int(body.get("width", 1024))
|
||||
steps = int(body.get("num_inference_steps", 11))
|
||||
guidance = float(body.get("guidance_scale", 0.0))
|
||||
seed = int(body.get("seed", -1))
|
||||
negative_prompt = (body.get("negative_prompt") or "").strip() or None
|
||||
return_format = (body.get("return_format") or "base64").strip().lower()
|
||||
|
||||
# 参数校验
|
||||
if height % 64 != 0 or width % 64 != 0:
|
||||
raise HTTPException(status_code=400, detail="height 和 width 须为 64 的整数倍")
|
||||
if height < 256 or width < 256 or height > 2048 or width > 2048:
|
||||
raise HTTPException(status_code=400, detail="尺寸须在 [256, 2048] 之间")
|
||||
|
||||
pipe, device = _get_pipe()
|
||||
|
||||
generator = None
|
||||
if seed >= 0:
|
||||
generator = torch.Generator(device).manual_seed(seed)
|
||||
|
||||
# 生成
|
||||
print(f"[Z-Image] 生成中: {width}x{height}, steps={steps}, prompt={prompt[:80]}...")
|
||||
t0 = time.time()
|
||||
result = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance,
|
||||
generator=generator,
|
||||
)
|
||||
image = result.images[0]
|
||||
elapsed = time.time() - t0
|
||||
print(f"[Z-Image] 生成完成，耗时 {elapsed:.1f}s")
|
||||
|
||||
# 编码输出
|
||||
buf = BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
buf.seek(0)
|
||||
|
||||
if return_format == "base64":
|
||||
img_b64 = base64.b64encode(buf.read()).decode("ascii")
|
||||
return JSONResponse({
|
||||
"success": True,
|
||||
"image": img_b64,
|
||||
"format": "png",
|
||||
"width": width,
|
||||
"height": height,
|
||||
"elapsed_seconds": round(elapsed, 2),
|
||||
})
|
||||
else:
|
||||
return Response(content=buf.getvalue(), media_type="image/png")
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# ======================== 启动入口 ========================
|
||||
_model_args: dict = {}
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Z-Image-Turbo Image Generation Server")
|
||||
parser.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"服务端口（默认 {DEFAULT_PORT}）")
|
||||
args = parser.parse_args()
|
||||
|
||||
_model_args["port"] = args.port
|
||||
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=args.port)
|
||||
Reference in new issue
Block a user