Merge remote-tracking branch 'refs/remotes/origin/develop'

app/service/generate_image/service_generate_image.py

@@ -10,21 +10,19 @@
 import json
 import logging
 import time
-from io import BytesIO
 
 import cv2
 import minio
+import numpy as np
 import redis
 import tritonclient.grpc as grpcclient
-import numpy as np
-from minio import Minio
 from tritonclient.utils import np_to_triton_dtype
 
 from app.core.config import *
 from app.schemas.generate_image import GenerateImageModel
-from app.service.generate_image.utils.adjust_contrast import adjust_contrast
-from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust, face_detect_pic
-from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_stain_png_sd
+from app.service.generate_image.utils.image_processing import remove_background, stain_detection, generate_category_recognition, autoLevels, luminance_adjust
+from app.service.generate_image.utils.upload_sd_image import upload_png_sd
+from app.service.utils.oss_client import oss_get_image
 
 logger = logging.getLogger()
 
@@ -36,22 +34,23 @@ class GenerateImage:
             self.channel = self.connection.channel()
         # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
         # self.channel = self.connection.channel()
-        self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
+        # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
         self.grpc_client = grpcclient.InferenceServerClient(url=GI_MODEL_URL)
         self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
         if request_data.mode == "img2img":
             # cv2 读图片是BGR PIL读图片是RGB
             self.image = self.get_image(request_data.image_url)
-            self.prompt = request_data.prompt
         else:
             self.image = np.random.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)
-            self.prompt = request_data.prompt
 
+        self.prompt = request_data.prompt
         self.tasks_id = request_data.tasks_id
         self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
         self.mode = request_data.mode
         self.batch_size = 1
         self.category = request_data.category
+        if self.category == "sketch":
+            self.prompt = f"{self.category},{self.prompt}"
         self.index = 0
         self.gender = request_data.gender
         self.generate_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': '', 'category': ''}
@@ -63,10 +62,13 @@ class GenerateImage:
         # Read data from response.
         # read image use cv2
         try:
-            response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
-            image_file = BytesIO(response.data)
-            image_array = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
-            image_cv2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
+            # response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
+            # image_file = BytesIO(response.data)
+            # image_array = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
+            # image_cv2 = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
+            # image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
+
+            image_cv2 = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="cv2")
             image_rbg = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
             image = cv2.resize(image_rbg, (1024, 1024))
         except minio.error.S3Error:
@@ -104,7 +106,7 @@ class GenerateImage:
                 image_result = not_smudge_image
             if is_smudge:  # 无污点
                 # image_result = adjust_contrast(image_result)
