first cimmit

app/server/generate_image/flux2_klein/client.py

@@ -0,0 +1,13 @@
+# This file is auto-generated by LitServe.
+# Disable auto-generation by setting `generate_client_file=False` in `LitServer.run()`.
+
+import requests
+
+request_data = {
+    "prompt": "Create realistic studio photo with real people model standing and wearing this garment, in white studio, Keep original model if present, or generate appropriate model, Standing pose, facing camera.",
+    "image_path":"aida-results/result_38151e0a-f83b-11f0-89f6-0242ac130002.png",
+    "infer_step":4,
+    "tasks_id":"123456-123"
+}
+response = requests.post("http://127.0.0.1:8012//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

@@ -0,0 +1,65 @@
+import io
+import torch
+import litserve as ls
+from diffusers import Flux2KleinPipeline
+from minio import Minio
+
+from app.config.config import settings
+from app.server.utils.minio_client import oss_get_image, oss_upload_image
+
+minio_client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ACCESS, secret_key=settings.MINIO_SECRET, secure=settings.MINIO_SECURE)
+
+
+class Flux2KleinServer(ls.LitAPI):
+    def setup(self, device):
+        # 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)
+        self.model.to(device)  # save some VRAM by offloading the model to CPU
+
+    def decode_request(self, request):
+        return request
+
+    def predict(self, request_data):
+        image_path = request_data.get("image_path", "")
+        prompt = request_data.get("prompt", "")
+        height = request_data.get("height", 768)
+        width = request_data.get("width", 512)
+        infer_step = request_data.get("infer_step", 4)
+        tasks_id = request_data.get("tasks_id", "test")
+
+        user_id = tasks_id[tasks_id.rfind('-') + 1:]
+        input_image = oss_get_image(oss_client=minio_client, path=image_path, data_type='pil')
+        with torch.no_grad():
+            images = self.model(
+                image=input_image,
+                prompt=prompt,
+                height=height,
+                width=width,
+                guidance_scale=1.0,
+                num_inference_steps=infer_step,
+                # generator=torch.Generator(device='cuda').manual_seed(3)
+            )[0]
+
+        # save image to minio
+        image = images[0]  # Assuming you want to retrieve the first image
+        image_data = io.BytesIO()
+        image.save(image_data, format='PNG')
+        image_data.seek(0)
+        image_bytes = image_data.read()
+        object_name = f'{user_id}/product_image/{tasks_id}.png'
+        req = oss_upload_image(oss_client=minio_client, bucket="aida-users", object_name=object_name, image_bytes=image_bytes)
+        image_url = f"aida-users/{object_name}"
+        return image_url
+
+    def encode_response(self, image_url):
+        return {"image_url": image_url}
+
+
+# Starting the server
+if __name__ == "__main__":
+    # Assume that an appropriate device (e.g., 'cuda', 'cpu') is specified
+    api = Flux2KleinServer()
+    server = ls.LitServer(api)
+    server.run(port=8011)

app/server/utils/minio_client.py

@@ -0,0 +1,97 @@
+import io
+import json
+import logging
+import time
+from io import BytesIO
+
+import cv2
+import numpy as np
+import urllib3
+from PIL import Image
+from minio import Minio
+
+from app.config.config import settings
+
+minio_client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ACCESS, secret_key=settings.MINIO_SECRET, secure=settings.MINIO_SECURE)
+
+
+# 自定义 Retry 类
+class CustomRetry(urllib3.Retry):
+    def increment(self, method=None, url=None, response=None, error=None, **kwargs):
+        # 调用父类的 increment 方法
+        new_retry = super(CustomRetry, self).increment(method, url, response, error, **kwargs)
+        # 打印重试信息
+        logger.info(f"重试连接: {method} {url}，错误: {error}，重试次数: {self.total - new_retry.total}")
+        return new_retry
+
+
+logger = logging.getLogger()
+timeout = urllib3.Timeout(connect=1, read=10.0)  # 连接超时 5 秒，读取超时 10 秒
+http_client = urllib3.PoolManager(
+    num_pools=10,  # 设置连接池大小
+    maxsize=10,
+    timeout=timeout,
+    cert_reqs='CERT_REQUIRED',  # 需要证书验证
+    retries=CustomRetry(
+        total=5,
+        backoff_factor=0.2,
+        status_forcelist=[500, 502, 503, 504],
+    ),
+)
+
+
+#  获取图片
+def oss_get_image(oss_client, path, data_type):
+    # cv2 默认全通道读取
+    bucket = path.split("/", 1)[0]
+    object_name = path.split("/", 1)[1]
+    image_object = None
+    try:
+        image_data = oss_client.get_object(bucket_name=bucket, object_name=object_name)
+        if data_type == "cv2":
+            image_bytes = image_data.read()
+            image_array = np.frombuffer(image_bytes, np.uint8)  # 转成8位无符号整型
+            image_object = cv2.imdecode(image_array, cv2.IMREAD_UNCHANGED)
+            if image_object.dtype == np.uint16:
+                image_object = (image_object / 256).astype('uint8')
+        else:
+            data_bytes = BytesIO(image_data.read())
+            image_object = Image.open(data_bytes)
+    except Exception as e:
+        logger.warning(f"获取图片出现异常 ######: {e}")
+    return image_object
+
+
+def oss_upload_image(oss_client, bucket, object_name, image_bytes):
+    req = None
+    try:
+        req = oss_client.put_object(bucket_name=bucket, object_name=object_name, data=io.BytesIO(image_bytes), length=len(image_bytes), content_type='image/png')
+    except Exception as e:
+        logger.warning(f"上传图片出现异常 ######: {e}")
+    return req
+
+
+if __name__ == '__main__':
+    # url = "aida-results/result_0002186a-e631-11ee-86a6-b48351119060.png"
+    # url = "aida-collection-element/11523/Moodboard/f60af0d2-94c2-48f9-90ff-74b8e8a481b5.jpg"
+    # url = "test/test.png"
+    # url = "aida-users/89/product_image/string-89.png"
+    # url = "test/845046c7-4f62-4f54-a4a9-c26d49c6969335b5b3a9-d335-4871-a46c-3cc3caf07da259629dfd1f1f555a2e2a9def7e719366.png"
+    # url = 'aida-users/89/relight_image/123-89.png'
+    # url = 'aida-users/89/relight_image/123-89.png'
+    # url = 'aida-users/89/relight_image/123-89.png'
+    # url = "aida-users/89/sketchboard/female/Dress/e6724ab7-8d3f-4677-abe0-c3e42ab7af85.jpeg"
+    # url = "aida-users/87/print/956614a2-7e75-4fbe-9ed0-c1831e37a2c9-4-87.png"
+    # url = "aida-users/89/single_logo/123-89.png"
+    # url = "aida-users/89/sketch/4e8fe37d-7068-400a-ac94-c01647fa5f6f.png"
+
+    url = "aida-users/123/product_image/123456-123.png"
+    read_type = "2"
+    if read_type == "cv2":
+        img = oss_get_image(oss_client=minio_client, path=url, data_type=read_type)
+        cv2.imshow("", img)
+        cv2.waitKey(0)
+    else:
+        img = oss_get_image(oss_client=minio_client, path=url, data_type=read_type)
+        img.show()
+        img.save("result.png")