From 2121f0ab46b937b66fdc787766860e2bdcac74f0 Mon Sep 17 00:00:00 2001 From: zhouchengrong Date: Mon, 3 Jun 2024 11:38:20 +0800 Subject: [PATCH] =?UTF-8?q?feat=20generate=20single=20logo=20=E9=83=A8?= =?UTF-8?q?=E7=BD=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/core/config.py | 8 + app/schemas/generate_image.py | 2 +- .../design/items/pipelines/painting.py | 8 +- app/service/design/utils/synthesis_item.py | 6 +- app/service/design/utils/upload_image.py | 218 +++++++++--------- .../service_generate_single_logo.py | 98 ++++++-- .../generate_image/utils/upload_sd_image.py | 34 +++ 7 files changed, 235 insertions(+), 139 deletions(-) diff --git a/app/core/config.py b/app/core/config.py index 08802a5..a5bc957 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -41,6 +41,11 @@ MINIO_ACCESS = 'vXKFLSJkYeEq2DrSZvkB' MINIO_SECRET = 'uKTZT3x7C43WvPN9QTc99DiRkwddWZrG9Uh3JVlR' MINIO_SECURE = True +# S3 配置 +S3_ACCESS_KEY = "AKIAVD3OJIMF6UJFLSHZ" +S3_AWS_SECRET_ACCESS_KEY = "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" +S3_REGION_NAME = "ap-east-1" + # redis 配置 REDIS_HOST = "10.1.1.240" REDIS_PORT = "6379" @@ -107,6 +112,9 @@ GI_SYS_IMAGE_URL = "aida-sys-image/generate_image/white_image.jpg" SLOGAN_RABBITMQ_QUEUES = os.getenv("SLOGAN_RABBITMQ_QUEUES", f"Slogan{RABBITMQ_ENV}") # Generate Single Logo service config +GSL_MODEL_URL = '10.1.1.240:10051' +GSL_MINIO_BUCKET = "aida-users" +GSL_MODEL_NAME = 'stable_diffusion_xl' GEN_SINGLE_LOGO_RABBITMQ_QUEUES = os.getenv("GEN_SINGLE_LOGO_RABBITMQ_QUEUES", f"GenSingleLogo{RABBITMQ_ENV}") # Generate Single Logo service config diff --git a/app/schemas/generate_image.py b/app/schemas/generate_image.py index b30e64e..49cf9ce 100644 --- a/app/schemas/generate_image.py +++ b/app/schemas/generate_image.py @@ -13,7 +13,7 @@ class GenerateImageModel(BaseModel): class GenerateSingleLogoImageModel(BaseModel): tasks_id: str prompt: str - image_url: str + seed: str class GenerateProductImageModel(BaseModel): diff --git a/app/service/design/items/pipelines/painting.py b/app/service/design/items/pipelines/painting.py index 0b48082..d1f6957 100644 --- a/app/service/design/items/pipelines/painting.py +++ b/app/service/design/items/pipelines/painting.py @@ -15,13 +15,7 @@ minio_client = Minio( secret_key=MINIO_SECRET, secure=MINIO_SECURE) - -# s3 = boto3.client( -# 's3', -# aws_access_key_id="AKIAVD3OJIMF6UJFLSHZ", -# aws_secret_access_key="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", -# region_name="ap-east-1" -# ) +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) @PIPELINES.register_module() diff --git a/app/service/design/utils/synthesis_item.py b/app/service/design/utils/synthesis_item.py index 91505bd..0cf844b 100644 --- a/app/service/design/utils/synthesis_item.py +++ b/app/service/design/utils/synthesis_item.py @@ -29,9 +29,9 @@ minio_client = Minio( s3 = boto3.client( 's3', - aws_access_key_id="AKIAVD3OJIMF6UJFLSHZ", - aws_secret_access_key="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", - region_name="ap-east-1" + aws_access_key_id=S3_ACCESS_KEY, + aws_secret_access_key=S3_AWS_SECRET_ACCESS_KEY, + region_name=S3_REGION_NAME ) diff --git a/app/service/design/utils/upload_image.py b/app/service/design/utils/upload_image.py index 7503adc..70b259c 100644 --- a/app/service/design/utils/upload_image.py +++ b/app/service/design/utils/upload_image.py @@ -25,137 +25,131 @@ minio_client = Minio( secure=MINIO_SECURE) """S3 上传""" -s3 = boto3.client( - 's3', - aws_access_key_id="AKIAVD3OJIMF6UJFLSHZ", - aws_secret_access_key="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", - region_name="ap-east-1" -) +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) -# @RunTime -# def upload_png_mask(front_image, object_name, mask=None): -# start_time = time.time() -# mask_url = None -# if mask is not None: -# # 反转掩模 -# mask_inverted = cv2.bitwise_not(mask) -# # 将掩模转换为 RGBA 格式 -# rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) -# rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] -# # 将图像数据保存到内存中的 BytesIO 对象中 -# image_bytes = io.BytesIO() -# image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) -# image_bytes.seek(0) -# try: -# key = f"mask/mask_{object_name}.png" -# mask_url = f"{AIDA_CLOTHING}/{key}" -# s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=image_bytes, ContentType='image/png') -# except Exception as e: -# print(f'上传到 S3 失败: {e}') -# with io.BytesIO() as output: -# front_image.save(output, format='PNG') -# data = output.getvalue() -# # 创建一个 S3 客户端 -# try: -# key = f"image/image_{object_name}.png" -# image_url = f"{AIDA_CLOTHING}/{key}" -# s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=data, ContentType='image/png') -# return front_image, image_url, mask_url -# except Exception as e: -# print(f'上传到 S3 失败: {e}') -# -# -# @RunTime -# def upload_layer_image(image, object_name): -# with io.BytesIO() as output: -# image.save(output, format='PNG') -# data = output.getvalue() -# # 创建一个 S3 客户端 -# try: -# key = f"image/image_{object_name}.png" -# image_url = f"{AIDA_CLOTHING}/{key}" -# s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=data, ContentType='image/png') -# return image_url -# except Exception as e: -# print(f'上传到 S3 失败: {e}') -# -# -# @RunTime -# def upload_mask_image(mask, object_name): -# # 反转掩模 -# mask_inverted = cv2.bitwise_not(mask) -# # 将掩模转换为 RGBA 格式 -# rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) -# rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] -# # 将图像数据保存到内存中的 BytesIO 对象中 -# image_bytes = io.BytesIO() -# image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) -# image_bytes.seek(0) -# try: -# key = f"mask/mask_{object_name}.png" -# mask_url = f"{AIDA_CLOTHING}/{key}" -# s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=image_bytes, ContentType='image/png') -# return mask_url -# except Exception as e: -# print(f'上传到 S3 失败: {e}') - - -"""minio 上传""" - @RunTime def upload_png_mask(front_image, object_name, mask=None): - start_time = time.time() + mask_url = None + if mask is not None: + # 反转掩模 + mask_inverted = cv2.bitwise_not(mask) + # 将掩模转换为 RGBA 格式 + rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) + rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] + # 将图像数据保存到内存中的 BytesIO 对象中 + image_bytes = io.BytesIO() + image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) + image_bytes.seek(0) + try: + key = f"mask/mask_{object_name}.png" + mask_url = f"{AIDA_CLOTHING}/{key}" + s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=image_bytes, ContentType='image/png') + except Exception as e: + print(f'上传到 S3 失败: {e}') + with io.BytesIO() as output: + front_image.save(output, format='PNG') + data = output.getvalue() + # 创建一个 S3 客户端 try: - mask_url = None - if mask is not None: - mask_inverted = cv2.bitwise_not(mask) - # 将掩模的3通道转换为4通道，白色部分不透明，黑色部分透明 - rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) - rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] - image_bytes = io.BytesIO() - image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) - - image_bytes.seek(0) - mask_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'mask/mask_{object_name}.png', image_bytes, len(image_bytes.getvalue()), content_type='image/png').object_name}" - - image_data = io.BytesIO() - front_image.save(image_data, format='PNG') - image_data.seek(0) - image_bytes = image_data.read() - image_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'image/image_{object_name}.png', io.BytesIO(image_bytes), len(image_bytes), content_type='image/png').object_name}" - # print(f"upload_png_mask {object_name} = {time.time() - start_time}") + key = f"image/image_{object_name}.png" + image_url = f"{AIDA_CLOTHING}/{key}" + s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=data, ContentType='image/png') return front_image, image_url, mask_url except Exception as e: - logging.warning(f"upload_png_mask runtime exception : {e}") + print(f'上传到 S3 失败: {e}') @RunTime def upload_layer_image(image, object_name): + with io.BytesIO() as output: + image.save(output, format='PNG') + data = output.getvalue() + # 创建一个 S3 客户端 try: - image_data = io.BytesIO() - image.save(image_data, format='PNG') - image_data.seek(0) - image_bytes = image_data.read() - image_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'image/image_{object_name}.png', io.BytesIO(image_bytes), len(image_bytes), content_type='image/png').object_name}" + key = f"image/image_{object_name}.png" + image_url = f"{AIDA_CLOTHING}/{key}" + s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=data, ContentType='image/png') return image_url except Exception as e: - logging.warning(f"upload_png_mask runtime exception : {e}") + print(f'上传到 S3 失败: {e}') @RunTime def upload_mask_image(mask, object_name): + # 反转掩模 + mask_inverted = cv2.bitwise_not(mask) + # 将掩模转换为 RGBA 格式 + rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) + rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] + # 将图像数据保存到内存中的 BytesIO 对象中 + image_bytes = io.BytesIO() + image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) + image_bytes.seek(0) try: - mask_inverted = cv2.bitwise_not(mask) - # 将掩模的3通道转换为4通道，白色部分不透明，黑色部分透明 - rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) - rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] - image_bytes = io.BytesIO() - image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) - - image_bytes.seek(0) - mask_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'mask/mask_{object_name}.png', image_bytes, len(image_bytes.getvalue()), content_type='image/png').object_name}" + key = f"mask/mask_{object_name}.png" + mask_url = f"{AIDA_CLOTHING}/{key}" + s3.put_object(Bucket=AIDA_CLOTHING, Key=key, Body=image_bytes, ContentType='image/png') return mask_url except Exception as e: - logging.warning(f"upload_png_mask runtime exception : {e}") + print(f'上传到 S3 失败: {e}') + + +"""minio 上传""" + +# @RunTime +# def upload_png_mask(front_image, object_name, mask=None): +# start_time = time.time() +# try: +# mask_url = None +# if mask is not None: +# mask_inverted = cv2.bitwise_not(mask) +# # 将掩模的3通道转换为4通道，白色部分不透明，黑色部分透明 +# rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) +# rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] +# image_bytes = io.BytesIO() +# image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) +# +# image_bytes.seek(0) +# mask_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'mask/mask_{object_name}.png', image_bytes, len(image_bytes.getvalue()), content_type='image/png').object_name}" +# +# image_data = io.BytesIO() +# front_image.save(image_data, format='PNG') +# image_data.seek(0) +# image_bytes = image_data.read() +# image_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'image/image_{object_name}.png', io.BytesIO(image_bytes), len(image_bytes), content_type='image/png').object_name}" +# # print(f"upload_png_mask {object_name} = {time.time() - start_time}") +# return front_image, image_url, mask_url +# except Exception as e: +# logging.warning(f"upload_png_mask runtime exception : {e}") +# +# +# @RunTime +# def upload_layer_image(image, object_name): +# try: +# image_data = io.BytesIO() +# image.save(image_data, format='PNG') +# image_data.seek(0) +# image_bytes = image_data.read() +# image_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'image/image_{object_name}.png', io.BytesIO(image_bytes), len(image_bytes), content_type='image/png').object_name}" +# return image_url +# except Exception as e: +# logging.warning(f"upload_png_mask runtime exception : {e}") +# +# +# @RunTime +# def upload_mask_image(mask, object_name): +# try: +# mask_inverted = cv2.bitwise_not(mask) +# # 将掩模的3通道转换为4通道，白色部分不透明，黑色部分透明 +# rgba_image = cv2.cvtColor(mask_inverted, cv2.COLOR_BGR2BGRA) +# rgba_image[rgba_image[:, :, 0] == 0] = [0, 0, 0, 0] +# image_bytes = io.BytesIO() +# image_bytes.write(cv2.imencode('.png', rgba_image)[1].tobytes()) +# +# image_bytes.seek(0) +# mask_url = f"{AIDA_CLOTHING}/{minio_client.put_object('aida-clothing', f'mask/mask_{object_name}.png', image_bytes, len(image_bytes.getvalue()), content_type='image/png').object_name}" +# return mask_url +# except Exception as e: +# logging.warning(f"upload_png_mask runtime exception : {e}") diff --git a/app/service/generate_image/service_generate_single_logo.py b/app/service/generate_image/service_generate_single_logo.py index 0bb38a0..ed25d74 100644 --- a/app/service/generate_image/service_generate_single_logo.py +++ b/app/service/generate_image/service_generate_single_logo.py @@ -9,25 +9,39 @@ """ 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_single_logo 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.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS)) - self.channel = self.connection.channel() + 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=GI_MODEL_URL) + 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)) @@ -37,20 +51,72 @@ class GenerateSingleLogoImage: 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=GSL_MODEL_NAME, + inputs=inputs, + callback=self.callback + ) + + def callback(self, result, error): + if error: + self.gen_single_logo_data['status'] = "FAILURE" + self.gen_single_logo_data['message'] = str(error) + # self.generate_data['data'] = 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_single_logo(image_result, user_id=self.user_id, category=f"{self.category}", object_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: - generate_data, _ = self.read_tasks_status() + # 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) + + # negative_prompts + text_obj_neg = np.array(self.negative_prompts, dtype="object").reshape((-1, 1)) + # print('text obj neg: ', text_obj_neg) + 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 + seed = np.array(self.seed, dtype="object").reshape((-1, 1)) + print('seed: ', self.seed) + 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.infer(inputs) + time_out = 600 + generate_data = None + while time_out > 0: + generate_data, _ = self.read_tasks_status() + # logger.info(generate_data) + if generate_data['status'] in ["REVOKED", "FAILURE"]: + ctx.cancel() + break + elif generate_data['status'] == "SUCCESS": + break + time_out -= 1 + time.sleep(0.1) + # logger.info(time_out, generate_data) return generate_data except Exception as e: - self.gen_single_logo_data['status'] = "FAILURE" - self.gen_single_logo_data['message'] = str(e) - self.redis_client.set(self.tasks_id, json.dumps(self.gen_single_logo_data)) 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) - self.channel.basic_publish(exchange='', routing_key=GEN_SINGLE_LOGO_RABBITMQ_QUEUES, body=str_generate_data) + if DEBUG is False: + self.channel.basic_publish(exchange='', routing_key=GEN_SINGLE_LOGO_RABBITMQ_QUEUES, body=str_generate_data) + # 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)}") @@ -64,9 +130,9 @@ def infer_cancel(tasks_id): if __name__ == '__main__': rd = GenerateSingleLogoImageModel( - tasks_id="123-8", - prompt='skeleton sitting by the side of a river looking soulful, concert poster, 4k, artistic', - image_url="", + tasks_id="123-89", + prompt='an apple', + seed="1", ) server = GenerateSingleLogoImage(rd) print(server.get_result()) diff --git a/app/service/generate_image/utils/upload_sd_image.py b/app/service/generate_image/utils/upload_sd_image.py index 0e8e542..7cb7f3e 100644 --- a/app/service/generate_image/utils/upload_sd_image.py +++ b/app/service/generate_image/utils/upload_sd_image.py @@ -10,6 +10,7 @@ import io import logging +import boto3 import cv2 from PIL import Image from minio import Minio @@ -17,6 +18,39 @@ from minio import Minio from app.core.config import * minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE) +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_single_logo(image, user_id, category, object_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}/{object_name}', + io.BytesIO(image_bytes), + len(image_bytes), + content_type='image/jpeg' + ) + image_url = f"aida-users/{minio_req.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, object_name):