comfyUI

import httpx, time
from fastapi import FastAPI, Response
from fastapi.middleware.cors import CORSMiddleware # 导入CORS中间件
from fastapi.openapi.docs import get_swagger_ui_html # 导入生成API文档的函数
from fastapi.openapi.utils import get_openapi # 导入生成API文档的函数
from pydantic import BaseModel, Field
import random
COMFY = "http://127.0.0.1:8188"  # ComfyUI 地址
# 局域网IP
LAN_IP = "http://192.168.1.140:8188"


# 免费的谷歌翻译
def google_free_translate(text: str, to_lang="en", from_lang="auto", retries=2, timeout=8) -> str:
    """
    使用 Google 非公开端点进行翻译（无 API key）。
    仅用于临时/个人用途，可能随时失效或被限流。  # 中文注释
    """
    url = "https://translate.googleapis.com/translate_a/single"
    params = {
        "client": "gtx",   # 网页客户端标识
        "sl": from_lang,   # 源语言
        "tl": to_lang,     # 目标语言
        "dt": "t",         # 只要翻译文本
        "q": text          # 待翻译内容
    }
    headers = {
        # 伪装成浏览器，降低被风控概率
        "User-Agent": "Mozilla/5.0"
    }
    for i in range(retries + 1):
        try:
            r = httpx.get(url, params=params, headers=headers, timeout=timeout)
            r.raise_for_status()
            data = r.json()
            print(data)
            # data[0] 是句子片段列表，每项形如 [译文, 原文, ...]
            return "".join(seg[0] for seg in data[0] if seg and seg[0])
        except Exception:
            if i == retries:
                raise
            print(f"翻译失败: {e}")
            time.sleep(0.6 * (i + 1))  # 简单退避

# 拼装提示词
def convert_prompt(age: int, gender: str, city: str,bio: str="") -> str:
    # 精英图片缩小年龄，让图片更年轻
    if age > 55:
        age = 50
    elif age > 45:
        age -= 5
    elif age > 40:
        age -= 3

    info = f"来自{city}的{age}岁{gender}"
    if bio:
        info += f",{bio}"
    return info


# 设置工作流，返回 prompt_json
def set_prompt(prompt: str) -> dict:
    # 随机种子
    seed = random.randint(1, 2**31 - 1)
    prompt_json = {
        "prompt": {
            "6": {
                "inputs": {
                    "text": prompt,
                    "clip": ["30", 1],
                },
                "class_type": "CLIPTextEncode",
                "_meta": {"title": "CLIP Text Encode (Positive Prompt)"},
            },
            "8": {
                "inputs": {"samples": ["31", 0], "vae": ["30", 2]},
                "class_type": "VAEDecode",
                "_meta": {"title": "VAE解码"},
            },
            "9": {
                "inputs": {"filename_prefix": "image", "images": ["8", 0]},
                "class_type": "SaveImage",
                "_meta": {"title": "保存图像"},
            },
            "27": {
                "inputs": {"width": 496, "height": 496, "batch_size": 1},
                "class_type": "EmptySD3LatentImage",
                "_meta": {"title": "空Latent图像（SD3）"},
            },
            "30": {
                "inputs": {"ckpt_name": "flux1-dev-fp8.safetensors"},
                "class_type": "CheckpointLoaderSimple",
                "_meta": {"title": "Checkpoint加载器（简易）"},
            },
            "31": {
                "inputs": {
                    "seed": seed,
                    "steps": 20,
                    "cfg": 1,
                    "sampler_name": "euler",
                    "scheduler": "simple",
                    "denoise": 1,
                    "model": ["30", 0],
                    "positive": ["35", 0],
                    "negative": ["33", 0],
                    "latent_image": ["27", 0],
                },
                "class_type": "KSampler",
                "_meta": {"title": "K采样器"},
            },
            "33": {
                "inputs": {"text": "cartoon, sad, scary, negative", "clip": ["30", 1]},
                "class_type": "CLIPTextEncode",
                "_meta": {"title": "CLIP Text Encode (Negative Prompt)"},
            },
            "35": {
                "inputs": {"guidance": 3.5, "conditioning": ["6", 0]},
                "class_type": "FluxGuidance",
                "_meta": {"title": "Flux引导"},
            },
        },
        "client_id": "demo-client",
    }
    return prompt_json

