'commit'

dsLightRag/Routes/QWenImageRoute.py

@@ -0,0 +1,147 @@
+import atexit
+import logging
+import os
+import shutil
+import tempfile  # 新增导入
+import uuid
+from typing import Optional
+
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel, Field
+
+from Config.Config import OBS_PREFIX, OBS_BUCKET, OBS_SERVER
+# 导入QwenImageGenerator类
+from QWenImage.QwenImageKit import QwenImageGenerator
+from Util.ObsUtil import ObsUploader
+
+# 创建路由路由器
+router = APIRouter(prefix="/api/qwenImage", tags=["千问生图"])
+
+# 配置日志
+logger = logging.getLogger(__name__)
+
+# 初始化图片生成器
+image_generator = QwenImageGenerator()
+
+
+class GenerateImageRequest(BaseModel):
+    """生成图片请求模型"""
+    prompt: str = Field(..., description="图片描述提示词")
+    n: int = Field(default=1, ge=1, le=4, description="生成图片数量，范围1-4")
+    size: str = Field(default="1328*1328", description="图片尺寸")
+    api_key: Optional[str] = Field(default=None, description="自定义API密钥")
+
+
+@router.post("/generate")
+async def generate_image(request: GenerateImageRequest):
+    """生成图片API接口
+    
+    Args:
+        request: 包含生成图片参数的请求对象
+        
+    Returns:
+        dict: 包含生成结果的字典
+    """
+    try:
+        logger.info(f"接收到图片生成请求: prompt={request.prompt[:50]}..., n={request.n}, size={request.size}")
+        
+        # 如果提供了自定义API密钥，创建新的生成器实例
+        generator = QwenImageGenerator(api_key=request.api_key) if request.api_key else image_generator
+
+        # 仅生成图片，不保存本地
+        result = generator.generate_image(
+            prompt=request.prompt,
+            n=request.n,
+            size=request.size
+        )
+        
+        # 处理结果
+        if result['success']:
+            logger.info(f"图片生成成功，返回{len(result['images'])}张图片")
+            
+            # 构造返回响应
+            response = {
+                "code": 200,
+                "message": "图片生成成功",
+                "data": {
+                    "images": result['images']
+                }
+            }
+            # 新增：无条件执行OBS上传（直接处理图片URL）
+            obs_urls = []
+            uploader = ObsUploader()
+            for image_url in result['images']:
+                try:
+                    # 直接从URL下载图片二进制数据
+                    import requests
+                    response_img = requests.get(image_url, timeout=10)
+                    response_img.raise_for_status()
+                    bytes_data = response_img.content
+
+                    # 生成UUID文件名并上传
+                    jpg_file_name = f"{str(uuid.uuid4())}.jpg"
+                    object_key = f"{OBS_PREFIX}/QWen3Image/{jpg_file_name}"
+                    success, upload_result = uploader.upload_base64_image(object_key, bytes_data)
+
+                    if success:
+                        obs_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{object_key}"
+                        obs_urls.append(obs_url)
+                        logger.info(f"图片上传OBS成功: {obs_url}")
+                    else:
+                        logger.error(f"图片上传OBS失败: {upload_result}")
+                except Exception as e:
+                    logger.error(f"处理图片URL {image_url} 时出错: {str(e)}")
+
+            response["data"]["obs_files"] = obs_urls
+            return response
+        else:
+            logger.error(f"图片生成失败: {result['error_msg']}")
+            raise HTTPException(
+                status_code=500,
+                detail={
+                    "code": 500,
+                    "message": "图片生成失败",
+                    "error_detail": result['error_msg']
+                }
+            )
+            
+    except Exception as e:
+        logger.exception(f"处理图片生成请求时发生异常: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "code": 500,
+                "message": "处理请求时发生异常",
+                "error_detail": str(e)
+            }
+        )
+
+
+@router.get("/config")
+async def get_image_config():
+    """获取图片生成配置信息
+    
+    Returns:
+        dict: 包含配置信息的字典
+    """
+    return {
+        "code": 200,
+        "message": "获取配置成功",
+        "data": {
+            "supported_sizes": ["1328*1328", "1024*1024", "768*1024", "1024*768"],
+            "max_images_per_request": 4
+        }
+    }
+
+
+# 注册程序退出时的清理函数
+@atexit.register
+def clean_temp_files():
+    temp_root = os.path.join(tempfile.gettempdir(), "qwen_images")
+    if os.path.exists(temp_root):
+        try:
+            shutil.rmtree(temp_root)
+            logger.info(f"临时图片目录已清理: {temp_root}")
+        except Exception as e:
+            logger.warning(f"清理临时文件失败: {str(e)}")
+

dsLightRag/Start.py

@@ -23,6 +23,7 @@ from Routes.JiMengRoute import router as jimeng_router
 from Routes.SunoRoute import router as suno_router
 from Routes.XueBanRoute import router as xueban_router
 from Routes.MjRoute import router as mj_router
+from Routes.QWenImageRoute import router as qwen_image_router
 from Util.LightRagUtil import *
 from contextlib import asynccontextmanager
 
@@ -36,8 +37,8 @@ logger.addHandler(handler)
 
 @asynccontextmanager
 async def lifespan(_: FastAPI):
-    #pool = await init_postgres_pool()
-    #app.state.pool = pool
+    pool = await init_postgres_pool()
+    app.state.pool = pool
 
     asyncio.create_task(train_document_task())
 
@@ -45,7 +46,7 @@ async def lifespan(_: FastAPI):
         yield
     finally:
         # 应用关闭时销毁连接池
