TacoLab/FLOW2API

1import os2import json3import re4import base645import aiohttp # Async test. Need to install6import asyncio7 8 9# --- 配置区域 ---10BASE_URL = os.getenv('GEMINI_FLOW2API_URL', 'http://127.0.0.1:8000')11BACKEND_URL = BASE_URL + "/v1/chat/completions"12API_KEY = os.getenv('GEMINI_FLOW2API_APIKEY', 'Bearer han1234')13if API_KEY is None:14    raise ValueError('[gemini flow2api] api key not set')15MODEL_LANDSCAPE = "gemini-3.0-pro-image-landscape"16MODEL_PORTRAIT = "gemini-3.0-pro-image-portrait"17 18# 修改: 增加 model 参数，默认为 None19async def request_backend_generation(20        prompt: str,21        images: list[bytes] = None,22        model: str = None) -> bytes | None:23    """24    请求后端生成图片。25    :param prompt: 提示词26    :param images: 图片二进制列表27    :param model: 指定模型名称 (可选)28    :return: 成功返回图片bytes，失败返回None29    """30    # 更新token31    images = images or []32    33    # 逻辑: 如果未指定 model，默认使用 Landscape34    use_model = model if model else MODEL_LANDSCAPE35 36    # 1. 构造 Payload37    if images:38        content_payload = [{"type": "text", "text": prompt}]39        print(f"[Backend] 正在处理 {len(images)} 张图片输入...")40        for img_bytes in images:41            b64_str = base64.b64encode(img_bytes).decode('utf-8')42            content_payload.append({43                "type": "image_url",44                "image_url": {"url": f"data:image/jpeg;base64,{b64_str}"}45            })46    else:47        content_payload = prompt48 49    payload = {50        "model": use_model,  # 使用选定的模型51        "messages": [{"role": "user", "content": content_payload}],52        "stream": True53    }54    55    headers = {56        "Authorization": API_KEY,57        "Content-Type": "application/json"58    }59 60    image_url = None61    print(f"[Backend] Model: {use_model} | 发起请求: {prompt[:20]}...") 62    63    try:64        async with aiohttp.ClientSession() as session:65            async with session.post(BACKEND_URL, json=payload, headers=headers, timeout=120) as response:66                if response.status != 200:67                    err_text = await response.text()68                    content = response.content69                    print(f"[Backend Error] Status {response.status}: {err_text} {content}")70                    raise Exception(f"API Error: {response.status}: {err_text}")71 72                async for line in response.content:73                    line_str = line.decode('utf-8').strip()74                    if line_str.startswith('{"error'):75                        chunk = json.loads(data_str)76                        delta = chunk.get("choices", [{}])[0].get("delta", {})77                        msg = delta['reasoning_content']78                        if '401' in msg:79                            msg += '\nAccess Token 已失效，需重新配置。'80                        elif '400' in msg:81                            msg += '\n返回内容被拦截。'82                        raise Exception(msg)83 84                    if not line_str or not line_str.startswith('data: '):85                        continue86                    87                    data_str = line_str[6:]88                    if data_str == '[DONE]':89                        break90                    91                    try:92                        chunk = json.loads(data_str)93                        delta = chunk.get("choices", [{}])[0].get("delta", {})94                        95                        # 打印思考过程96                        if "reasoning_content" in delta:97                            print(delta['reasoning_content'], end="", flush=True)98 99                        # 提取内容中的图片链接100                        if "content" in delta:101                            content_text = delta["content"]102                            img_match = re.search(r'!\[.*?\]\((.*?)\)', content_text)103                            if img_match:104                                image_url = img_match.group(1)105                                print(f"\n[Backend] 捕获图片链接: {image_url}")106                    except json.JSONDecodeError:107                        continue108            109            # 3. 下载生成的图片110            if image_url:111                async with session.get(image_url) as img_resp:112                    if img_resp.status == 200:113                        image_bytes = await img_resp.read()114                        return image_bytes115                    else:116                        print(f"[Backend Error] 图片下载失败: {img_resp.status}")117    except Exception as e:118        print(f"[Backend Exception] {e}")119        raise e 120        121    return None122 123if __name__ == '__main__':124    async def main():125        print("=== AI 绘图接口测试 ===")126        user_prompt = input("请输入提示词 (例如 '一只猫'): ").strip()127        if not user_prompt:128            user_prompt = "A cute cat in the garden"129        130        print(f"正在请求: {user_prompt}")131        132        # 这里的 images 传空列表用于测试文生图133        # 如果想测试图生图，你需要手动读取本地文件：134        # with open("output_test.jpg", "rb") as f: img_data = f.read()135        # result = await request_backend_generation(user_prompt, [img_data])136        137        result = await request_backend_generation(user_prompt)138        139        if result:140            filename = "output_test.jpg"141            with open(filename, "wb") as f:142                f.write(result)143            print(f"\n[Success] 图片已保存为 {filename}，大小: {len(result)} bytes")144        else:145            print("\n[Failed] 生成失败")146 147    # 运行测试148    if os.name == 'nt':  # Windows 兼容性149        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())150    asyncio.run(main())