#新建man.py文件，加入代码
from modelscope import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, snapshot_download import torch import math import os from pathlib import Path from datetime import datetime import base64 from io import BytesIO from typing import Optional, List, Union import time import uuid from fastapi import FastAPI, HTTPException from pydantic import BaseModel import uvicorn# 设置HF_ENDPOINT使用国内镜像加速模型下载 os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"# 减少显存碎片化 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"# 启用推理模式，减少显存占用 torch.inference_mode()# 定义请求和响应模型 class ImageGenerationRequest(BaseModel):prompt: strnegative_prompt: Optional[str] = ""width: Optional[int] = 1024height: Optional[int] = 1024num_inference_steps: Optional[int] = 8seed: Optional[int] = 0class ImageData(BaseModel):b64_json: strrevised_prompt: strclass ImageGenerationResponse(BaseModel):created: intdata: List[ImageData]# 初始化模型 def initialize_model():print("=" * 50)print("开始初始化模型...")print("=" * 50)scheduler_config = {'base_image_seq_len': 256,'base_shift': math.log(3),'invert_sigmas': False,'max_image_seq_len': 8192,'max_shift': math.log(3),'num_train_timesteps': 1000,'shift': 1.0,'shift_terminal': None,'stochastic_sampling': False,'time_shift_type': 'exponential','use_beta_sigmas': False,'use_dynamic_shifting': True,'use_exponential_sigmas': False,'use_karras_sigmas': False,}print("正在配置调度器...")scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)print("调度器配置完成")# 检测可用GPU数量num_gpus = torch.cuda.device_count()print(f"检测到 {num_gpus} 个GPU设备")# 根据项目规范，对于DiffusionPipeline模型，我们使用max_memory参数进行显存分配# 而不是手动指定每层的设备映射if num_gpus > 1:print("检测到多个GPU，正在进行显存分配...")# 获取每个GPU的显存信息并计算可分配的显存量max_memory = {}for i in range(num_gpus):free_mem, total_mem = torch.cuda.mem_get_info(i)# 按70%的空闲显存计算分配量，同时确保不超过22GBallocated_mem = min(int(free_mem * 0.7), 22 * 1024**3, free_mem)max_memory[i] = allocated_memprint(f"GPU {i}: 分配 {(allocated_mem / 1024**3):.2f} GB 显存")# 加载模型并指定显存分配print("正在加载模型到多个GPU...")pipe = DiffusionPipeline.from_pretrained('Qwen/Qwen-Image',scheduler=scheduler,torch_dtype=torch.bfloat16,max_memory=max_memory, # 为每个GPU分配显存 )print("多GPU模型加载完成")else:# 单GPU情况print("检测到单个GPU，正在加载模型...")pipe = DiffusionPipeline.from_pretrained('Qwen/Qwen-Image',scheduler=scheduler,torch_dtype=torch.bfloat16,)pipe = pipe.to("cuda")print("单GPU模型加载完成")print(f"模型已分配到{num_gpus}个GPU设备上")print("=" * 50)print("模型初始化完成")print("=" * 50)return pipe# 提前下载LoRA权重 def download_lora_weights():print("=" * 50)print("开始下载LoRA权重...")print("=" * 50)# 使用ModelScope的snapshot_download下载LoRA权重model_dir = snapshot_download('lightx2v/Qwen-Image-Lightning')print(f"LoRA权重已下载到: {model_dir}")# 查找.pt或.safetensors文件lora_files = list(Path(model_dir).glob("*.safetensors")) + list(Path(model_dir).glob("*.pt"))if not lora_files:raise FileNotFoundError("在下载的LoRA权重目录中未找到.safetensors或.pt文件")lora_file_path = lora_files[0] # 使用找到的第一个文件print(f"使用LoRA文件: {lora_file_path}")print("=" * 50)print("LoRA权重下载完成")print("=" * 50)return str(lora_file_path)# 图像生成函数 def generate_image(request: ImageGenerationRequest, pipe):"""生成图像"""print("=" * 50)print("开始生成图像...")print("=" * 50)print(f"输入参数:")print(f" - 提示词: {request.prompt}")print(f" - 负面提示词: {request.negative_prompt}")print(f" - 图像尺寸: {request.width}x{request.height}")print(f" - 推理步数: {request.num_inference_steps}")print(f" - 随机种子: {request.seed}")start_time = time.time()print("开始图像生成过程...")print("正在调用模型生成图像...")image = pipe(prompt=request.prompt,negative_prompt=request.negative_prompt,width=request.width,height=request.height,num_inference_steps=request.num_inference_steps,true_cfg_scale=1.0,generator=torch.manual_seed(request.seed)).images[0]generation_time = time.time() - start_timeprint(f"图像生成完成，耗时: {generation_time:.2f} 秒")print("正在将图像转换为base64编码...")# 将图像转换为base64编码buffered = BytesIO()image.save(buffered, format="PNG")img_str = base64.b64encode(buffered.getvalue()).decode()print("图像转换完成")print("=" * 50)print("图像生成流程结束")print("=" * 50)return img_str# 初始化FastAPI应用 app = FastAPI(title="Qwen-Image API",description="基于Qwen-Image模型的OpenAI兼容图像生成API",version="1.0.0" )# 全局模型实例 pipe = None@app.on_event("startup") async def startup_event():global pipe# 初始化模型pipe = initialize_model()# 下载LoRA权重lora_file_path = download_lora_weights()print("正在加载LoRA权重...")pipe.load_lora_weights(lora_file_path)print("LoRA权重加载完成")@app.post("/v1/images/generations", response_model=ImageGenerationResponse) async def create_image(request: ImageGenerationRequest):try:# 生成图像image_data = generate_image(request, pipe)# 构造响应response = ImageGenerationResponse(created=int(time.time()),data=[ImageData(b64_json=image_data,revised_prompt=request.prompt)])return responseexcept Exception as e:raise HTTPException(status_code=500, detail=str(e))if __name__ == "__main__":uvicorn.run(app, host="0.0.0.0", port=8800)
测试调用
curl -X POST http://localhost:8800/v1/images/generations -H "Content-Type: application/json" -d '{"prompt": "a lovely cat"}' | grep -q '"data"' && echo "API 工作正常" || echo "API 出现问题"
调用时控制台截图：
