# /root/code/python/ieops-python/sdk/src/handler/server/diffuser.py
"""
Diffuser Server - Hugging Face Diffusers 推理服务

直接集成 Diffusers 库，无需子进程。提供：
- 文生图 (txt2img)
- 图生图 (img2img)
- 图像修复 (inpaint)
- 流式进度回调
- LoRA 加载

架构（直接集成，无子进程）：
┌─────────────────┐     ┌────────────────────────────────┐
│  外部客户端      │────▶│  DiffuserServer (FastAPI)      │
│  (HTTP/SSE)     │     │  内置 Diffusers Pipeline       │
└─────────────────┘     └────────────────────────────────┘
"""

import asyncio
import json
import os
import base64
import gc
import time
import uuid
from io import BytesIO
from typing import Any, AsyncGenerator, Dict, List, Optional, Callable
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field

from .server import BaseServer
from ...utils.log import uvicorn_logger, log
from ...utils.util import APP_ID, arandstr

# ========== 环境变量配置 ==========
_MODEL_SERVER_SOCKET = os.getenv("MODEL_SERVER_SOCKET", f"{APP_ID}.sock")
_MODEL_SERVER_PORT = int(os.getenv("MODEL_SERVER_PORT", 8001))
_MODEL_SERVER_HOST = os.getenv("MODEL_SERVER_HOST", "127.0.0.1")
_GRACEFUL_SHUTDOWN_TIME = int(os.getenv("GRACEFUL_SHUTDOWN_TIME", 3))
_MODEL_TIMEOUT = int(os.getenv("MODEL_TIMEOUT", 600))

# Diffuser 配置
_MODEL_PATH = os.getenv("MODEL_PATH", "stabilityai/stable-diffusion-xl-base-1.0")
_MODEL_TYPE = os.getenv("MODEL_TYPE", "auto")
_TORCH_DTYPE = os.getenv("TORCH_DTYPE", "float16")
_DEVICE = os.getenv("DEVICE", "auto")
_ENABLE_VAE_SLICING = os.getenv("ENABLE_VAE_SLICING", "true").lower() == "true"
_ENABLE_VAE_TILING = os.getenv("ENABLE_VAE_TILING", "false").lower() == "true"
_ENABLE_MODEL_CPU_OFFLOAD = os.getenv("ENABLE_MODEL_CPU_OFFLOAD", "false").lower() == "true"
_ENABLE_XFORMERS = os.getenv("ENABLE_XFORMERS", "false").lower() == "true"
_COMPILE_MODEL = os.getenv("COMPILE_MODEL", "false").lower() == "true"

SERVER_CODE_OK = 0
SERVER_CODE_ERROR = 1
SERVER_CODE_STOP = 2


# ========== 延迟导入 Diffusers（避免未安装时报错）==========
def _import_diffusers():
    """延迟导入 diffusers 相关模块"""
    import torch
    from PIL import Image
    from diffusers import (
        AutoPipelineForText2Image,
        AutoPipelineForImage2Image,
        AutoPipelineForInpainting,
        DPMSolverMultistepScheduler,
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
        DDIMScheduler,
        UniPCMultistepScheduler,
        LCMScheduler,
    )
    return {
        "torch": torch,
        "Image": Image,
        "AutoPipelineForText2Image": AutoPipelineForText2Image,
        "AutoPipelineForImage2Image": AutoPipelineForImage2Image,
        "AutoPipelineForInpainting": AutoPipelineForInpainting,
        "schedulers": {
            "dpm++_2m": DPMSolverMultistepScheduler,
            "dpm++_2m_karras": DPMSolverMultistepScheduler,
            "euler": EulerDiscreteScheduler,
            "euler_a": EulerAncestralDiscreteScheduler,
            "ddim": DDIMScheduler,
            "unipc": UniPCMultistepScheduler,
            "lcm": LCMScheduler,
        }
    }