-        #await close_postgres_pool(pool)
+        await close_postgres_pool(pool)
         pass
 
 
@@ -66,6 +67,8 @@ app.include_router(jimeng_router)  # 即梦路由
 app.include_router(suno_router)  # Suno路由
 app.include_router(xueban_router)  # 学伴路由
 app.include_router(mj_router)  # Midjourney路由
+app.include_router(qwen_image_router) # Qwen Image 路由
+
 
 # Teaching Model 相关路由
 # 登录相关（不用登录）
@@ -84,4 +87,4 @@ app.include_router(teaching_model_router, prefix="/api/teaching/model", tags=["t
 app.include_router(teaching_model_router, prefix="/api/teaching/model", tags=["teacher_model"])
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8100)
+    uvicorn.run(app, host="0.0.0.0", port=8200)

dsLightRag/Test/TestQWen3Image.py

@@ -0,0 +1,140 @@
+import requests
+import json
+import logging
+import time
+import os
+
+# 配置日志
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger("TestQWenImage")
+
+# API基础URL
+base_url = "http://localhost:8200/api/qwenImage"
+
+
+def test_get_config():
+    """测试获取配置接口"""
+    try:
+        url = f"{base_url}/config"
+        logger.info(f"调用获取配置接口: {url}")
+        response = requests.get(url)
+        
+        if response.status_code == 200:
+            result = response.json()
+            logger.info(f"获取配置成功: {json.dumps(result, ensure_ascii=False, indent=2)}")
+            return True, result
+        else:
+            logger.error(f"获取配置失败: HTTP状态码={response.status_code}, 响应内容={response.text}")
+            return False, None
+    except Exception as e:
+        logger.exception(f"获取配置时发生异常: {str(e)}")
+        return False, None
+
+
+def test_generate_image(prompt, n=1, size='1328*1328', save_local=True):
+    """测试生成图片接口
+    
+    Args:
+        prompt: 图片描述提示词
+        n: 生成图片数量
+        size: 图片尺寸
+        save_local: 是否保存到本地
+        
+    Returns:
+        tuple: (是否成功, 结果数据)
+    """
+    try:
+        url = f"{base_url}/generate"
+        headers = {"Content-Type": "application/json"}
+        data = {
+            "prompt": prompt,
+            "n": n,
+            "size": size,
+            "save_local": save_local
+        }
+        
+        logger.info(f"调用生成图片接口: {url}")
+        logger.info(f"请求参数: prompt={prompt[:50]}..., n={n}, size={size}, save_local={save_local}")
+        
+        # 记录开始时间
+        start_time = time.time()
+        
+        # 发送请求
+        response = requests.post(url, headers=headers, data=json.dumps(data, ensure_ascii=False))
+        
+        # 计算耗时
+        elapsed_time = time.time() - start_time
+        logger.info(f"请求耗时: {elapsed_time:.2f}秒")
+        
+        if response.status_code == 200:
+            result = response.json()
+            logger.info(f"生成图片成功: {json.dumps(result, ensure_ascii=False, indent=2)}")
+            
+            # 检查返回的数据
+            if result.get("code") == 200 and "data" in result:
+                images = result["data"].get("images", [])
+                logger.info(f"成功生成{len(images)}张图片")
+                
+                # 如果保存了本地文件，检查文件是否存在
+                if save_local and "local_file_paths" in result["data"]:
+                    for file_path in result["data"]["local_file_paths"]:
+                        full_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), file_path.lstrip('.'))
+                        if os.path.exists(full_path):
+                            logger.info(f"本地文件已保存: {full_path}")
+                        else:
+                            logger.warning(f"本地文件不存在: {full_path}")
+                
+            return True, result
+        else:
+            logger.error(f"生成图片失败: HTTP状态码={response.status_code}, 响应内容={response.text}")
+            return False, None
+    except Exception as e:
+        logger.exception(f"生成图片时发生异常: {str(e)}")
+        return False, None
+
+
+def main():
+    """主函数，运行单元测试"""
+    logger.info("===== 开始测试QWenImage接口 ====")
+    
+    # 1. 测试获取配置接口
+    logger.info("\n1. 测试获取配置接口")
+    config_success, config_data = test_get_config()
+    
+    # 2. 测试生成图片接口 - 基本测试
+    logger.info("\n2. 测试生成图片接口 - 基本测试")
+    basic_prompt = "一只可爱的小猫在草地上玩耍，阳光明媚，高清细节"
+    generate_success, generate_data = test_generate_image(
+        prompt=basic_prompt,
+        n=1,
+        size="1328*1328",
+        save_local=True
+    )
+    
+    # 3. 测试生成图片接口 - 不同参数
+    if config_success:
+        supported_sizes = config_data["data"].get("supported_sizes", ["1328*1328"])
+        max_images = min(2, config_data["data"].get("max_images_per_request", 1))  # 为了测试效率，最多请求2张
+        
+        logger.info(f"\n3. 测试生成图片接口 - 不同参数(size={supported_sizes[0]}, n={max_images})")
+        different_prompt = "美丽的山水风景画，中国水墨画风格"
+        test_generate_image(
+            prompt=different_prompt,
+            n=max_images,
+            size=supported_sizes[0],
+            save_local=True
+        )
+    
+    logger.info("\n===== QWenImage接口测试完成 =====")
+    
+    # 输出测试结果摘要
+    success_count = sum([config_success, generate_success])
+    total_count = 2  # 基本测试的接口数量
+    logger.info(f"测试结果: {success_count}/{total_count} 接口测试成功")
+
+
+if __name__ == "__main__":
+    main()