first cimmit

app/config/config.py

@@ -17,7 +17,10 @@ class Settings(BaseSettings):
         extra='ignore'  # 忽略环境变量中多余的键
     )
     # 启动端口
-    SERVE_PROD: int = Field(default=8000, description='')
+    SERVE_PORT: int = Field(default=10090, description='')
+
+    FLUX2_KLEIN_MODEL_PATH: str = Field(default="/models", description='')
+
     # 调试配饰
     LOCAL: int = Field(default=0, description="是否在本地运行，1表示本地运行，0表示生产环境运行")
 

app/logging_env.py

@@ -0,0 +1,50 @@
+LOGS_PATH = './logs/'
+LOGGER_CONFIG_DICT = {
+    'version': 1,
+    'disable_existing_loggers': False,
+    'formatters': {
+        'simple': {'format': '%(asctime)s %(filename)s [line:%(lineno)d] %(levelname)s %(message)s'}
+    },
+    'handlers': {
+        'console': {
+            'class': 'logging.StreamHandler',
+            'level': 'INFO',
+            'formatter': 'simple',
+            'stream': 'ext://sys.stdout',
+        },
+        'info_file_handler': {
+            'class': 'logging.handlers.RotatingFileHandler',
+            'level': 'INFO',
+            'formatter': 'simple',
+            'filename': f'{LOGS_PATH}info.log',
+            'maxBytes': 10485760,
+            'backupCount': 50,
+            'encoding': 'utf8',
+        },
+        'error_file_handler': {
+            'class': 'logging.handlers.RotatingFileHandler',
+            'level': 'ERROR',
+            'formatter': 'simple',
+            'filename': f'{LOGS_PATH}error.log',
+            'maxBytes': 10485760,
+            'backupCount': 20,
+            'encoding': 'utf8',
+        },
+        'debug_file_handler': {
+            'class': 'logging.handlers.RotatingFileHandler',
+            'level': 'DEBUG',
+            'formatter': 'simple',
+            'filename': f'{LOGS_PATH}debug.log',
+            'maxBytes': 10485760,
+            'backupCount': 50,
+            'encoding': 'utf8',
+        },
+    },
+    'loggers': {
+        'my_module': {'level': 'INFO', 'handlers': ['console'], 'propagate': 'no'}
+    },
+    'root': {
+        'level': 'INFO',
+        'handlers': ['error_file_handler', 'info_file_handler', 'debug_file_handler', 'console'],
+    },
+}

app/main.py

@@ -0,0 +1,24 @@
+import logging
+import os
+
+import litserve as ls
+
+from app.logging_env import LOGGER_CONFIG_DICT
+from app.server.generate_image.flux2_klein.server import Flux2KleinServer
+
+logger = logging.getLogger(__name__)
+# 判断目录是否存在
+logs = 'logs'
+if not os.path.exists(logs):
+    # 不存在则创建目录（parents=True 允许创建多级目录，exist_ok=True 避免目录已存在时报错）
+    os.makedirs(logs, exist_ok=True)
+    logger.info(f"目录 {logs} 创建成功")
+else:
+    logger.info(f"目录 {logs} 已存在")
+logging.config.dictConfig(LOGGER_CONFIG_DICT)
+
+if __name__ == '__main__':
+    to_product = Flux2KleinServer(api_path='/api/v1/to_product')
+
+    server = ls.LitServer([to_product])
+    server.run(port=8888)

app/server/generate_image/flux2_klein/client.py

@@ -9,5 +9,5 @@ request_data = {
     "infer_step":4,
     "tasks_id":"123456-123"
 }
-response = requests.post("http://127.0.0.1:8012//api/v1/to_product", json=request_data)
+response = requests.post("http://127.0.0.1:10090//api/v1/to_product", json=request_data)
 print(f"Status: {response.status_code}\nResponse:\n {response.text}")

app/server/generate_image/flux2_klein/server.py

@@ -1,4 +1,6 @@
 import io
+import os
+
 import torch
 import litserve as ls
 from diffusers import Flux2KleinPipeline
@@ -15,7 +17,8 @@ class Flux2KleinServer(ls.LitAPI):
         # Load the model
         dtype = torch.bfloat16
         self.device = device
-        self.model = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", torch_dtype=dtype, is_distilled=False)
+        model_path = os.path.join(settings.FLUX2_KLEIN_MODEL_PATH, "FLUX.2-klein-4B")
+        self.model = Flux2KleinPipeline.from_pretrained(model_path, torch_dtype=dtype, is_distilled=False)
         self.model.to(device)  # save some VRAM by offloading the model to CPU
 
     def decode_request(self, request):