feat 产品图打光模型部署

fix

app/service/generate_image/service_generate_relight_image.py

@@ -22,6 +22,7 @@ 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()
 
@@ -31,43 +32,34 @@ class GenerateRelightImage:
         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.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 = "12345"
+        self.seed = "1"
         self.negative_prompt = 'lowres, bad anatomy, bad hands, cropped, worst quality'
         self.direction = "Right Light"
-        # TODO aida design 结果图背景改为白色
-        self.image = self.get_image(request_data.image_url)
-        # TODO image 填充并resize成512*768
-
+        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 get_image(self, image_url):
-        response = self.minio_client.get_object(image_url.split('/')[0], image_url[image_url.find('/') + 1:])
-        image = cv2.imdecode(np.frombuffer(response.data, np.uint8), 1)
-        return image
-
     def callback(self, result, error):
         if error:
             self.gen_product_data['status'] = "FAILURE"
             self.gen_product_data['message'] = str(error)
-            # self.gen_product_data['data'] = str(error)
             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
         else:
             # pil图像转成numpy数组
             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}", object_name=f"{self.tasks_id}.png")
-            # logger.info(f"upload image SUCCESS ： {image_url}")
+            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)
@@ -77,13 +69,6 @@ class GenerateRelightImage:
         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=GRI_MODEL_NAME,
-            inputs=inputs,
-            callback=self.callback
-        )
-
     def get_result(self):
         try:
             prompts = [self.prompt] * self.batch_size
@@ -114,11 +99,10 @@ class GenerateRelightImage:
 
             inputs = [input_text, input_natext, input_image, input_seed, input_direction]
 
-            ctx = self.infer(inputs)
+            ctx = self.grpc_client.async_infer(model_name=GRI_MODEL_NAME, inputs=inputs, callback=self.callback)
             time_out = 600
             while time_out > 0:
                 gen_product_data, _ = self.read_tasks_status()
-                # logger.info(gen_product_data)
                 if gen_product_data['status'] in ["REVOKED", "FAILURE"]:
                     ctx.cancel()
                     break
@@ -126,7 +110,6 @@ class GenerateRelightImage:
                     break
                 time_out -= 1
                 time.sleep(0.1)
-                # logger.info(time_out, gen_product_data)
             gen_product_data, _ = self.read_tasks_status()
             return gen_product_data
         except Exception as e:
@@ -138,7 +121,6 @@ class GenerateRelightImage:
             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)
-            # self.channel.basic_publish(exchange='', routing_key=GEN_PRODUCT_IMAGE_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)}")
 
 
@@ -154,7 +136,7 @@ if __name__ == '__main__':
     rd = GenerateRelightImageModel(
         tasks_id="123-89",
         # prompt="beautiful woman, detailed face, sunshine, outdoor, warm atmosphere",
-        prompt="",
+        prompt="Colorful black",
         image_url='aida-users/89/product_image/123-89.png'
     )
     server = GenerateRelightImage(rd)