-                image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", object_name=f"{self.tasks_id}.png")
+                image_url = upload_png_sd(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
                 # logger.info(f"upload image SUCCESS ： {image_url}")
                 self.generate_data['status'] = "SUCCESS"
                 self.generate_data['message'] = "success"
@@ -121,13 +123,6 @@ class GenerateImage:
         status_data = self.redis_client.get(self.tasks_id)
         return json.loads(status_data), status_data
 
-    def infer(self, inputs):
-        return self.grpc_client.async_infer(
-            model_name=GI_MODEL_NAME,
-            inputs=inputs,
-            callback=self.callback
-        )
-
     def get_result(self):
         try:
             prompts = [self.prompt] * self.batch_size
@@ -147,7 +142,7 @@ class GenerateImage:
             input_mode.set_data_from_numpy(mode_obj)
 
             inputs = [input_text, input_image, input_mode]
-            ctx = self.infer(inputs)
+            ctx = self.grpc_client.async_infer(model_name=GI_MODEL_NAME, inputs=inputs, callback=self.callback)
             time_out = 600
             generate_data = None
             while time_out > 0:
@@ -187,9 +182,10 @@ if __name__ == '__main__':
     rd = GenerateImageModel(
         tasks_id="123-89",
         prompt='skeleton sitting by the side of a river looking soulful, concert poster, 4k, artistic',
-        image_url="",
+        image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg",
         mode='txt2img',
-        category="test"
+        category="test",
+        gender="male"
     )
     server = GenerateImage(rd)
     print(server.get_result())

app/service/generate_image/service_generate_product_image.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+"""
+@Project ：trinity_client
+@File    ：service_att_recognition.py
+@Author  ：周成融
+@Date    ：2023/7/26 12:01:05
+@detail  ：
+"""
+import json
+import logging
+import time
+
+import cv2
+import numpy as np
+import redis
+import tritonclient.grpc as grpcclient
+from PIL import Image, ImageOps
+from tritonclient.utils import np_to_triton_dtype
+
+from app.core.config import *
+from app.schemas.generate_image import GenerateProductImageModel
+from app.service.generate_image.utils.upload_sd_image import upload_SDXL_image
+from app.service.utils.oss_client import oss_get_image
+
+logger = logging.getLogger()
+
+
+class GenerateProductImage:
+    def __init__(self, request_data):
+        if DEBUG is False:
+            self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+            self.channel = self.connection.channel()
+        # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+        # self.channel = self.connection.channel()
+        # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
+        self.grpc_client = grpcclient.InferenceServerClient(url=GPI_MODEL_URL)
+        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+        self.category = "product_image"
+        self.image_strength = request_data.image_strength
+        self.batch_size = 1
+        self.product_type = request_data.product_type
+        self.prompt = request_data.prompt
+        self.image, self.image_size = pre_processing_image(request_data.image_url)
+        self.tasks_id = request_data.tasks_id
+        self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
+        self.gen_product_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
+        self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+        self.redis_client.expire(self.tasks_id, 600)
+
+    def callback(self, result, error):
+        if error:
+            self.gen_product_data['status'] = "FAILURE"
+            self.gen_product_data['message'] = str(error)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+        else:
+            # pil图像转成numpy数组
+            if self.product_type == "single":
+                image = result.as_numpy("generated_cnet_image")
+            else:
+                image = result.as_numpy("generated_inpaint_image")
+            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8))).resize(self.image_size)
+            image_url = upload_SDXL_image(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
+            self.gen_product_data['status'] = "SUCCESS"
+            self.gen_product_data['message'] = "success"
+            self.gen_product_data['image_url'] = str(image_url)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+
+    def read_tasks_status(self):
+        status_data = self.redis_client.get(self.tasks_id)
+        return json.loads(status_data), status_data
+
+    def get_result(self):
+        try:
+            prompts = [self.prompt] * self.batch_size
+            self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
+            self.image = cv2.resize(self.image, (512, 768))
+            images = [self.image.astype(np.uint8)] * self.batch_size
+
+            if self.product_type == "single":
+                text_obj = np.array(prompts, dtype="object").reshape(-1, 1)
+                image_obj = np.array(images, dtype=np.uint8).reshape((-1, 768, 512, 3))
+                image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape(-1, 1)
+            else:
+                text_obj = np.array(prompts, dtype="object").reshape(1)
+                image_obj = np.array(images, dtype=np.uint8).reshape((768, 512, 3))
+                image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1))
+
+            # 假设 prompts、images 和 self.image_strength 已经定义
+
+            input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
+            input_image = grpcclient.InferInput("input_image", image_obj.shape, "UINT8")
+            input_image_strength = grpcclient.InferInput("image_strength", image_strength_obj.shape, np_to_triton_dtype(image_strength_obj.dtype))
+
+            input_text.set_data_from_numpy(text_obj)
+            input_image.set_data_from_numpy(image_obj)
+            inputs = [input_text, input_image, input_image_strength]
+            input_image_strength.set_data_from_numpy(image_strength_obj)
+
+            if self.product_type == "single":
+                ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_SINGLE, inputs=inputs, callback=self.callback)
+            else:
+                ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback)
+
+            time_out = 600
+            while time_out > 0:
+                gen_product_data, _ = self.read_tasks_status()
+                if gen_product_data['status'] in ["REVOKED", "FAILURE"]:
+                    ctx.cancel()
+                    break
+                elif gen_product_data['status'] == "SUCCESS":
+                    break
+                time_out -= 1
+                time.sleep(0.1)
+            gen_product_data, _ = self.read_tasks_status()
+            return gen_product_data
+        except Exception as e:
+            self.gen_product_data['status'] = "FAILURE"
+            self.gen_product_data['message'] = str(e)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+            raise Exception(str(e))
+        finally:
+            dict_gen_product_data, str_gen_product_data = self.read_tasks_status()
+            if DEBUG is False:
+                self.channel.basic_publish(exchange='', routing_key=GPI_RABBITMQ_QUEUES, body=str_gen_product_data)
+            logger.info(f" [x] Sent to： {GPI_RABBITMQ_QUEUES} data：@@@@  {json.dumps(dict_gen_product_data, indent=4)}")
+
+
+def infer_cancel(tasks_id):
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+    data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
+    gen_product_data = json.dumps(data)
+    redis_client.set(tasks_id, gen_product_data)
+    return data
+
+
+def pre_processing_image(image_url):
+    image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="PIL")
+    # 原始图片的尺寸
+    width, height = image.size
+
+    # 计算长宽比为 3:2 的新尺寸
+    desired_ratio = 2 / 3
+    current_ratio = width / height
+
+    if current_ratio > desired_ratio:
+        # 原始图片更宽，需要在上下添加 padding
+        new_width = width
+        new_height = int(width / desired_ratio)
+    else:
+        # 原始图片更高或者长宽比已经为 3:2
+        new_height = height
+        new_width = int(height * desired_ratio)
+
+    # 创建一个新的画布，大小为添加 padding 后的尺寸，并设置为白色背景
+    pad_image = Image.new('RGBA', (new_width, new_height), (0, 0, 0, 0))
+
+    # 将原始图片粘贴到新的画布中心
+    left = (new_width - width) // 2
+    top = (new_height - height) // 2
+    pad_image.paste(image, (left, top))
+
+    # 将画布 resize 成宽度 500，长度 750
+    resized_image = pad_image.resize((500, 750))
+    image_size = (512, 768)
+
+    if resized_image.mode in ('RGBA', 'LA') or (resized_image.mode == 'P' and 'transparency' in resized_image.info):
+        # 创建白色背景
+        background = Image.new("RGB", image_size, (255, 255, 255))
+        # 将图片粘贴到白色背景上
+        background.paste(resized_image, mask=resized_image.split()[3])
+        image = np.array(background)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return image, image_size
+
+
+if __name__ == '__main__':
+    rd = GenerateProductImageModel(
+        tasks_id="123-89",
+        # prompt="",
+        image_strength=0.9,
+        prompt=" the best quality, masterpiece. detailed, high-res, simple background, studio photography, extremely detailed, updo, detailed face, face, close-up, HDR, UHD, 8K realistic, Highly detailed, simple background, Studio lighting",
+        image_url="aida-results/result_00097282-ebb2-11ee-a822-b48351119060.png",
+        product_type="overall"
+    )
+    server = GenerateProductImage(rd)
+    print(server.get_result())

