src/inference/api.py

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()