from fastapi import FastAPI, UploadFile, File, Form, HTTPException from fastapi.responses import JSONResponse, FileResponse, StreamingResponse from pydantic import BaseModel from typing import Optional, List, Dict, Any import torch import io from PIL import Image import base64 import json import time from datetime import datetime import asyncio from concurrent.futures import ThreadPoolExecutor import uuid import os from .sampler import TextToImagePipeline, SamplerFactory from ..data.text_encoder import CLIPTextEncoderWrapper # 请求/响应模型 class TextToImageRequest(BaseModel): prompt: str negative_prompt: Optional[str] = "" width: int = 512 height: int = 512 num_steps: int = 50 guidance_scale: float = 7.5 num_images: int = 1 seed: Optional[int] = None sampler: str = "ddim" class ImageResponse(BaseModel): images: List[str] # base64编码的图像 metadata: Dict[str, Any] request_id: str generation_time: float class BatchRequest(BaseModel): requests: List[TextToImageRequest] priority: int = 0 class StatusResponse(BaseModel): status: str queue_length: int active_tasks: int gpu_memory_usage: float uptime: float # API应用 class LuminaAPI: """Lumina API服务器""" def __init__( self, model, diffusion, text_encoder, vae_decoder=None, host: str = "0.0.0.0", port: int = 8000, max_queue_size: int = 100, max_workers: int = 2 ): self.model = model self.diffusion = diffusion self.text_encoder = text_encoder self.vae_decoder = vae_decoder self.host = host self.port = port # 创建FastAPI应用 self.app = FastAPI( title="Lumina Image Generation API", description="轻量级图像生成模型API", version="1.0.0" ) # 任务队列 self.task_queue = asyncio.Queue(maxsize=max_queue_size) self.executor = ThreadPoolExecutor(max_workers=max_workers) self.active_tasks = 0 # 请求历史 self.request_history = [] self.max_history = 1000 # 统计信息 self.start_time = time.time() self.total_requests = 0 self.total_images = 0 # 初始化管道 self.pipeline = TextToImagePipeline( model=model, diffusion=diffusion, text_encoder=text_encoder, vae_decoder=vae_decoder, sampler_type="ddim" ) # 设置路由 self._setup_routes() def _setup_routes(self): """设置API路由""" @self.app.get("/") async def root(): return { "message": "Lumina Image Generation API", "version": "1.0.0", "docs": "/docs", "endpoints": [ "/generate", "/batch_generate", "/status", "/health" ] } @self.app.get("/health") async def health_check(): """健康检查""" gpu_available = torch.cuda.is_available() gpu_memory = torch.cuda.memory_allocated() / 1024**3 if gpu_available else 0 return { "status": "healthy", "gpu_available": gpu_available, "gpu_memory_gb": gpu_memory, "model_loaded": self.model is not None, "text_encoder_loaded": self.text_encoder is not None } @self.app.get("/status") async def get_status(): """获取服务状态""" uptime = time.time() - self.start_time # GPU内存使用 if torch.cuda.is_available(): gpu_memory = torch.cuda.memory_allocated() / 1024**3 else: gpu_memory = 0 return StatusResponse( status="running", queue_length=self.task_queue.qsize(), active_tasks=self.active_tasks, gpu_memory_usage=gpu_memory, uptime=uptime ) @self.app.post("/generate", response_model=ImageResponse) async def generate_image(request: TextToImageRequest): """生成单个图像""" request_id = str(uuid.uuid4()) # 记录请求 self.request_history.append({ "request_id": request_id, "prompt": request.prompt, "timestamp": datetime.now().isoformat() }) # 限制历史记录大小 if len(self.request_history) > self.max_history: self.request_history = self.request_history[-self.max_history:] # 生成图像 start_time = time.time() try: # 在线程池中运行生成任务 loop = asyncio.get_event_loop() images = await loop.run_in_executor( self.executor, self._generate_sync, request ) generation_time = time.time() - start_time # 转换为base64 image_b64_list = [] for img_tensor in images: if isinstance(img_tensor, torch.Tensor): # 转换为PIL图像 if img_tensor.dim() == 4: img_tensor = img_tensor.squeeze(0) # 归一化到[0, 255] img_tensor = torch.clamp(img_tensor * 255, 0, 255).byte() # 转换为numpy数组 if img_tensor.shape[0] == 3: # CHW格式 img_array = img_tensor.permute(1, 2, 0).cpu().numpy() else: img_array = img_tensor.cpu().numpy() # 转换为PIL图像 img = Image.fromarray(img_array) # 转换为base64 buffered = io.BytesIO() img.save(buffered, format="PNG") img_b64 = base64.b64encode(buffered.getvalue()).decode() image_b64_list.append(img_b64) else: image_b64_list.append("") # 更新统计信息 self.total_requests += 1 self.total_images += len(images) return ImageResponse( images=image_b64_list, metadata={ "prompt": request.prompt, "negative_prompt": request.negative_prompt, "width": request.width, "height": request.height, "num_steps": request.num_steps, "guidance_scale": request.guidance_scale, "seed": request.seed, "sampler": request.sampler }, request_id=request_id, generation_time=generation_time ) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @self.app.post("/batch_generate") async def batch_generate(batch_request: BatchRequest): """批量生成图像""" request_ids = [str(uuid.uuid4()) for _ in batch_request.requests] results = [] # 为每个请求生成任务 tasks = [] for req, req_id in zip(batch_request.requests, request_ids): task = asyncio.create_task(self._process_single_request(req, req_id)) tasks.append(task) # 等待所有任务完成 results = await asyncio.gather(*tasks, return_exceptions=True) # 处理结果 successful_results = [] failed_results = [] for result in results: if isinstance(result, Exception): failed_results.append({"error": str(result)}) else: successful_results.append(result) return { "successful": successful_results, "failed": failed_results, "total_requests": len(batch_request.requests), "successful_count": len(successful_results), "failed_count": len(failed_results) } @self.app.get("/history") async def get_history(limit: int = 50): """获取请求历史""" return self.request_history[-limit:] @self.app.post("/txt2img") # Stable Diffusion兼容端点 async def txt2img( prompt: str = Form(...), negative_prompt: str = Form(""), width: int = Form(512), height: int = Form(512), steps: int = Form(50), cfg_scale: float = Form(7.5), seed: int = Form(-1) ): """兼容Stable Diffusion WebUI的端点""" request = TextToImageRequest( prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, num_steps=steps, guidance_scale=cfg_scale, seed=seed if seed != -1 else None ) response = await generate_image(request) # 返回第一个图像 if response.images: return StreamingResponse( io.BytesIO(base64.b64decode(response.images[0])), media_type="image/png" ) else: raise HTTPException(status_code=500, detail="生成失败") @self.app.get("/stats") async def get_stats(): """获取统计信息""" uptime = time.time() - self.start_time return { "total_requests": self.total_requests, "total_images": self.total_images, "requests_per_minute": self.total_requests / (uptime / 60) if uptime > 0 else 0, "avg_generation_time": None, # 可以添加计时逻辑 "uptime_seconds": uptime, "queue_size": self.task_queue.qsize(), "active_workers": self.active_tasks } def _generate_sync(self, request: TextToImageRequest) -> List[torch.Tensor]: """同步生成图像（在单独的线程中运行）""" # 更新管道采样器 self.pipeline.sampler = SamplerFactory.create_sampler( request.sampler, self.model, self.diffusion, request.num_steps ) # 生成图像 images = self.pipeline( prompt=request.prompt, negative_prompt=request.negative_prompt, height=request.height, width=request.width, num_inference_steps=request.num_steps, guidance_scale=request.guidance_scale, num_images=request.num_images, seed=request.seed, progress_bar=False ) return images async def _process_single_request(self, request: TextToImageRequest, request_id: str) -> Dict: """处理单个请求""" try: # 在队列中添加任务 await self.task_queue.put((request, request_id)) # 更新活动任务计数 self.active_tasks += 1 # 处理任务 loop = asyncio.get_event_loop() images = await loop.run_in_executor( self.executor, self._generate_sync, request ) # 转换图像 image_b64_list = [] for img_tensor in images: # 简化的转换逻辑 img_b64 = "placeholder" # 实际应该转换为base64 image_b64_list.append(img_b64) # 更新活动任务计数 self.active_tasks -= 1 return { "request_id": request_id, "images": image_b64_list, "success": True } except Exception as e: self.active_tasks -= 1 return { "request_id": request_id, "error": str(e), "success": False } def run(self): """运行API服务器""" import uvicorn print(f"启动Lumina API服务器在 http://{self.host}:{self.port}") print(f"API文档: http://{self.host}:{self.port}/docs") uvicorn.run( self.app, host=self.host, port=self.port, log_level="info" ) class SimpleWebUI: """简单的Web UI（使用Gradio）""" def __init__(self, pipeline: TextToImagePipeline): self.pipeline = pipeline def create_interface(self): """创建Gradio界面""" try: import gradio as gr except ImportError: print("Gradio未安装，无法创建Web UI") return None def generate_image_ui( prompt, negative_prompt, width, height, num_steps, guidance_scale, seed, sampler ): """UI生成函数""" # 设置种子 if seed and seed > 0: torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) # 生成图像 images = self.pipeline( prompt=prompt, negative_prompt=negative_prompt, height=height, width=width, num_inference_steps=num_steps, guidance_scale=guidance_scale, num_images=1, seed=seed if seed > 0 else None, progress_bar=False ) # 转换为PIL图像 if images: img_tensor = images[0] if isinstance(img_tensor, torch.Tensor): if img_tensor.dim() == 4: img_tensor = img_tensor.squeeze(0) # 归一化到[0, 255] img_tensor = torch.clamp(img_tensor * 255, 0, 255).byte() # 转换为numpy数组 if img_tensor.shape[0] == 3: # CHW格式 img_array = img_tensor.permute(1, 2, 0).cpu().numpy() else: img_array = img_tensor.cpu().numpy() # 转换为PIL图像 from PIL import Image img = Image.fromarray(img_array) return img return None # 创建界面 with gr.Blocks(title="Lumina Image Generator") as interface: gr.Markdown("# 🎨 Lumina - 轻量级图像生成") gr.Markdown("基于扩散模型的文本到图像生成系统") with gr.Row(): with gr.Column(): prompt = gr.Textbox( label="提示词", placeholder="输入描述图像的文本...", lines=3 ) negative_prompt = gr.Textbox( label="负面提示词", placeholder="不想在图像中出现的内容...", lines=2 ) with gr.Row(): width = gr.Slider( minimum=256, maximum=1024, value=512, step=64, label="宽度" ) height = gr.Slider( minimum=256, maximum=1024, value=512, step=64, label="高度" ) with gr.Row(): num_steps = gr.Slider( minimum=1, maximum=100, value=30, step=1, label="采样步数" ) guidance_scale = gr.Slider( minimum=1.0, maximum=20.0, value=7.5, step=0.5, label="引导强度" ) with gr.Row(): seed = gr.Number( value=-1, label="随机种子 (-1为随机)" ) sampler = gr.Dropdown( choices=["ddim", "dpm", "lcm"], value="ddim", label="采样器" ) generate_btn = gr.Button("生成图像", variant="primary") with gr.Column(): output_image = gr.Image( label="生成的图像", type="pil" ) # 示例 gr.Markdown("### 示例提示词") examples = gr.Examples( examples=[ ["A beautiful sunset over mountains, digital art", "", 512, 512, 30, 7.5, -1], ["A cute cat playing with a ball of yarn", "blurry, deformed", 512, 512, 25, 8.0, -1], ["An astronaut riding a horse on Mars", "cartoon, anime", 512, 512, 40, 7.0, -1] ], inputs=[prompt, negative_prompt, width, height, num_steps, guidance_scale, seed] ) # 事件处理 generate_btn.click( fn=generate_image_ui, inputs=[prompt, negative_prompt, width, height, num_steps, guidance_scale, seed, sampler], outputs=output_image ) return interface def launch(self, share: bool = False, server_name: str = "0.0.0.0", server_port: int = 7860): """启动Web UI""" interface = self.create_interface() if interface: interface.launch( share=share, server_name=server_name, server_port=server_port ) def create_api_server(config: dict, model, diffusion, text_encoder, vae_decoder=None): """创建API服务器""" # 确定主机和端口 host = config.get('host', '0.0.0.0') port = config.get('port', 8000) # 创建API服务器 api_server = LuminaAPI( model=model, diffusion=diffusion, text_encoder=text_encoder, vae_decoder=vae_decoder, host=host, port=port, max_queue_size=config.get('max_queue_size', 100), max_workers=config.get('max_workers', 2) ) return api_server def test_api(): """测试API""" import torch.nn as nn # 创建模拟组件 class MockModel(nn.Module): def forward(self, x, t, context): return torch.randn_like(x) class MockDiffusion: pass class MockTextEncoder: def encode(self, texts): return torch.randn(len(texts), 77, 768) model = MockModel() diffusion = MockDiffusion() text_encoder = MockTextEncoder() # 创建API服务器 api = LuminaAPI( model=model, diffusion=diffusion, text_encoder=text_encoder ) print("API服务器创建成功") print("端点:") print(" POST /generate - 生成图像") print(" GET /health - 健康检查") print(" GET /status - 状态信息") return api if __name__ == '__main__': # 测试API api = test_api() # 注意：实际运行需要调用 api.run()