From e1f19f62b37ad7542c1b055dcf7960439761e0a7 Mon Sep 17 00:00:00 2001 From: zhouchengrong Date: Sun, 1 Dec 2024 14:24:48 +0800 Subject: [PATCH] =?UTF-8?q?feat=20=20generate=20img=20fast=20version=20?= =?UTF-8?q?=E9=83=A8=E7=BD=B2=20fix?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/core/config.py | 1 + app/schemas/generate_image.py | 1 + .../generate_image/service_generate_image.py | 9 ++++-- .../service_generate_product_image.py | 32 +++++++++++++------ app/service/utils/new_oss_client.py | 2 +- 5 files changed, 32 insertions(+), 13 deletions(-) diff --git a/app/core/config.py b/app/core/config.py index ff6f0a1..97e014d 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -101,6 +101,7 @@ SR_RABBITMQ_QUEUES = os.getenv("SR_RABBITMQ_QUEUES", f"SuperResolution{RABBITMQ_ # GenerateImage service config GI_MODEL_NAME = 'stable_diffusion_xl' +FAST_GI_MODEL_URL = '10.1.1.243:10011' GI_MODEL_URL = '10.1.1.240:10041' GI_MINIO_BUCKET = "aida-users" GI_RABBITMQ_QUEUES = os.getenv("GI_RABBITMQ_QUEUES", f"GenerateImage{RABBITMQ_ENV}") diff --git a/app/schemas/generate_image.py b/app/schemas/generate_image.py index 3dd7cf8..11e295f 100644 --- a/app/schemas/generate_image.py +++ b/app/schemas/generate_image.py @@ -8,6 +8,7 @@ class GenerateImageModel(BaseModel): mode: str category: str gender: str + version: str class GenerateSingleLogoImageModel(BaseModel): diff --git a/app/service/generate_image/service_generate_image.py b/app/service/generate_image/service_generate_image.py index dac211c..7d2937b 100644 --- a/app/service/generate_image/service_generate_image.py +++ b/app/service/generate_image/service_generate_image.py @@ -35,7 +35,11 @@ class GenerateImage: # 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) + if request_data.version == "fast": + self.grpc_client = grpcclient.InferenceServerClient(url=FAST_GI_MODEL_URL) + else: + 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 @@ -185,7 +189,8 @@ if __name__ == '__main__': image_url="aida-collection-element/87/Printboard/842c09cf-7297-42d9-9e6e-9c17d4a13cb5.jpg", mode='txt2img', category="test", - gender="male" + gender="male", + version="fast" ) server = GenerateImage(rd) print(server.get_result()) diff --git a/app/service/generate_image/service_generate_product_image.py b/app/service/generate_image/service_generate_product_image.py index 16b814b..ebffb05 100644 --- a/app/service/generate_image/service_generate_product_image.py +++ b/app/service/generate_image/service_generate_product_image.py @@ -55,7 +55,10 @@ class GenerateProductImage: self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data)) else: # pil图像转成numpy数组 - image = result.as_numpy("generated_inpaint_image") + if self.product_type == "single": + image = result.as_numpy("generated_inpaint_image") + else: + image = result.as_numpy("generated_cnet_image") image_result = Image.fromarray(np.squeeze(image.astype(np.uint8))).resize(self.image_size) cropped_image = post_processing_image(image_result, self.left, self.top) image_url = upload_SDXL_image(cropped_image, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png") @@ -75,9 +78,14 @@ class GenerateProductImage: self.image = cv2.resize(self.image, (1024, 1024)) images = [self.image.astype(np.uint8)] * self.batch_size - text_obj = np.array(prompts, dtype="object").reshape((1)) - image_obj = np.array(images, dtype=np.uint8).reshape((1024, 1024, 3)) - image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1)) + 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, 1024, 1024, 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((1024, 1024, 3)) + image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1)) # 假设 prompts、images 和 self.image_strength 已经定义 @@ -91,7 +99,11 @@ class GenerateProductImage: inputs = [input_text, input_image, input_image_strength] - ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback) + + if self.product_type == "single": + ctx = self.grpc_client.async_infer(model_name="stable_diffusion_xl_cnet_inpaint", 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: @@ -173,10 +185,10 @@ if __name__ == '__main__': rd = GenerateProductImageModel( tasks_id="123-89", # prompt="", - image_strength=0.65, - prompt="The best quality, masterpiece, real image. A handsome man wearing blouse, outwear, trousers, 8K realistic, HUD", - image_url="aida-results/result_00097282-ebb2-11ee-a822-b48351119060.png", - product_type="overall" + image_strength=0.7, + prompt="The best quality, masterpiece,outwear, 8K realistic, HUD", + image_url="aida-results/result_836dce70-ad59-11ef-86ab-0242ac130002.png", + product_type="single" ) server = GenerateProductImage(rd) - print(server.get_result()) + print(server.get_result()) \ No newline at end of file diff --git a/app/service/utils/new_oss_client.py b/app/service/utils/new_oss_client.py index 5067d15..6dd22bd 100644 --- a/app/service/utils/new_oss_client.py +++ b/app/service/utils/new_oss_client.py @@ -82,7 +82,7 @@ if __name__ == '__main__': # 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-results/result_00097282-ebb2-11ee-a822-b48351119060.png" + url ="aida-users/89/test/123-89.png" # url = "aida-collection-element/12148/Sketchboard/95ea577b-305b-4a62-b30a-39c0dd3ddb3f.png" read_type = "2"