# ========== Pydantic 模型定义 ==========
class Txt2ImgRequest(BaseModel):
    """文生图请求"""
    prompt: str = Field(..., description="正向提示词")
    negative_prompt: str = Field(default="", description="负向提示词")
    width: int = Field(default=1024, ge=256, le=4096, description="图像宽度")
    height: int = Field(default=1024, ge=256, le=4096, description="图像高度")
    num_inference_steps: int = Field(default=30, ge=1, le=150, description="推理步数")
    guidance_scale: float = Field(default=7.5, ge=0, le=30, description="引导系数")
    num_images: int = Field(default=1, ge=1, le=4, description="生成图片数量")
    seed: Optional[int] = Field(default=None, description="随机种子")
    scheduler: Optional[str] = Field(default=None, description="调度器类型")
    lora_weights: Optional[List[Dict[str, Any]]] = Field(default=None, description="LoRA 权重配置")


class Img2ImgRequest(BaseModel):
    """图生图请求"""
    prompt: str = Field(..., description="正向提示词")
    negative_prompt: str = Field(default="", description="负向提示词")
    image: str = Field(..., description="输入图像 (base64)")
    strength: float = Field(default=0.75, ge=0, le=1, description="去噪强度")
    num_inference_steps: int = Field(default=30, ge=1, le=150, description="推理步数")
    guidance_scale: float = Field(default=7.5, ge=0, le=30, description="引导系数")
    num_images: int = Field(default=1, ge=1, le=4, description="生成图片数量")
    seed: Optional[int] = Field(default=None, description="随机种子")


class InpaintRequest(BaseModel):
    """图像修复请求"""
    prompt: str = Field(..., description="正向提示词")
    negative_prompt: str = Field(default="", description="负向提示词")
    image: str = Field(..., description="输入图像 (base64)")
    mask_image: str = Field(..., description="遮罩图像 (base64)")
    width: Optional[int] = Field(default=None, description="输出宽度")
    height: Optional[int] = Field(default=None, description="输出高度")
    num_inference_steps: int = Field(default=30, ge=1, le=150, description="推理步数")
    guidance_scale: float = Field(default=7.5, ge=0, le=30, description="引导系数")
    strength: float = Field(default=0.75, ge=0, le=1, description="去噪强度")
    num_images: int = Field(default=1, ge=1, le=4, description="生成图片数量")
    seed: Optional[int] = Field(default=None, description="随机种子")