app/service/generate_image/service_generate_relight_image.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+"""
+@Project ：trinity_client
+@File    ：service_att_recognition.py
+@Author  ：周成融
+@Date    ：2023/7/26 12:01:05
+@detail  ：
+"""
+import json
+import logging
+import time
+
+import cv2
+import numpy as np
+import redis
+import tritonclient.grpc as grpcclient
+from PIL import Image
+from tritonclient.utils import np_to_triton_dtype
+
+from app.core.config import *
+from app.schemas.generate_image import GenerateRelightImageModel
+from app.service.generate_image.utils.upload_sd_image import upload_SDXL_image
+from app.service.utils.oss_client import oss_get_image
+
+logger = logging.getLogger()
+
+
+class GenerateRelightImage:
+    def __init__(self, request_data):
+        if DEBUG is False:
+            self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+            self.channel = self.connection.channel()
+        # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
+        self.grpc_client = grpcclient.InferenceServerClient(url=GRI_MODEL_URL)
+        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+        self.category = "relight_image"
+        self.batch_size = 1
+        self.prompt = request_data.prompt
+        self.seed = "1"
+        self.product_type = request_data.product_type
+        self.negative_prompt = 'lowres, bad anatomy, bad hands, cropped, worst quality'
+        self.direction = request_data.direction
+        self.image_url = request_data.image_url
+        self.image = oss_get_image(bucket=self.image_url.split('/')[0], object_name=self.image_url[self.image_url.find('/') + 1:], data_type="cv2")
+        self.tasks_id = request_data.tasks_id
+        self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
+        self.gen_product_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
+        self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+        self.redis_client.expire(self.tasks_id, 600)
+
+    def callback(self, result, error):
+        if error:
+            self.gen_product_data['status'] = "FAILURE"
+            self.gen_product_data['message'] = str(error)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+        else:
+            # pil图像转成numpy数组
+            if self.product_type == 'single':
+                image = result.as_numpy("generated_relight_image")
+            else:
+                image = result.as_numpy("generated_inpaint_image")
+
+            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
+
+            image_url = upload_SDXL_image(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
+            self.gen_product_data['status'] = "SUCCESS"
+            self.gen_product_data['message'] = "success"
+            self.gen_product_data['image_url'] = str(image_url)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+
+    def read_tasks_status(self):
+        status_data = self.redis_client.get(self.tasks_id)
+        return json.loads(status_data), status_data
+
+    def get_result(self):
+        try:
+            prompts = [self.prompt] * self.batch_size
+            image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
+            image = cv2.resize(image, (512, 768))
+            images = [image.astype(np.uint8)] * self.batch_size
+            seeds = [self.seed] * self.batch_size
+            nagetive_prompts = [self.negative_prompt] * self.batch_size
+            directions = [self.direction] * self.batch_size
+
+            if self.product_type == 'single':
+                text_obj = np.array(prompts, dtype="object").reshape((-1, 1))
+                image_obj = np.array(images, dtype=np.uint8).reshape((-1, 768, 512, 3))
+                na_text_obj = np.array(nagetive_prompts, dtype="object").reshape((-1, 1))
+                seed_obj = np.array(seeds, dtype="object").reshape((-1, 1))
+                direction_obj = np.array(directions, dtype="object").reshape((-1, 1))
+            else:
+                text_obj = np.array(prompts, dtype="object").reshape((1))
+                image_obj = np.array(images, dtype=np.uint8).reshape((768, 512, 3))
+                na_text_obj = np.array(nagetive_prompts, dtype="object").reshape((1))
+                seed_obj = np.array(seeds, dtype="object").reshape((1))
+                direction_obj = np.array(directions, dtype="object").reshape((1))
+
+            input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
+            input_image = grpcclient.InferInput("input_image", image_obj.shape, "UINT8")
+            input_natext = grpcclient.InferInput("negative_prompt", na_text_obj.shape, np_to_triton_dtype(na_text_obj.dtype))
+            input_seed = grpcclient.InferInput("seed", seed_obj.shape, np_to_triton_dtype(seed_obj.dtype))
+            input_direction = grpcclient.InferInput("direction", direction_obj.shape, np_to_triton_dtype(direction_obj.dtype))
+
+            input_text.set_data_from_numpy(text_obj)
+            input_image.set_data_from_numpy(image_obj)
+            input_natext.set_data_from_numpy(na_text_obj)
+            input_seed.set_data_from_numpy(seed_obj)
+            input_direction.set_data_from_numpy(direction_obj)
+
+            inputs = [input_text, input_natext, input_image, input_seed, input_direction]
+            if self.product_type == 'single':
+                ctx = self.grpc_client.async_infer(model_name=GRI_MODEL_NAME_SINGLE, inputs=inputs, callback=self.callback)
+            else:
+                ctx = self.grpc_client.async_infer(model_name=GRI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback)
+
+            time_out = 600
+            while time_out > 0:
+                gen_product_data, _ = self.read_tasks_status()
+                if gen_product_data['status'] in ["REVOKED", "FAILURE"]:
+                    ctx.cancel()
+                    break
+                elif gen_product_data['status'] == "SUCCESS":
+                    break
+                time_out -= 1
+                time.sleep(0.1)
+            gen_product_data, _ = self.read_tasks_status()
+            return gen_product_data
+        except Exception as e:
+            self.gen_product_data['status'] = "FAILURE"
+            self.gen_product_data['message'] = str(e)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
+            raise Exception(str(e))
+        finally:
+            dict_gen_product_data, str_gen_product_data = self.read_tasks_status()
+            if DEBUG is False:
+                self.channel.basic_publish(exchange='', routing_key=GRI_RABBITMQ_QUEUES, body=str_gen_product_data)
+            logger.info(f" [x] Sent to： {GRI_RABBITMQ_QUEUES} data：@@@@  {json.dumps(dict_gen_product_data, indent=4)}")
+
+
+def infer_cancel(tasks_id):
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+    data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
+    gen_product_data = json.dumps(data)
+    redis_client.set(tasks_id, gen_product_data)
+    return data
+
+
+if __name__ == '__main__':
+    rd = GenerateRelightImageModel(
+        tasks_id="123-89",
+        # prompt="beautiful woman, detailed face, sunshine, outdoor, warm atmosphere",
+        prompt="Colorful black",
+        image_url='aida-results/result_0000b606-1902-11ef-9424-0242ac180002.png',
+        direction="Right Light",
+        product_type="single"
+    )
+    server = GenerateRelightImage(rd)
+    print(server.get_result())

app/service/generate_image/service_generate_single_logo.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+"""
+@Project ：trinity_client
+@File    ：service_att_recognition.py
+@Author  ：周成融
+@Date    ：2023/7/26 12:01:05
+@detail  ：
+"""
+import json
+import logging
+import time
+
+import cv2
+import numpy as np
+import redis
+from PIL import Image
+from minio import Minio
+from tritonclient.utils import np_to_triton_dtype
+
+from app.core.config import *
+import tritonclient.grpc as grpcclient
+from app.schemas.generate_image import GenerateSingleLogoImageModel
+from app.service.generate_image.utils.upload_sd_image import upload_png_sd, upload_SDXL_image
+
+logger = logging.getLogger()
+
+
+class GenerateSingleLogoImage:
+    def __init__(self, request_data):
+        if DEBUG is False:
+            self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
+            self.channel = self.connection.channel()
+        self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
+        self.grpc_client = grpcclient.InferenceServerClient(url=GSL_MODEL_URL)
+        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+        self.batch_size = 1
+        self.category = "single_logo"
+        self.negative_prompts = "bad, ugly"
+        self.seed = request_data.seed
+        self.tasks_id = request_data.tasks_id
+        self.prompt = request_data.prompt
+        self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
+        self.gen_single_logo_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
+        self.redis_client.set(self.tasks_id, json.dumps(self.gen_single_logo_data))
+        self.redis_client.expire(self.tasks_id, 600)
+
+    def read_tasks_status(self):
+        status_data = self.redis_client.get(self.tasks_id)
+        return json.loads(status_data), status_data
+
+    def callback(self, result, error):
+        if error:
+            self.gen_single_logo_data['status'] = "FAILURE"
+            self.gen_single_logo_data['message'] = str(error)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_single_logo_data))
+        else:
+            image = result.as_numpy("generated_image")
+            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
+            image_url = upload_SDXL_image(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
+            self.gen_single_logo_data['status'] = "SUCCESS"
+            self.gen_single_logo_data['message'] = "success"
+            self.gen_single_logo_data['image_url'] = str(image_url)
+            self.redis_client.set(self.tasks_id, json.dumps(self.gen_single_logo_data))
+
+    def get_result(self):
+        try:
+            # prompt
+            prompts = [self.prompt] * self.batch_size
+            text_obj = np.array(prompts, dtype="object").reshape((-1, 1))
+            input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
+            input_text.set_data_from_numpy(text_obj)
+
+            text_obj_neg = np.array(self.negative_prompts, dtype="object").reshape((-1, 1))
+            input_text_neg = grpcclient.InferInput("negative_prompt", text_obj_neg.shape, np_to_triton_dtype(text_obj_neg.dtype))
+            input_text_neg.set_data_from_numpy(text_obj_neg)
+
+            seed = np.array(self.seed, dtype="object").reshape((-1, 1))
+            input_seed = grpcclient.InferInput("seed", seed.shape, np_to_triton_dtype(seed.dtype))
+            input_seed.set_data_from_numpy(seed)
+            inputs = [input_text, input_text_neg, input_seed]
+            ctx = self.grpc_client.async_infer(model_name=GSL_MODEL_NAME, inputs=inputs, callback=self.callback)
+            time_out = 600
+            generate_data = None
+            while time_out > 0:
+                generate_data, _ = self.read_tasks_status()
+                if generate_data['status'] in ["REVOKED", "FAILURE"]:
+                    ctx.cancel()
+                    break
+                elif generate_data['status'] == "SUCCESS":
+                    break
+                time_out -= 1
+                time.sleep(0.1)
+            return generate_data
+        except Exception as e:
+            raise Exception(str(e))
+        finally:
+            dict_generate_data, str_generate_data = self.read_tasks_status()
+            if DEBUG is False:
+                self.channel.basic_publish(exchange='', routing_key=GI_RABBITMQ_QUEUES, body=str_generate_data)
+            logger.info(f" [x] Sent {json.dumps(dict_generate_data, indent=4)}")
+
+
+def infer_cancel(tasks_id):
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
+    data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
+    generate_data = json.dumps(data)
+    redis_client.set(tasks_id, generate_data)
+    return data
+
+
+if __name__ == '__main__':
+    rd = GenerateSingleLogoImageModel(
+        tasks_id="123-89",
+        prompt='an apple',
+        seed="2",
+    )
+    server = GenerateSingleLogoImage(rd)
+    print(server.get_result())

app/service/generate_image/utils/image_processing.py

@@ -381,7 +381,7 @@ if __name__ == '__main__':
     remove_bg_img = remove_background(luminance)
     # cv2.imwrite("remove_bg_img.png", remove_bg_img)
 
-    print(1)
+    # print(1)
     cv2.imshow("source", img)
     cv2.imshow("levels", equAuto)
     cv2.imshow("luminance", luminance)

app/service/generate_image/utils/upload_sd_image.py

@@ -10,26 +10,61 @@
 import io
 import logging
 
+# import boto3
 import cv2
 from PIL import Image
 from minio import Minio
 
 from app.core.config import *
+from app.service.utils.oss_client import oss_upload_image
 
 minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
 
 
-def upload_png_sd(image, user_id, category, object_name):
+# s3 = boto3.client('s3', aws_access_key_id=S3_ACCESS_KEY, aws_secret_access_key=S3_AWS_SECRET_ACCESS_KEY, region_name=S3_REGION_NAME)
+
+
+# def upload_single_logo(image, user_id, category, object_name):
+#     with io.BytesIO() as output:
+#         image.save(output, format='PNG')
+#         data = output.getvalue()
+#     # 创建一个 S3 客户端
+#     try:
+#         key = f'{user_id}/{category}/{object_name}'
+#         image_url = f"{AIDA_CLOTHING}/{key}"
+#         s3.put_object(Bucket=GSL_MINIO_BUCKET, Key=key, Body=data, ContentType='image/png')
+#         return image_url
+#     except Exception as e:
+#         print(f'上传到 S3 失败: {e}')
+
+def upload_SDXL_image(image, user_id, category, file_name):
+    try:
+        image_data = io.BytesIO()
+        image.save(image_data, format='PNG')
+        image_data.seek(0)
+        image_bytes = image_data.read()
+
+        # minio_req = minio_client.put_object(
+        #     GI_MINIO_BUCKET,
+        #     f'{user_id}/{category}/{file_name}',
+        #     io.BytesIO(image_bytes),
+        #     len(image_bytes),
+        #     content_type='image/jpeg'
+        # )
+        object_name = f'{user_id}/{category}/{file_name}'
+        req = oss_upload_image(bucket=GI_MINIO_BUCKET, object_name=object_name, image_bytes=image_bytes)
+        image_url = f"aida-users/{object_name}"
+        return image_url
+    except Exception as e:
+        logging.warning(f"upload_png_mask runtime exception : {e}")
+
+
+def upload_png_sd(image, user_id, category, file_name):
     try:
         _, img_byte_array = cv2.imencode('.jpg', image)
-        minio_req = minio_client.put_object(
-            GI_MINIO_BUCKET,
-            f'{user_id}/{category}/{object_name}',
-            io.BytesIO(img_byte_array),
-            len(img_byte_array),
-            content_type='image/jpeg'
-        )
-        image_url = f"aida-users/{minio_req.object_name}"
+        object_name = f'{user_id}/{category}/{file_name}'
+        req = oss_upload_image(bucket=GI_MINIO_BUCKET, object_name=object_name, image_bytes=img_byte_array)
+        image_url = f"aida-users/{object_name}"
         return image_url
     except Exception as e:
         logging.warning(f"upload_png_mask runtime exception : {e}")