# 提交任务
def submit(payload) -> str:
    """
    提交任务
    返回队列里的 prompt_id
    """
    r = httpx.post(f"{COMFY}/prompt", json=payload, timeout=30)
    r.raise_for_status()
    return r.json()["prompt_id"]  # 返回队列里的 prompt_id


# 等待结果
def wait_result(prompt_id, timeout_s=120):
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        h = httpx.get(f"{COMFY}/history/{prompt_id}", timeout=15)
        if h.status_code == 200:
            data = h.json().get(prompt_id, {})
            outs = data.get("outputs", {})
            urls = []
            for _, node_out in outs.items():
                for img in node_out.get("images", []):
                    filename = img["filename"]
                    subfolder = img.get("subfolder", "")
                    typ = img.get("type", "output")
                    # 通过 /view 直接访问生成图
                    urls.append(
                        f"{COMFY}/view?filename={filename}&subfolder={subfolder}&type={typ}"
                    )
            if urls:
                return urls
        time.sleep(0.5)
    raise TimeoutError("等待结果超时")

app = FastAPI()

# 配置CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有源
    allow_credentials=True, # 允许跨域请求
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

# 定义请求参数模型
class Text2ImageRequest(BaseModel):
    """请求参数模型"""
    age: int = Field(..., description="年龄")
    gender: str = Field(..., description="性别")
    city: str = Field(..., description="城市")
    bio: str = Field(default="", description="个人简介，选填")
    # 传入语言
    language: str = Field(default="zh", description="传入的语言，en则不会启动翻译")

# 定义响应模型
@app.post("/text2image")
async def text2imaget(data: Text2ImageRequest,response: Response) -> str:
    # 拼接提示词
    prompt = convert_prompt(data.age, data.gender, data.city, data.bio)
    
    # 翻译
    try:
        if data.language != "en":
            prompt = google_free_translate(prompt)
        prompt = f"An elite professional,{prompt}"
    except Exception as e:
        print(f"翻译失败: {e}")
        response.status_code = 500
        return f"翻译失败: {e}"
    
    # 设置工作流
    prompt_json = set_prompt(prompt)
    # 提交任务
    pid = submit(prompt_json)
    # 返回任务ID
    return pid

# 获取任务结果
@app.get("/text2image/{pid}")
async def get_text2image_result(pid: str, response: Response) -> str:
    """
    获取图像生成结果
    返回单个图像URL，不是数组
    """
    h = httpx.get(f"{COMFY}/history/{pid}", timeout=15)
    if h.status_code == 200:
        data = h.json().get(pid, {})
        outs = data.get("outputs", {})

        for _, node_out in outs.items():
            for img in node_out.get("images", []):
                filename = img["filename"]
                subfolder = img.get("subfolder", "")
                typ = img.get("type", "output")
                # 通过 /view 直接访问生成图，返回第一个找到的图片
                return f"{LAN_IP}/view?filename={filename}&subfolder={subfolder}&type={typ}"
    # 如果没有找到图片，设置响应状态码并返回错误信息
    response.status_code = 400
    return "图片尚未生成完成或任务不存在"

# 自定义API文档
def custom_openapi():
    """生成自定义OpenAPI架构"""
    if app.openapi_schema:
        return app.openapi_schema
    
    openapi_schema = get_openapi(
        title="示例API",
        version="1.0.0",
        description="这是一个FastAPI示例应用，展示API文档生成",
        routes=app.routes,
    )
    
    # 自定义徽标信息
    openapi_schema["info"]["x-logo"] = {
        "url": "https://fastapi.tiangolo.com/img/logo-margin/logo-teal.png"
    }
    
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app.openapi = custom_openapi

# 自定义API文档路径
@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html():
    """自定义Swagger UI文档页面"""
    return get_swagger_ui_html(
        openapi_url=app.openapi_url,
        title=f"{app.title} - Swagger UI",
        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
        swagger_js_url="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5/swagger-ui-bundle.js",
        swagger_css_url="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5/swagger-ui.css",
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
    # convert_prompt(45,"男性","北京")

parser.add_argument("--listen", type=str, default="0.0.0.0", metavar="IP", nargs="?", const="0.0.0.0,::")