| 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]
|
| 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
|
|
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| image_b64_list = []
|
| for img_tensor in images:
|
| if isinstance(img_tensor, torch.Tensor):
|
|
|
| if img_tensor.dim() == 4:
|
| img_tensor = img_tensor.squeeze(0)
|
|
|
|
|
| img_tensor = torch.clamp(img_tensor * 255, 0, 255).byte()
|
|
|
|
|
| if img_tensor.shape[0] == 3:
|
| img_array = img_tensor.permute(1, 2, 0).cpu().numpy()
|
| else:
|
| img_array = img_tensor.cpu().numpy()
|
|
|
|
|
| img = Image.fromarray(img_array)
|
|
|
|
|
| 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")
|
| 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"
|
| 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
|
| )
|
|
|
|
|
| if images:
|
| img_tensor = images[0]
|
| if isinstance(img_tensor, torch.Tensor):
|
| if img_tensor.dim() == 4:
|
| img_tensor = img_tensor.squeeze(0)
|
|
|
|
|
| img_tensor = torch.clamp(img_tensor * 255, 0, 255).byte()
|
|
|
|
|
| if img_tensor.shape[0] == 3:
|
| img_array = img_tensor.permute(1, 2, 0).cpu().numpy()
|
| else:
|
| img_array = img_tensor.cpu().numpy()
|
|
|
|
|
| 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_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 = 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 = test_api()
|
|
|
| |
