LiuJuan Z-Image Generator代码实例：API化封装供内部系统调用的FastAPI示例

LiuJuan Z-Image Generator代码实例API化封装供内部系统调用的FastAPI示例1. 项目背景与需求如果你正在使用LiuJuan Z-Image Generator这个强大的本地图片生成工具可能会遇到这样一个场景团队里的设计师、运营同事或者公司内部的其他系统都希望能直接调用这个工具来生成图片而不是每个人都去学习怎么启动Streamlit界面、怎么配置参数。这时候把工具封装成一个API服务就成了一个很自然的需求。通过API任何有权限的系统或用户只需要发送一个简单的HTTP请求就能获得一张高质量的生成图片。这不仅能提升协作效率还能让图片生成能力无缝集成到现有的工作流中。本文将手把手带你将一个原本基于Streamlit交互的LiuJuan Z-Image Generator改造成一个稳定、易用的FastAPI服务。我们会从环境准备开始一步步讲解核心代码的封装、API接口的设计、错误处理最后完成部署和测试。即使你之前没有接触过FastAPI也能跟着做下来。2. 环境准备与项目结构在开始改造之前我们需要确保基础环境已经就绪并规划好新的项目结构。2.1 环境依赖假设你已经成功运行过原始的LiuJuan Z-Image Generator项目。我们需要在此基础上安装FastAPI及其相关依赖。创建一个新的requirements_api.txt文件或者直接在原有环境安装fastapi0.104.1 uvicorn[standard]0.24.0 python-multipart0.0.6 pydantic2.5.0使用pip安装它们pip install -r requirements_api.txt2.2 项目结构调整为了让代码更清晰我们建议对项目目录进行简单的调整。改造后的结构可能如下所示liujuan_zimage_api/ ├── app/ │ ├── __init__.py │ ├── main.py # FastAPI应用主入口 │ ├── core/ │ │ ├── __init__.py │ │ ├── generator.py # 封装图片生成核心逻辑 │ │ └── models.py # Pydantic数据模型定义 │ └── utils/ │ ├── __init__.py │ └── safety_checker.py # 可选的图片安全过滤器 ├── weights/ │ └── liujuan.safetensors # 你的自定义权重文件 ├── outputs/ # 图片输出目录 ├── requirements.txt # 原始项目依赖 ├── requirements_api.txt # API服务额外依赖 └── README.md核心思想是将业务逻辑图片生成、Web服务逻辑API接口和工具函数分离开这样代码更容易维护和扩展。3. 核心生成逻辑的封装首先我们需要把原始Streamlit应用中的图片生成代码抽离成一个独立的、可被调用的类或函数。这是API服务的“发动机”。在app/core/generator.py中我们创建ZImageGenerator类import torch from diffusers import DiffusionPipeline from typing import Optional, List import logging import os logging.basicConfig(levellogging.INFO) logger logging.getLogger(__name__) class ZImageGenerator: LiuJuan Z-Image 生成器核心类。 封装了模型加载、权重注入和图片生成的全流程。 def __init__(self, model_base: str ali-vilab/z-image, weight_path: str ./weights/liujuan.safetensors, device: Optional[str] None, torch_dtypetorch.bfloat16): 初始化生成器。 参数: model_base: 基础模型名称默认为阿里云通义Z-Image。 weight_path: LiuJuan自定义Safetensors权重文件路径。 device: 指定运行设备如cuda或cpu。为None时自动选择。 torch_dtype: 模型计算精度默认为BF16以优化显存和速度。 self.model_base model_base self.weight_path weight_path self.torch_dtype torch_dtype # 自动选择设备 if device is None: self.device cuda if torch.cuda.is_available() else cpu else: self.device device logger.info(f初始化生成器设备: {self.device}, 精度: {torch_dtype}) # 初始化时先不加载模型采用懒加载模式 self.pipe None def _load_model(self): 内部方法加载并配置扩散模型管道。 if self.pipe is not None: return logger.info(开始加载基础模型...) try: # 1. 加载基础模型管道 self.pipe DiffusionPipeline.from_pretrained( self.model_base, torch_dtypeself.torch_dtype, safety_checkerNone, # 可选禁用内置安全检查以提升速度 requires_safety_checkerFalse, ) # 2. 应用核心优化配置 # 显存碎片治理对于多轮生成尤其有效 if self.device cuda: torch.cuda.empty_cache() # 设置max_split_size_mb有助于减少显存碎片 # 注意此设置可能因CUDA版本而异实践中128是一个常用值 torch.cuda.set_per_process_memory_fraction(1.0) # 更通用的碎片治理方式是及时清理缓存 # 3. 加载自定义权重 logger.info(f正在注入自定义权重: {self.weight_path}) from safetensors.torch import load_file custom_weights load_file(self.weight_path) # 权重键名清洗移除可能的前缀以匹配基础模型结构 cleaned_weights {} for key, value in custom_weights.items(): # 移除常见的权重前缀如transformer.或model. new_key key.replace(transformer., ).replace(model., ) cleaned_weights[new_key] value # 宽松模式加载权重兼容非完全匹配的键 load_result self.pipe.unet.load_state_dict(cleaned_weights, strictFalse) if load_result.missing_keys: logger.warning(f权重加载缺失键: {load_result.missing_keys[:5]}...) # 只打印前5个 if load_result.unexpected_keys: logger.warning(f权重加载意外键: {load_result.unexpected_keys[:5]}...) # 4. 将整个管道移至目标设备并启用CPU卸载优化 self.pipe.to(self.device) if self.device cuda: # 启用模型CPU卸载将非活跃组件移至CPU显著降低峰值显存 self.pipe.enable_model_cpu_offload() logger.info(模型加载与优化完成。) except Exception as e: logger.error(f模型加载失败: {e}) self.pipe None raise def generate_image(self, prompt: str, negative_prompt: Optional[str] None, num_inference_steps: int 12, guidance_scale: float 2.0, height: int 768, width: int 768, num_images_per_prompt: int 1, seed: Optional[int] None) - List: 生成图片的核心方法。 参数: prompt: 正面提示词描述想要生成的图片内容。 negative_prompt: 负面提示词描述不希望出现的内容。 num_inference_steps: 去噪步数影响细节和生成时间。Z-Image推荐10-15。 guidance_scale: 提示词引导系数。Z-Image推荐较低值如2.0。 height: 生成图片高度。 width: 生成图片宽度。 num_images_per_prompt: 一次生成图片的数量。 seed: 随机种子用于复现结果。为None时随机生成。 返回: 一个PIL.Image对象的列表。 # 确保模型已加载懒加载 if self.pipe is None: self._load_model() # 设置随机种子以保证可复现性 if seed is not None: generator torch.Generator(deviceself.device).manual_seed(seed) else: generator None logger.info(f开始生成图片提示词: {prompt[:50]}...) try: # 调用管道生成图片 images self.pipe( promptprompt, negative_promptnegative_prompt, num_inference_stepsnum_inference_steps, guidance_scaleguidance_scale, heightheight, widthwidth, num_images_per_promptnum_images_per_prompt, generatorgenerator, ).images logger.info(f图片生成成功共{len(images)}张。) return images except torch.cuda.OutOfMemoryError as oom_error: logger.error(显存不足(OOM)尝试清理缓存并重试...) torch.cuda.empty_cache() # 这里可以添加降级策略例如降低分辨率或批次大小 raise RuntimeError(生成失败显存不足。请尝试降低图片分辨率或减少生成数量。) from oom_error except Exception as e: logger.error(f图片生成过程中发生错误: {e}) raise这个类做了几件关键事情懒加载模型只在第一次生成时加载避免API启动过慢。封装了所有优化包括BF16精度、权重清洗、CPU卸载等你无需再关心底层配置。清晰的错误处理对显存不足等常见问题做了捕获和友好提示。参数化所有生成参数都通过方法参数暴露方便API调用。4. 构建FastAPI应用与接口设计接下来我们使用FastAPI来构建Web服务。FastAPI以其高性能、易于使用和自动生成API文档的特点而闻名。在app/main.py中我们创建主应用from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi.responses import FileResponse, JSONResponse from fastapi.middleware.cors import CORSMiddleware import uuid import os from datetime import datetime from typing import List, Optional from app.core.generator import ZImageGenerator from app.core.models import GenerateRequest, GenerateResponse, TaskStatus from app.utils.safety_checker import content_safety_check # 可选的安全检查 # 初始化FastAPI应用 app FastAPI( titleLiuJuan Z-Image Generator API, description基于阿里云通义Z-Image和LiuJuan权重的图片生成API服务。, version1.0.0 ) # 添加CORS中间件允许前端跨域调用根据实际情况配置 app.add_middleware( CORSMiddleware, allow_origins[*], # 生产环境应替换为具体的前端地址 allow_credentialsTrue, allow_methods[*], allow_headers[*], ) # 全局生成器实例单例模式避免重复加载模型 _generator: Optional[ZImageGenerator] None # 用于存储任务状态的内存字典生产环境建议用Redis或数据库 _task_status_store {} def get_generator() - ZImageGenerator: 获取或创建全局生成器实例。 global _generator if _generator is None: _generator ZImageGenerator() return _generator app.on_event(startup) async def startup_event(): 应用启动时可选地预加载模型以加速第一个请求。 # 注释掉下一行则采用懒加载模式第一个请求会稍慢 # get_generator()._load_model() print(LiuJuan Z-Image Generator API 服务已启动。) app.get(/) async def root(): 根路径返回服务基本信息。 return { service: LiuJuan Z-Image Generator API, status: running, documentation: /docs, health_check: /health } app.get(/health) async def health_check(): 健康检查端点。 try: # 简单检查CUDA是否可用如果使用GPU import torch cuda_available torch.cuda.is_available() return { status: healthy, timestamp: datetime.now().isoformat(), cuda_available: cuda_available, model_loaded: _generator is not None and _generator.pipe is not None } except Exception as e: raise HTTPException(status_code500, detailfHealth check failed: {str(e)}) app.post(/generate, response_modelGenerateResponse) async def generate_image(request: GenerateRequest): 同步生成图片接口。 接收生成参数立即处理并返回图片。 适用于快速、轻量的生成请求。 generator get_generator() # 可选对提示词进行安全检查或过滤 # safe_prompt content_safety_check(request.prompt) try: # 调用核心生成逻辑 images generator.generate_image( promptrequest.prompt, negative_promptrequest.negative_prompt, num_inference_stepsrequest.num_inference_steps, guidance_scalerequest.guidance_scale, heightrequest.height, widthrequest.width, num_images_per_promptrequest.num_images_per_prompt, seedrequest.seed, ) # 保存图片到本地文件系统生产环境建议使用对象存储 saved_paths [] os.makedirs(outputs, exist_okTrue) for idx, img in enumerate(images): filename f{uuid.uuid4().hex[:8]}_{int(datetime.now().timestamp())}_{idx}.png filepath os.path.join(outputs, filename) img.save(filepath) saved_paths.append(filepath) print(f图片已保存: {filepath}) # 构建返回结果 # 注意这里直接返回文件路径。更常见的做法是上传到云存储后返回URL。 return GenerateResponse( successTrue, message图片生成成功, image_pathssaved_paths, task_idNone # 同步任务无task_id ) except Exception as e: raise HTTPException(status_code500, detailf图片生成失败: {str(e)}) app.post(/generate/async, response_modelGenerateResponse) async def generate_image_async(request: GenerateRequest, background_tasks: BackgroundTasks): 异步生成图片接口。 立即返回一个任务ID图片生成在后台进行。 适用于耗时较长或批量生成任务。 task_id uuid.uuid4().hex # 初始化任务状态 _task_status_store[task_id] TaskStatus( task_idtask_id, statuspending, message任务已接收等待处理, create_timedatetime.now() ) # 将生成任务添加到后台 background_tasks.add_task(process_generation_task, task_id, request) return GenerateResponse( successTrue, message异步任务已创建, image_paths[], task_idtask_id ) def process_generation_task(task_id: str, request: GenerateRequest): 后台任务处理函数。 try: _task_status_store[task_id].status processing _task_status_store[task_id].message 正在生成图片... generator get_generator() images generator.generate_image( promptrequest.prompt, negative_promptrequest.negative_prompt, num_inference_stepsrequest.num_inference_steps, guidance_scalerequest.guidance_scale, heightrequest.height, widthrequest.width, num_images_per_promptrequest.num_images_per_prompt, seedrequest.seed, ) saved_paths [] os.makedirs(outputs, exist_okTrue) for idx, img in enumerate(images): filename f{task_id}_{idx}.png filepath os.path.join(outputs, filename) img.save(filepath) saved_paths.append(filepath) _task_status_store[task_id].status completed _task_status_store[task_id].message 图片生成完成 _task_status_store[task_id].image_paths saved_paths _task_status_store[task_id].finish_time datetime.now() except Exception as e: _task_status_store[task_id].status failed _task_status_store[task_id].message f任务失败: {str(e)} _task_status_store[task_id].finish_time datetime.now() app.get(/task/{task_id}) async def get_task_status(task_id: str): 查询异步任务状态。 if task_id not in _task_status_store: raise HTTPException(status_code404, detail任务ID不存在) return _task_status_store[task_id] app.get(/image/{filename}) async def get_image(filename: str): 获取已生成图片的接口。 通过文件名从outputs目录返回图片文件。 filepath os.path.join(outputs, filename) if not os.path.exists(filepath): raise HTTPException(status_code404, detail图片未找到) return FileResponse(filepath, media_typeimage/png)同时我们需要在app/core/models.py中定义数据模型这有助于FastAPI自动进行请求验证和生成API文档from pydantic import BaseModel, Field from typing import Optional, List from datetime import datetime class GenerateRequest(BaseModel): 图片生成请求体模型。 prompt: str Field(..., description正面提示词描述想要生成的图片内容。, examplephotograph of a beautiful girl, close up, natural skin texture, soft lighting, 8k, masterpiece) negative_prompt: Optional[str] Field(, description负面提示词描述不希望出现的内容。, examplensfw, low quality, text, watermark, bad anatomy, blurry) num_inference_steps: int Field(12, ge1, le50, description去噪步数影响细节和生成时间。Z-Image推荐10-15。) guidance_scale: float Field(2.0, ge1.0, le20.0, description提示词引导系数。Z-Image推荐较低值如2.0。) height: int Field(768, ge512, le1024, description生成图片的高度。) width: int Field(768, ge512, le1024, description生成图片的宽度。) num_images_per_prompt: int Field(1, ge1, le4, description一次生成图片的数量。注意数量增加会显著增加显存消耗。) seed: Optional[int] Field(None, description随机种子用于复现结果。留空则随机生成。) class GenerateResponse(BaseModel): 图片生成响应模型。 success: bool message: str image_paths: List[str] [] # 生产环境建议改为图片URL列表 task_id: Optional[str] None # 异步任务ID class TaskStatus(BaseModel): 异步任务状态模型。 task_id: str status: str # pending, processing, completed, failed message: str image_paths: List[str] [] create_time: datetime finish_time: Optional[datetime] None5. 运行、测试与部署5.1 运行API服务在项目根目录下使用Uvicorn启动服务# 开发模式支持热重载 uvicorn app.main:app --reload --host 0.0.0.0 --port 8000 # 生产模式使用更多worker # uvicorn app.main:app --host 0.0.0.0 --port 8000 --workers 2启动后你可以通过浏览器访问http://localhost:8000/docs查看自动生成的交互式API文档Swagger UI并直接在那里测试接口。5.2 调用API示例使用Python调用同步接口import requests import json api_url http://localhost:8000/generate payload { prompt: photograph of a beautiful girl, close up, natural skin texture, soft lighting, 8k, masterpiece, negative_prompt: nsfw, low quality, text, watermark, num_inference_steps: 12, guidance_scale: 2.0, height: 768, width: 768, num_images_per_prompt: 1, seed: 42 } headers { Content-Type: application/json } response requests.post(api_url, datajson.dumps(payload), headersheaders) if response.status_code 200: result response.json() print(生成成功) print(f消息: {result[message]}) for path in result[image_paths]: print(f图片路径: {path}) # 你可以再调用 /image/{filename} 接口来下载图片 else: print(f请求失败: {response.status_code}) print(response.text)使用cURL调用异步接口curl -X POST http://localhost:8000/generate/async \ -H Content-Type: application/json \ -d { prompt: a cute cat wearing a hat, cartoon style, negative_prompt: ugly, deformed, num_inference_steps: 15 }如果成功你会得到一个包含task_id的响应。然后可以用这个ID查询状态curl http://localhost:8000/task/{你的task_id}5.3 生产环境部署建议使用Gunicorn/Uvicorn Worker对于生产环境建议使用Gunicorn管理多个Uvicorn worker进程以提高并发能力。反向代理使用Nginx或Apache作为反向代理处理SSL/TLS、静态文件和负载均衡。进程管理使用Systemd或Supervisor来管理服务进程确保异常退出后能自动重启。存储分离将生成的图片上传到云存储如AWS S3、阿里云OSS、MinIO并在响应中返回可访问的URL而不是服务器本地路径。任务队列对于高并发场景将异步任务放入Redis或RabbitMQ队列由独立的Worker进程消费避免API进程被阻塞。监控与日志集成Prometheus、Grafana进行监控并配置好日志收集如ELK栈。安全加固配置严格的CORS策略。添加API密钥认证可使用FastAPI的HTTPBearer。对用户输入的提示词进行内容安全过滤防止生成不当内容。6. 总结通过本文的步骤我们成功地将一个本地的LiuJuan Z-Image Generator工具封装成了一个功能完整的FastAPI服务。现在这个服务可以轻松地被其他系统集成内部管理系统可以增加一个“生成宣传图”的按钮直接调用此API。设计平台设计师可以在平台上输入描述后台调用API生成初稿。自动化脚本运营人员可以编写脚本批量生成不同风格的图片。这个改造过程的核心思想是解耦和封装将生成逻辑与交互界面分离并通过定义清晰的接口API来提供服务。这样做的好处非常明显可集成性任何支持HTTP调用的语言或平台都能使用。可维护性API服务器和生成逻辑可以独立升级、扩展。资源优化可以集中管理GPU资源避免每个用户都在本地运行沉重的模型。当然本文提供的示例是一个起点。在实际生产环境中你可能还需要考虑更多方面比如接口限流、缓存策略、更完善的错误处理以及高可用部署。但有了这个基础框架你可以根据实际需求灵活地进行扩展和优化。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。