Create main.py

main.py

@@ -0,0 +1,493 @@
+# -*- coding: utf-8 -*-
+
+import os
+import random
+import re
+import json
+import time
+import aiohttp
+import asyncio
+from aiohttp import web
+import unicodedata
+from dataclasses import dataclass, asdict
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Config:
+    # SILICONFLOW-Key
+    AUTH_TOKEN: str = 'sk-xxxxxxxxx'
+    # One-API/New-API 中转地址
+    OPENAI_CHAT_API: str = 'https://xxxxxxxxxxxx/v1/chat/completions'
+    # key
+    OPENAI_CHAT_API_KEY: str = 'sk-xxxxxxxxxx'
+    # 默认的翻译模型
+    DEFAULT_TRANSLATE_MODEL: str = 'deepseek-chat'
+    # 增强的翻译模型
+    DEFAULT_PROMPT_MODEL: str = 'Qwen2-72B-Instruct'
+
+
+config = Config()
+
+URLS = {
+    'API_FLUX1_API4GPT_COM': 'https://api-flux1.api4gpt.com',
+    'FLUXAIWEB_COM_TOKEN': 'https://fluxaiweb.com/flux/getToken',
+    'FLUXAIWEB_COM_GENERATE': 'https://fluxaiweb.com/flux/generateImage',
+    'FLUXIMG_COM': 'https://fluximg.com/api/image/generateImage',
+    'API_SILICONFLOW_CN': 'https://api.siliconflow.cn/v1/chat/completions'
+}
+
+URL_MAP = {
+    'flux': "https://api.siliconflow.cn/v1/black-forest-labs/FLUX.1-schnell/text-to-image",
+    'sd3': "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-3-medium/text-to-image",
+    'sdxl': "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-xl-base-1.0/text-to-image",
+    'sd2': "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-2-1/text-to-image",
+    'sdt': "https://api.siliconflow.cn/v1/stabilityai/sd-turbo/text-to-image",
+    'sdxlt': "https://api.siliconflow.cn/v1/stabilityai/sdxl-turbo/text-to-image",
+    'sdxll': "https://api.siliconflow.cn/v1/ByteDance/SDXL-Lightning/text-to-image"
+}
+
+IMG_URL_MAP = {
+    'sdxl': "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-xl-base-1.0/image-to-image",
+    'sd2': "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-2-1/image-to-image",
+    'sdxll': "https://api.siliconflow.cn/v1/ByteDance/SDXL-Lightning/image-to-image",
+    'pm': "https://api.siliconflow.cn/v1/TencentARC/PhotoMaker/image-to-image"
+}
+
+RATIO_MAP = {
+    "1:1": "1024x1024",
+    "1:2": "1024x2048",
+    "3:2": "1536x1024",
+    "4:3": "1536x2048",
+    "16:9": "2048x1152",
+    "9:16": "1152x2048"
+}
+
+SYSTEM_ASSISTANT = """作为 Stable Diffusion Prompt 提示词专家，您将从关键词中创建提示，通常来自 Danbooru 等数据库。
+提示通常描述图像，使用常见词汇，按重要性排列，并用逗号分隔。避免使用"-"或"."，但可以接受空格和自然语言。避免词汇重复。
+
+为了强调关键词，请将其放在括号中以增加其权重。例如，"(flowers)"将'flowers'的权重增加1.1倍，而"(((flowers)))"将其增加1.331倍。使用"(flowers:1.5)"将'flowers'的权重增加1.5倍。只为重要的标签增加权重。
+
+提示包括三个部分：**前缀**（质量标签+风格词+效果器）+ **主题**（图像的主要焦点）+ **场景**（背景、环境）。
+
+*   前缀影响图像质量。像"masterpiece"、"best quality"、"4k"这样的标签可以提高图像的细节。像"illustration"、"lensflare"这样的风格词定义图像的风格。像"bestlighting"、"lensflare"、"depthoffield"这样的效果器会影响光照和深度。
+
+*   主题是图像的主要焦点，如角色或场景。对主题进行详细描述可以确保图像丰富而详细。增加主题的权重以增强其清晰度。对于角色，描述面部、头发、身体、服装、姿势等特征。
+
+*   场景描述环境。没有场景，图像的背景是平淡的，主题显得过大。某些主题本身包含场景（例如建筑物、风景）。像"花草草地"、"阳光"、"河流"这样的环境词可以丰富场景。你的任务是设计图像生成的提示。请按照以下步骤进行操作：
+
+1.  我会发送给您一个图像场景。需要你生成详细的图像描述
+2.  图像描述必须是英文，输出为Positive Prompt。
+
+示例：
+
+我发送：二战时期的护士。
+您回复只回复：
+A WWII-era nurse in a German uniform, holding a wine bottle and stethoscope, sitting at a table in white attire, with a table in the background, masterpiece, best quality, 4k, illustration style, best lighting, depth of field, detailed character, detailed environment.
+"""
+
+
+async def select_random_image_generator():
+    generators = [generate_image1, generate_image2, generate_image3, generate_image4]
+    generator = random.choice(generators)
+    if generator == generate_image4:
+        return lambda prompt, size, model: generator(prompt, size, model)
+    else:
+        return lambda prompt, size, model=None: generator(prompt, size)
+
+
+def extract_size_and_model_from_prompt(prompt):
+    size_match = re.search(r'--ar\s+(\S+)', prompt)
+    model_match = re.search(r'--m\s+(\S+)', prompt)
+
+    size = size_match.group(1) if size_match else '1:1'
+    model = model_match.group(1) if model_match else ''
+
+    clean_prompt = re.sub(r'--ar\s+\S+', '', prompt).strip()
+    clean_prompt = re.sub(r'--m\s+\S+', '', clean_prompt).strip()
+
+    return {'size': size, 'model': model, 'clean_prompt': clean_prompt}
+
+
+def is_chinese(char):
+    return 'CJK' in unicodedata.name(char, '')
+
+
+async def translate_prompt(prompt):
+    if all(not is_chinese(char) for char in prompt):
+        logger.info('Prompt is already in English, skipping translation')
+        return prompt
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(config.OPENAI_CHAT_API, json={
+                'model': config.DEFAULT_TRANSLATE_MODEL,  # 使用config中的model
+                'messages': [
+                    {'role': 'system', 'content': SYSTEM_ASSISTANT},
+                    {'role': 'user', 'content': prompt}
+                ],
+            }, headers={
+                'Content-Type': 'application/json',
+                'Authorization': f'Bearer {config.OPENAI_CHAT_API_KEY}'
+            }) as response:
+                if response.status != 200:
+                    error_text = await response.text()
+                    logger.error(f'HTTP error! status: {response.status}, body: {error_text}')
+                    raise Exception(f'HTTP error! status: {response.status}')
+
+                if 'application/json' not in response.headers.get('Content-Type', ''):
+                    error_text = await response.text()
+                    logger.error(f'Unexpected content type: {response.headers.get("Content-Type")}, body: {error_text}')
+                    raise Exception(f'Unexpected content type: {response.headers.get("Content-Type")}')
+
+                result = await response.json()
+                return result['choices'][0]['message']['content']
+
+    except Exception as e:
+        logger.error('Translation error:', e)
+        return prompt
+
+
+async def handle_request(request):
+    if request.method != 'POST' or not request.url.path.endswith('/v1/chat/completions'):
+        return web.Response(text='Not Found', status=404)
+
+    try:
+        data = await request.json()
+        messages = data.get('messages', [])
+        stream = data.get('stream', False)
+
+        user_message = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), None)
+
+        if not user_message:
+            return web.json_response({'error': "未找到用户消息"}, status=400)
+
+        size_and_model = extract_size_and_model_from_prompt(user_message)
+        translated_prompt = await translate_prompt(size_and_model['clean_prompt'])
+
+        selected_generator = await select_random_image_generator()
+        attempts = 0
+        max_attempts = 3
+
+        while attempts < max_attempts:
+            try:
+                image_data = await selected_generator(translated_prompt, size_and_model['size'],
+                                                      size_and_model['model'])
+                break
+            except Exception as e:
+                logger.error(f"Error generating image with generator {selected_generator.__name__}: {e}")
+                selected_generator = await select_random_image_generator()
+                attempts += 1
+
+        if attempts == max_attempts:
+            logger.error("Failed to generate image after multiple attempts")
+            return web.json_response({'error': "生成图像失败"}, status=500)
+
+        unique_id = f"chatcmpl-{int(time.time())}"
+        created_timestamp = int(time.time())
+        model_name = "flux"
+        system_fingerprint = "fp_" + ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=9))
+
+        if stream:
+            return await handle_stream_response(request, unique_id, image_data, size_and_model['clean_prompt'],
+                                                translated_prompt, size_and_model['size'], created_timestamp,
+                                                model_name, system_fingerprint)
+        else:
+            return handle_non_stream_response(unique_id, image_data, size_and_model['clean_prompt'], translated_prompt,
+                                              size_and_model['size'], created_timestamp, model_name, system_fingerprint)
+    except Exception as e:
+        logger.error('Error handling request:', e)
+        return web.json_response({'error': f"处理请求失败: {str(e)}"}, status=500)
+
+
+async def handle_stream_response(request, unique_id, image_data, original_prompt, translated_prompt, size, created,
+                                 model, system_fingerprint):
+    logger.debug("Starting stream response")
+    response = web.StreamResponse(
+        status=200,
+        reason='OK',
+        headers={
+            'Content-Type': 'text/event-stream',
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive'
+        }
+    )
+    await response.prepare(request)
+    logger.debug("Response prepared")
+
+    chunks = [
+        f"原始提示词：\n{original_prompt}\n",
+        f"翻译后的提示词：\n{translated_prompt}\n",
+        f"图像规格���{size}\n",
+        "正在根据提示词生成图像...\n",
+        "图像正在处理中...\n",
+        "即将完成...\n",
+        f"生成成功！\n图像生成完毕，以下是结果：\n\n![生成的图像]({image_data['data'][0]['url']})"
+    ]
+
+    for i, chunk in enumerate(chunks):
+        json_chunk = json.dumps({
+            "id": unique_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": model,
+            "system_fingerprint": system_fingerprint,
+            "choices": [{
+                "index": 0,
+                "delta": {"content": chunk},
+                "logprobs": None,
+                "finish_reason": None
+            }]
+        })
+        try:
+            await response.write(f"data: {json_chunk}\n\n".encode('utf-8'))
+            logger.debug(f"Chunk {i + 1} sent")
+        except Exception as e:
+            logger.error(f"Error sending chunk {i + 1}: {str(e)}")
+        await asyncio.sleep(0.5)  # 模拟生成时间
+
+    final_chunk = json.dumps({
+        "id": unique_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+        "system_fingerprint": system_fingerprint,
+        "choices": [{
+            "index": 0,
+            "delta": {},
+            "logprobs": None,
+            "finish_reason": "stop"
+        }]
+    })
+    try:
+        await response.write(f"data: {final_chunk}\n\n".encode('utf-8'))
+        logger.debug("Final chunk sent")
+    except Exception as e:
+        logger.error(f"Error sending final chunk: {str(e)}")
+
+    await response.write_eof()
+    logger.debug("Stream response completed")
+    return response
+
+
+def handle_non_stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model,
+                               system_fingerprint):
+    content = (
+        f"原始提示词：{original_prompt}\n"
+        f"翻译后的提示词：{translated_prompt}\n"
+        f"图像规格：{size}\n"
+        f"图像生成成功！\n"
+        f"以下是结果：\n\n"
+        f"![生成的图像]({image_data['data'][0]['url']})"
+    )
+
+    response = {
+        'id': unique_id,
+        'object': "chat.completion",
+        'created': created,
+        'model': model,
+        'system_fingerprint': system_fingerprint,
+        'choices': [{
+            'index': 0,
+            'message': {
+                'role': "assistant",
+                'content': content
+            },
+            'finish_reason': "stop"
+        }],
+        'usage': {
+            'prompt_tokens': len(original_prompt),
+            'completion_tokens': len(content),
+            'total_tokens': len(original_prompt) + len(content)
+        }
+    }
+
+    return web.json_response(response)
+
+
+async def generate_image1(prompt, size):
+    # 调用 get_prompt 函数来增强提示词
+    enhanced_prompt = await get_prompt(prompt)
+
+    prompt_without_spaces = enhanced_prompt.replace(" ", "")
+    image_url = f"{URLS['API_FLUX1_API4GPT_COM']}/?prompt={prompt_without_spaces}&size={size}"
+    return {
+        'data': [{'url': image_url}],
+        'size': size
+    }
+
+
+async def generate_image2(prompt, size):
+    random_ip = generate_random_ip()
+    # 调用 get_prompt 来增强提示词
+    enhanced_prompt = await get_prompt(prompt)
+    async with aiohttp.ClientSession() as session:
+        async with session.get(URLS['FLUXAIWEB_COM_TOKEN'],
+                               headers={'X-Forwarded-For': random_ip}) as token_response:
+            token_data = await token_response.json()
+            token = token_data['data']['token']
+
+        async with session.post(URLS['FLUXAIWEB_COM_GENERATE'], headers={
+            'Content-Type': 'application/json',
+            'token': token,
+            'X-Forwarded-For': random_ip
+        }, json={
+            'prompt': enhanced_prompt,
+            'aspectRatio': size,
+            'outputFormat': 'webp',
+            'numOutputs': 1,
+            'outputQuality': 90
+        }) as image_response:
+            image_data = await image_response.json()
+            return {
+                'data': [{'url': image_data['data']['image']}],
+                'size': size
+            }
+
+
+async def generate_image3(prompt, size):
+    json_body = {
+        'textStr': prompt,
+        'model': "black-forest-labs/flux-schnell",
+        'size': size
+    }
+
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(URLS['FLUXIMG_COM'], data=json.dumps(json_body),
+                                        headers={'Content-Type': 'text/plain;charset=UTF-8'}) as response:
+                    if response.status == 200:
+                        image_url = await response.text()
+                        return {
+                            'data': [{'url': image_url}],
+                            'size': size
+                        }
+                    else:
+                        logger.error(
+                            f"Unexpected response status: {response.status}, response text: {await response.text()}")
+        except aiohttp.ClientConnectorError as e:
+            logger.error(f"Connection error on attempt {attempt + 1}: {e}")
+            await asyncio.sleep(2 ** attempt)  # Exponential backoff
+
+    logger.error("Failed to generate image after multiple attempts")
+    return {
+        'data': [{'url': "https://via.placeholder.com/640x480/428675/ffffff?text=Error"}],
+        'size': size
+    }
+
+
+async def generate_image4(prompt, size, model):
+    if not config.AUTH_TOKEN:
+        raise Exception("AUTH_TOKEN is required for this method")
+
+    api_url = URL_MAP.get(model, URL_MAP['flux'])
+    clean_prompt = re.sub(r'--m\s+\S+', '', prompt).strip()
+
+    # 调用 get_prompt 函数来增强提示词
+    enhanced_prompt = await get_prompt(clean_prompt)
+
+    json_body = {
+        'prompt': enhanced_prompt,
+        'image_size': RATIO_MAP.get(size, "1024x1024"),
+        'num_inference_steps': 50
+    }
+
+    if model and model != "flux":
+        json_body['batch_size'] = 1
+        json_body['guidance_scale'] = 7.5
+
+        if model in ["sdt", "sdxlt"]:
+            json_body['num_inference_steps'] = 6
+            json_body['guidance_scale'] = 1
+        elif model == "sdxll":
+            json_body['num_inference_steps'] = 4
+            json_body['guidance_scale'] = 1
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(api_url, headers={
+            'authorization': config.AUTH_TOKEN if config.AUTH_TOKEN.startswith(
+                'Bearer ') else f'Bearer {config.AUTH_TOKEN}',
+            'Accept': 'application/json',
+            'Content-Type': 'application/json'
+        }, json=json_body) as response:
+            if response.status != 200:
+                raise Exception(f'Unexpected response {response.status}')
+
+            json_response = await response.json()
+            return {
+                'data': [{'url': json_response['images'][0]['url']}],
+                'size': size
+            }
+
+
+def generate_random_ip():
+    return f"{random.randint(0, 255)}.{random.randint(0, 255)}.{random.randint(0, 255)}.{random.randint(0, 255)}"
+
+
+async def get_prompt(prompt):
+    logger.info(f"Original Prompt: {prompt}")  # 记录输入的原始提示词
+    request_body_json = json.dumps({
+        'model': config.DEFAULT_PROMPT_MODEL,  # 使用config中的model
+        'messages': [
+            {
+                'role': "system",
+                'content': SYSTEM_ASSISTANT
+            },
+            {
+                'role': "user",
+                'content': prompt
+            }
+        ],
+        'stream': False,
+        'max_tokens': 512,
+        'temperature': 0.7,
+        'top_p': 0.7,
+        'top_k': 50,
+        'frequency_penalty': 0.5,
+        'n': 1
+    })
+
+    # 打印出请求的详细信息
+    logger.debug(f"Request Body: {request_body_json}")
+    request_headers = {
+        'accept': 'application/json',
+        'authorization': config.OPENAI_CHAT_API_KEY if config.OPENAI_CHAT_API_KEY.startswith(
+            'Bearer ') else f'Bearer {config.OPENAI_CHAT_API_KEY}',
+        'content-type': 'application/json'
+    }
+    logger.debug(f"Request Headers: {request_headers}")
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(config.OPENAI_CHAT_API, headers=request_headers,
+                                    data=request_body_json) as response:
+                if response.status != 200:
+                    error_text = await response.text()
+                    logger.error(f"Failed to get response, status code: {response.status}, response: {error_text}")
+                    return prompt
+
+                json_response = await response.json()
+                logger.debug(f"API Response: {json_response}")  # 记录API的完整响应
+
+                if 'choices' in json_response and len(json_response['choices']) > 0:
+                    enhanced_prompt = json_response['choices'][0]['message']['content']
+                    logger.info(f"Enhanced Prompt: {enhanced_prompt}")  # 记录增强后的提示词
+                    return enhanced_prompt
+                else:
+                    logger.warning("No enhanced prompt found in the response, returning original prompt.")
+                    return prompt
+    except Exception as e:
+        logger.error(f"Exception occurred: {e}")
+        return prompt
+
+
+app = web.Application()
+app.router.add_post('/v1/chat/completions', handle_request)
+
+if __name__ == '__main__':
+    web.run_app(app, host='0.0.0.0', port=6666)