class DiffuserServer(BaseServer):
    """
    Diffuser 推理服务器 - 直接集成 Diffusers 库
    
    无需子进程，直接在当前进程中加载和运行 Diffusers Pipeline。
    """
    
    def __init__(
        self,
        model_path: Optional[str] = None,
        model_type: Optional[str] = None,
        torch_dtype: Optional[str] = None,
        device: Optional[str] = None,
        on_request_received: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
        enable_vae_slicing: Optional[bool] = None,
        enable_vae_tiling: Optional[bool] = None,
        enable_model_cpu_offload: Optional[bool] = None,
        enable_xformers: Optional[bool] = None,
        compile_model: Optional[bool] = None,
    ) -> None:
        """
        Args:
            model_path: 模型路径或 HuggingFace 模型 ID
            model_type: 模型类型 (auto, sd15, sdxl, sd3, flux)
            torch_dtype: 数据类型 (float16, float32, bfloat16)
            device: 设备 (cuda, cpu, auto)
            on_request_received: 请求预处理回调
            enable_vae_slicing: 启用 VAE slicing（减少显存）
            enable_vae_tiling: 启用 VAE tiling（支持大图）
            enable_model_cpu_offload: 启用模型 CPU 卸载
            enable_xformers: 启用 xformers 加速
            compile_model: 启用 PyTorch 2.0 compile
        """
        self._model_path = model_path or _MODEL_PATH
        self._model_type = model_type or _MODEL_TYPE
        self._torch_dtype = torch_dtype or _TORCH_DTYPE
        self._device = device or _DEVICE
        self._on_request_received = on_request_received
        
        # 优化选项
        self._enable_vae_slicing = enable_vae_slicing if enable_vae_slicing is not None else _ENABLE_VAE_SLICING
        self._enable_vae_tiling = enable_vae_tiling if enable_vae_tiling is not None else _ENABLE_VAE_TILING
        self._enable_model_cpu_offload = enable_model_cpu_offload if enable_model_cpu_offload is not None else _ENABLE_MODEL_CPU_OFFLOAD
        self._enable_xformers = enable_xformers if enable_xformers is not None else _ENABLE_XFORMERS
        self._compile_model = compile_model if compile_model is not None else _COMPILE_MODEL
        
        # Pipeline（启动时加载）
        self._modules: Optional[Dict[str, Any]] = None
        self._txt2img_pipe = None
        self._img2img_pipe = None
        self._inpaint_pipe = None
        self._loaded_loras: List[str] = []
        self._tasks: Dict[str, Dict[str, Any]] = {}
        
        self._app = FastAPI(lifespan=self._lifespan)
        self._server: Optional[uvicorn.Server] = None
        
        self._setup_routes()

    @asynccontextmanager
    async def _lifespan(self, app: FastAPI):
        """FastAPI 生命周期管理"""
        uvicorn_logger.info("Diffuser Server starting...")
        
        # 加载模型
        await self._load_model()
        
        yield
        
        # 清理
        uvicorn_logger.info("Diffuser Server shutting down...")
        await self._unload_model()
        
        # 清理 socket 文件
        if _MODEL_SERVER_SOCKET and os.path.exists(_MODEL_SERVER_SOCKET):
            os.remove(_MODEL_SERVER_SOCKET)

    async def _load_model(self):
        """加载 Diffusers 模型"""
        uvicorn_logger.info(f"Loading model: {self._model_path}")
        start_time = time.time()
        
        try:
            # 延迟导入
            self._modules = _import_diffusers()
            torch = self._modules["torch"]
            
            # 确定设备
            if self._device == "auto":
                self._device = "cuda" if torch.cuda.is_available() else "cpu"
            
            # 确定数据类型
            dtype_map = {
                "float16": torch.float16,
                "float32": torch.float32,
                "bfloat16": torch.bfloat16,
            }
            torch_dtype = dtype_map.get(self._torch_dtype, torch.float16)
            
            # 加载 txt2img pipeline
            AutoPipelineForText2Image = self._modules["AutoPipelineForText2Image"]
            self._txt2img_pipe = AutoPipelineForText2Image.from_pretrained(
                self._model_path,
                torch_dtype=torch_dtype,
                use_safetensors=True,
                variant="fp16" if torch_dtype == torch.float16 else None,
            )
            
            # 应用优化
            if not self._enable_model_cpu_offload:
                self._txt2img_pipe = self._txt2img_pipe.to(self._device)
            
            if self._enable_vae_slicing and hasattr(self._txt2img_pipe, 'enable_vae_slicing'):
                self._txt2img_pipe.enable_vae_slicing()
                
            if self._enable_vae_tiling and hasattr(self._txt2img_pipe, 'enable_vae_tiling'):
                self._txt2img_pipe.enable_vae_tiling()
                
            if self._enable_model_cpu_offload and hasattr(self._txt2img_pipe, 'enable_model_cpu_offload'):
                self._txt2img_pipe.enable_model_cpu_offload()
                
            if self._enable_xformers:
                try:
                    self._txt2img_pipe.enable_xformers_memory_efficient_attention()
                except Exception as e:
                    uvicorn_logger.warning(f"Could not enable xformers: {e}")
            
            if self._compile_model and hasattr(torch, "compile"):
                if hasattr(self._txt2img_pipe, 'unet'):
                    self._txt2img_pipe.unet = torch.compile(self._txt2img_pipe.unet, mode="reduce-overhead")
            
            # 创建 img2img 和 inpaint pipeline
            try:
                AutoPipelineForImage2Image = self._modules["AutoPipelineForImage2Image"]
                self._img2img_pipe = AutoPipelineForImage2Image.from_pipe(self._txt2img_pipe)
            except Exception as e:
                uvicorn_logger.warning(f"Could not create img2img pipeline: {e}")
            
            try:
                AutoPipelineForInpainting = self._modules["AutoPipelineForInpainting"]
                self._inpaint_pipe = AutoPipelineForInpainting.from_pipe(self._txt2img_pipe)
            except Exception as e:
                uvicorn_logger.warning(f"Could not create inpaint pipeline: {e}")
            
            load_time = time.time() - start_time
            uvicorn_logger.info(f"Model loaded in {load_time:.2f}s")
            
        except Exception as e:
            uvicorn_logger.error(f"Failed to load model: {e}")
            raise

    async def _unload_model(self):
        """卸载模型，释放显存"""
        self._txt2img_pipe = None
        self._img2img_pipe = None
        self._inpaint_pipe = None
        
        gc.collect()
        
        if self._modules:
            torch = self._modules["torch"]
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def shutdown(self) -> None:
        """请求服务器优雅关闭"""
        if self._server:
            self._server.should_exit = True

    async def serve(self) -> None:
        """启动服务器"""
        if _MODEL_SERVER_SOCKET:
            config = uvicorn.Config(
                self._app,
                uds=_MODEL_SERVER_SOCKET,
                log_level="info",
                loop="asyncio",
                timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN_TIME
            )
        else:
            config = uvicorn.Config(
                self._app,
                host=_MODEL_SERVER_HOST,
                port=_MODEL_SERVER_PORT,
                log_level="info",
                loop="asyncio",
                timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN_TIME
            )
        
        self._server = uvicorn.Server(config)
        await self._server.serve()

    def _setup_routes(self) -> None:
        """设置 API 路由"""
        
        @self._app.get("/health")
        async def health():
            """健康检查"""
            if self._txt2img_pipe is None:
                return JSONResponse({"status": "initializing"}, status_code=503)
            return {"status": "ok"}

        @self._app.get("/")
        async def root():
            """根路径"""
            return {
                "service": "Diffuser Server",
                "model": self._model_path,
                "status": "ready" if self._txt2img_pipe else "initializing",
                "endpoints": {
                    "health": "/health",
                    "txt2img": "/v1/txt2img",
                    "img2img": "/v1/img2img",
                    "inpaint": "/v1/inpaint",
                    "model_info": "/v1/model/info",
                }
            }

        @self._app.get("/v1/model/info")
        async def get_model_info():
            """获取模型信息"""
            torch = self._modules["torch"] if self._modules else None
            info = {
                "model_path": self._model_path,
                "model_type": self._model_type,
                "device": self._device,
                "torch_dtype": self._torch_dtype,
                "loaded": self._txt2img_pipe is not None,
                "loaded_loras": self._loaded_loras,
                "features": {
                    "txt2img": self._txt2img_pipe is not None,
                    "img2img": self._img2img_pipe is not None,
                    "inpaint": self._inpaint_pipe is not None,
                },
                "optimizations": {
                    "vae_slicing": self._enable_vae_slicing,
                    "vae_tiling": self._enable_vae_tiling,
                    "model_cpu_offload": self._enable_model_cpu_offload,
                    "xformers": self._enable_xformers,
                    "compile": self._compile_model,
                }
            }
            if torch and torch.cuda.is_available():
                info["cuda_memory_allocated"] = f"{torch.cuda.memory_allocated() / 1024**3:.2f} GB"
                info["cuda_memory_reserved"] = f"{torch.cuda.memory_reserved() / 1024**3:.2f} GB"
            return info

        @self._app.get("/v1/schedulers")
        async def get_schedulers():
            """获取可用调度器"""
            schedulers = list(self._modules["schedulers"].keys()) if self._modules else []
            return {"schedulers": schedulers}

        # ========== 文生图 ==========
        @self._app.post("/v1/txt2img")
        async def txt2img(request: Txt2ImgRequest):
            """文生图（同步）"""
            result = None
            async for data in self._txt2img(request, stream=False):
                if data.get("type") == "result":
                    result = data
                elif data.get("type") == "error":
                    raise HTTPException(status_code=500, detail=data["error"])
            return result

        @self._app.post("/v1/txt2img/stream")
        async def txt2img_stream(request: Txt2ImgRequest):
            """文生图（流式）"""
            return StreamingResponse(
                self._txt2img_sse(request),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        # ========== 图生图 ==========
        @self._app.post("/v1/img2img")
        async def img2img(request: Img2ImgRequest):
            """图生图（同步）"""
            result = None
            async for data in self._img2img(request, stream=False):
                if data.get("type") == "result":
                    result = data
                elif data.get("type") == "error":
                    raise HTTPException(status_code=500, detail=data["error"])
            return result

        @self._app.post("/v1/img2img/stream")
        async def img2img_stream(request: Img2ImgRequest):
            """图生图（流式）"""
            return StreamingResponse(
                self._img2img_sse(request),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        # ========== 图像修复 ==========
        @self._app.post("/v1/inpaint")
        async def inpaint(request: InpaintRequest):
            """图像修复（同步）"""
            result = None
            async for data in self._inpaint(request, stream=False):
                if data.get("type") == "result":
                    result = data
                elif data.get("type") == "error":
                    raise HTTPException(status_code=500, detail=data["error"])
            return result

        @self._app.post("/v1/inpaint/stream")
        async def inpaint_stream(request: InpaintRequest):
            """图像修复（流式）"""
            return StreamingResponse(
                self._inpaint_sse(request),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        # ========== 兼容 ieops 协议 ==========
        @self._app.post("/sse/infer")
        async def infer_sse(request: Request):
            """SSE 推理接口"""
            body = await request.body()
            return StreamingResponse(
                self._ieops_sse_generator(body),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no",
                }
            )

        @self._app.post("/infer")
        async def infer(request: Request):
            """同步推理接口"""
            body = await request.body()
            results: List[Any] = []
            last_code = SERVER_CODE_OK
            
            async for item in self._ieops_generator(body):
                code = item["code"]
                data = item["data"]
                if code == SERVER_CODE_ERROR:
                    last_code = code
                    results.append({"error": data})
                    break
                if code == SERVER_CODE_STOP:
                    break
                try:
                    results.append(json.loads(data) if isinstance(data, str) else data)
                except json.JSONDecodeError:
                    results.append(data)
            
            return {"code": last_code, "results": results}

    # ========== 工具方法 ==========
    def _decode_base64_image(self, base64_str: str):
        """解码 base64 图像"""
        Image = self._modules["Image"]
        if base64_str.startswith("data:"):
            base64_str = base64_str.split(",", 1)[1]
        image_data = base64.b64decode(base64_str)
        return Image.open(BytesIO(image_data)).convert("RGB")

    def _encode_image_to_base64(self, image, format: str = "PNG") -> str:
        """编码图像为 base64"""
        buffer = BytesIO()
        image.save(buffer, format=format)
        return base64.b64encode(buffer.getvalue()).decode("utf-8")

    def _set_scheduler(self, pipe, scheduler_name: Optional[str]):
        """设置调度器"""
        if not scheduler_name or not self._modules:
            return
        
        schedulers = self._modules["schedulers"]
        if scheduler_name not in schedulers:
            return
        
        scheduler_class = schedulers[scheduler_name]
        config = pipe.scheduler.config
        
        if "karras" in scheduler_name:
            pipe.scheduler = scheduler_class.from_config(config, use_karras_sigmas=True)
        else:
            pipe.scheduler = scheduler_class.from_config(config)

    def _load_lora(self, pipe, lora_weights: List[Dict[str, Any]]):
        """加载 LoRA"""
        for lora in lora_weights:
            lora_path = lora.get("path")
            lora_scale = lora.get("scale", 1.0)
            adapter_name = lora.get("name", os.path.basename(lora_path) if lora_path else "default")
            
            if lora_path and lora_path not in self._loaded_loras:
                try:
                    pipe.load_lora_weights(lora_path, adapter_name=adapter_name)
                    pipe.set_adapters([adapter_name], adapter_weights=[lora_scale])
                    self._loaded_loras.append(lora_path)
                    uvicorn_logger.info(f"Loaded LoRA: {lora_path}")
                except Exception as e:
                    uvicorn_logger.error(f"Failed to load LoRA: {e}")

    def _create_progress_callback(self, task_id: str, total_steps: int):
        """创建进度回调"""
        def callback(pipe, step_index, timestep, callback_kwargs):
            progress = (step_index + 1) / total_steps
            if task_id in self._tasks:
                self._tasks[task_id]["progress"] = progress
                self._tasks[task_id]["step"] = step_index + 1
            return callback_kwargs
        return callback

    # ========== 推理方法 ==========
    async def _txt2img(self, request: Txt2ImgRequest, stream: bool = False) -> AsyncGenerator[Dict[str, Any], None]:
        """文生图推理"""
        if not self._txt2img_pipe:
            yield {"type": "error", "error": "Model not loaded"}
            return
        
        torch = self._modules["torch"]
        task_id = str(uuid.uuid4())
        self._tasks[task_id] = {"progress": 0, "step": 0}
        
        try:
            # 种子
            seed = request.seed if request.seed is not None else int(torch.randint(0, 2**32 - 1, (1,)).item())
            generator = torch.Generator(device=self._device).manual_seed(seed)
            
            # 调度器
            if request.scheduler:
                self._set_scheduler(self._txt2img_pipe, request.scheduler)
            
            # LoRA
            if request.lora_weights:
                self._load_lora(self._txt2img_pipe, request.lora_weights)
            
            # 进度回调
            callback = self._create_progress_callback(task_id, request.num_inference_steps) if stream else None
            
            # 推理（在线程池中运行避免阻塞）
            loop = asyncio.get_event_loop()
            
            def run_inference():
                return self._txt2img_pipe(
                    prompt=request.prompt,
                    negative_prompt=request.negative_prompt or None,
                    width=request.width,
                    height=request.height,
                    num_inference_steps=request.num_inference_steps,
                    guidance_scale=request.guidance_scale,
                    num_images_per_prompt=request.num_images,
                    generator=generator,
                    callback_on_step_end=callback,
                )
            
            if stream:
                inference_task = loop.run_in_executor(None, run_inference)
                last_step = 0
                while not inference_task.done():
                    await asyncio.sleep(0.1)
                    current_step = self._tasks[task_id].get("step", 0)
                    if current_step > last_step:
                        yield {
                            "type": "progress",
                            "progress": self._tasks[task_id]["progress"],
                            "step": current_step,
                            "total_steps": request.num_inference_steps,
                        }
                        last_step = current_step
                output = await inference_task
            else:
                output = await loop.run_in_executor(None, run_inference)
            
            images_base64 = [self._encode_image_to_base64(img) for img in output.images]
            
            yield {
                "type": "result",
                "images": images_base64,
                "seed": seed,
                "task_id": task_id,
            }
            
        except Exception as e:
            uvicorn_logger.error(f"txt2img error: {e}")
            yield {"type": "error", "error": str(e)}
        finally:
            del self._tasks[task_id]
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    async def _txt2img_sse(self, request: Txt2ImgRequest) -> AsyncGenerator[str, None]:
        """文生图 SSE 生成器"""
        async for data in self._txt2img(request, stream=True):
            yield f"data: {json.dumps(data)}\n\n"

    async def _img2img(self, request: Img2ImgRequest, stream: bool = False) -> AsyncGenerator[Dict[str, Any], None]:
        """图生图推理"""
        if not self._img2img_pipe:
            yield {"type": "error", "error": "img2img pipeline not available"}
            return
        
        torch = self._modules["torch"]
        task_id = str(uuid.uuid4())
        self._tasks[task_id] = {"progress": 0, "step": 0}
        
        try:
            init_image = self._decode_base64_image(request.image)
            seed = request.seed if request.seed is not None else int(torch.randint(0, 2**32 - 1, (1,)).item())
            generator = torch.Generator(device=self._device).manual_seed(seed)
            
            callback = self._create_progress_callback(task_id, request.num_inference_steps) if stream else None
            
            loop = asyncio.get_event_loop()
            
            def run_inference():
                return self._img2img_pipe(
                    prompt=request.prompt,
                    negative_prompt=request.negative_prompt or None,
                    image=init_image,
                    strength=request.strength,
                    num_inference_steps=request.num_inference_steps,
                    guidance_scale=request.guidance_scale,
                    num_images_per_prompt=request.num_images,
                    generator=generator,
                    callback_on_step_end=callback,
                )
            
            output = await loop.run_in_executor(None, run_inference)
            images_base64 = [self._encode_image_to_base64(img) for img in output.images]
            
            yield {
                "type": "result",
                "images": images_base64,
                "seed": seed,
                "task_id": task_id,
            }
            
        except Exception as e:
            uvicorn_logger.error(f"img2img error: {e}")
            yield {"type": "error", "error": str(e)}
        finally:
            del self._tasks[task_id]
            gc.collect()

    async def _img2img_sse(self, request: Img2ImgRequest) -> AsyncGenerator[str, None]:
        """图生图 SSE 生成器"""
        async for data in self._img2img(request, stream=True):
            yield f"data: {json.dumps(data)}\n\n"

    async def _inpaint(self, request: InpaintRequest, stream: bool = False) -> AsyncGenerator[Dict[str, Any], None]:
        """图像修复推理"""
        if not self._inpaint_pipe:
            yield {"type": "error", "error": "inpaint pipeline not available"}
            return
        
        torch = self._modules["torch"]
        task_id = str(uuid.uuid4())
        self._tasks[task_id] = {"progress": 0, "step": 0}
        
        try:
            init_image = self._decode_base64_image(request.image)
            mask_image = self._decode_base64_image(request.mask_image)
            
            width = request.width or init_image.width
            height = request.height or init_image.height
            
            seed = request.seed if request.seed is not None else int(torch.randint(0, 2**32 - 1, (1,)).item())
            generator = torch.Generator(device=self._device).manual_seed(seed)
            
            callback = self._create_progress_callback(task_id, request.num_inference_steps) if stream else None
            
            loop = asyncio.get_event_loop()
            
            def run_inference():
                return self._inpaint_pipe(
                    prompt=request.prompt,
                    negative_prompt=request.negative_prompt or None,
                    image=init_image,
                    mask_image=mask_image,
                    width=width,
                    height=height,
                    strength=request.strength,
                    num_inference_steps=request.num_inference_steps,
                    guidance_scale=request.guidance_scale,
                    num_images_per_prompt=request.num_images,
                    generator=generator,
                    callback_on_step_end=callback,
                )
            
            output = await loop.run_in_executor(None, run_inference)
            images_base64 = [self._encode_image_to_base64(img) for img in output.images]
            
            yield {
                "type": "result",
                "images": images_base64,
                "seed": seed,
                "task_id": task_id,
            }
            
        except Exception as e:
            uvicorn_logger.error(f"inpaint error: {e}")
            yield {"type": "error", "error": str(e)}
        finally:
            del self._tasks[task_id]
            gc.collect()

    async def _inpaint_sse(self, request: InpaintRequest) -> AsyncGenerator[str, None]:
        """图像修复 SSE 生成器"""
        async for data in self._inpaint(request, stream=True):
            yield f"data: {json.dumps(data)}\n\n"

    # ========== ieops 协议兼容 ==========
    async def _ieops_sse_generator(self, request: bytes) -> AsyncGenerator[str, None]:
        """ieops 协议 SSE 生成器"""
        async for item in self._ieops_generator(request):
            code = item["code"]
            data = item["data"]
            yield f"{json.dumps({'code': code, 'data': data}, ensure_ascii=False)}\n\n"

    async def _ieops_generator(self, request: bytes) -> AsyncGenerator[Dict[str, Any], None]:
        """ieops 协议生成器"""
        start = int(time.time() * 1000)
        
        try:
            request_data = json.loads(request.decode("utf-8"))
            headers = request_data.get("headers", {})
            payload = request_data.get("payload", request_data)
        except json.JSONDecodeError:
            payload = json.loads(request.decode("utf-8"))
            headers = {}
        
        trace_id = str(headers.get("x-trace-id", f"traceid-{await arandstr(8)}"))
        logger = log.get_logger(trace_id=trace_id)
        logger.info("Diffuser request received")
        
        if isinstance(payload, dict):
            task_type = payload.get("type", "txt2img")
            task_params = payload.get("params", payload)
        else:
            task_params = json.loads(payload) if isinstance(payload, str) else payload
            task_type = task_params.get("type", "txt2img")
        
        if self._on_request_received:
            task_params = self._on_request_received(task_params)
        
        try:
            if task_type == "txt2img":
                req = Txt2ImgRequest(**task_params)
                generator = self._txt2img(req, stream=True)
            elif task_type == "img2img":
                req = Img2ImgRequest(**task_params)
                generator = self._img2img(req, stream=True)
            elif task_type == "inpaint":
                req = InpaintRequest(**task_params)
                generator = self._inpaint(req, stream=True)
            else:
                yield {"code": SERVER_CODE_ERROR, "data": f"Unknown task type: {task_type}"}
                return
            
            async for data in generator:
                if data.get("type") == "error":
                    yield {"code": SERVER_CODE_ERROR, "data": data["error"]}
                    return
                elif data.get("type") == "progress":
                    yield {"code": SERVER_CODE_OK, "data": json.dumps(data, ensure_ascii=False)}
                elif data.get("type") == "result":
                    yield {"code": SERVER_CODE_OK, "data": json.dumps(data, ensure_ascii=False)}
                    yield {"code": SERVER_CODE_STOP, "data": ""}
                    return
            
            yield {"code": SERVER_CODE_STOP, "data": ""}
            
        except Exception as e:
            logger.error(f"Error: {e}")
            yield {"code": SERVER_CODE_ERROR, "data": str(e)}
        finally:
            cost = int(time.time() * 1000) - start
            logger.info(f"Generation completed, cost: {cost}ms")
