Spaces:
Running
Running
| import os | |
| from oss_utils import ossService | |
| import requests | |
| import random | |
| import json | |
| import time | |
| from diffusers.utils import load_image | |
| from PIL import Image, ImageDraw, ImageFont | |
| import numpy as np | |
| # oss config | |
| BUCKET = os.environ.get('BUCKET', '') | |
| ENDPOINT = os.environ.get('ENDPOINT', '') | |
| PREFIX = os.environ.get('OSS_DIR_PREFIX', '') | |
| AK = os.environ.get('AK', '') | |
| SK = os.environ.get('SK', '') | |
| DASHONE_SERVICE_ID = os.environ.get('SERVICE_ID', '') | |
| URL = os.environ.get('URL', '') | |
| GET_URL = os.environ.get('GET_URL', '') | |
| OUTPUT_PATH = './output' | |
| os.makedirs(OUTPUT_PATH, exist_ok=True) | |
| oss_service = ossService(AK, SK, ENDPOINT, BUCKET, PREFIX) | |
| def async_request_and_query(data, request_id): | |
| url = URL | |
| headers = {'Content-Type': 'application/json'} | |
| # 1.发起一个异步请求 | |
| print('Start sending request') | |
| response = requests.post(url, headers=headers, data=json.dumps(data)) | |
| if response.status_code != requests.codes.ok: | |
| response.raise_for_status() | |
| response_json = json.loads(response.content.decode("utf-8")) | |
| print('Finish sending request') | |
| # 2.异步查询结果 | |
| is_running = True | |
| running_print_count = 0 | |
| sign_oss_path = None | |
| task_id = response_json['header']['task_id'] | |
| get_url = GET_URL | |
| get_header = headers | |
| get_data = {"header": {"request_id":request_id,"service_id":DASHONE_SERVICE_ID,"task_id":f"{task_id}"}} | |
| print('Start querying results') | |
| while is_running: | |
| response2 = requests.post(get_url, headers=get_header, data=json.dumps(get_data)) | |
| if response2.status_code != requests.codes.ok: | |
| response2.raise_for_status() | |
| response2_json = json.loads(response2.content.decode("utf-8")) | |
| task_status = response2_json['header']['task_status'] | |
| if task_status == 'SUCCESS': | |
| sign_oss_path = response2_json['payload']['output']['res'] | |
| break | |
| elif task_status in ['FAILED', 'ERROR'] or running_print_count >= 120: | |
| raise ValueError(f'Task Failed') | |
| else: | |
| time.sleep(1) | |
| running_print_count += 1 | |
| continue | |
| print('Task succeeded!') | |
| return sign_oss_path | |
| # def add_transparent_watermark(pil_image, watermark_text, position, opacity, font_path, font_size): | |
| # # 加载字体 | |
| # font = ImageFont.truetype(font_path, font_size) | |
| # # 创建一个半透明的水印图层 | |
| # watermark_layer = Image.new("RGBA", pil_image.size) | |
| # draw = ImageDraw.Draw(watermark_layer) | |
| # # 文本颜色和透明度 | |
| # text_color = (255, 255, 255, opacity) # 白色文本 | |
| # outline_color = (0, 0, 0, opacity) # 黑色轮廓 | |
| # # 获取文本尺寸 | |
| # text_width = draw.textlength(watermark_text, font=font) | |
| # text_height = text_width // 5 | |
| # # 计算水印位置 | |
| # img_width, img_height = pil_image.size | |
| # x = img_width - text_width - position[0] | |
| # y = img_height - text_height - position[1] | |
| # outline_range = 1 # 轮廓的粗细 | |
| # for adj in range(-outline_range, outline_range+1): | |
| # for ord in range(-outline_range, outline_range+1): | |
| # if adj != 0 or ord != 0: # 避免中心位置，那是真正的文本 | |
| # draw.text((x+adj, y+ord), watermark_text, font=font, fill=outline_color) | |
| # # 将文本绘制到水印层上 | |
| # draw.text((x, y), watermark_text, font=font, fill=text_color) | |
| # # 将水印层叠加到原始图像上 | |
| # pil_image_with_watermark = Image.alpha_composite(pil_image.convert("RGBA"), watermark_layer) | |
| # # 返回添加了水印的图像 | |
| # return pil_image_with_watermark | |
| def inference(input_image, upscale): | |
| # process alpha channel | |
| alpha_channel = input_image.split()[-1] | |
| input_image = input_image.convert('RGB') | |
| local_save_path = os.path.join(OUTPUT_PATH, 'tmp.png') | |
| local_output_save_path = os.path.join(OUTPUT_PATH, 'out.png') | |
| input_image.save(local_save_path) | |
| # generate image url | |
| oss_key = os.path.join(PREFIX, 'tmp.png') | |
| _, image_url = oss_service.uploadOssFile(oss_key, local_save_path) | |
| # rm local file | |
| if os.path.isfile(local_save_path): | |
| os.remove(local_save_path) | |
| data = {} | |
| data_header = {} | |
| data_payload = {} | |
| data_input = {} | |
| data_para = {} | |
| data_header['request_id'] = "".join(random.sample("0123456789abcdefghijklmnopqrstuvwxyz", 10)) | |
| data_header['service_id'] = DASHONE_SERVICE_ID | |
| data_input['image_url'] = image_url | |
| data_para['upscale'] = upscale | |
| data_para['platform'] = 'modelscope' | |
| data_payload['input'] = data_input | |
| data_payload['parameters'] = data_para | |
| data['header'] = data_header | |
| data['payload'] = data_payload | |
| try: | |
| output_url = async_request_and_query(data=data, request_id=data_header['request_id']) | |
| download_status = oss_service.downloadFile(output_url, local_output_save_path) | |
| if not download_status: | |
| raise ValueError(f'Download output image failed') | |
| except: | |
| output_image = load_image('./error.png') | |
| return [output_image], 'The input image format or resolution does not meet the requirements. Please change the image or resize it and try again.' | |
| output_image = load_image(local_output_save_path) | |
| if os.path.isfile(local_output_save_path): | |
| os.remove(local_output_save_path) | |
| out_width, out_height = output_image.size | |
| # add alpha channel | |
| output_alpha_channel = alpha_channel.resize(output_image.size, resample=Image.LANCZOS) | |
| # merge | |
| output_image_alpha = Image.merge("RGBA", (*output_image.split(), output_alpha_channel)) | |
| # new_width = out_width // 4 | |
| # new_height = out_height // 4 | |
| # resized_image = np.array(input_image.resize((new_width, new_height))) | |
| # np_output_image = np.array(output_image) | |
| # np_output_image[0:new_height, 0:new_width] = resized_image.copy() | |
| # new_output_image = Image.fromarray(np_output_image) | |
| # new_output_image = add_transparent_watermark( | |
| # pil_image=new_output_image, | |
| # watermark_text="追影-放大镜", | |
| # position=(5, 5), | |
| # opacity=200, | |
| # font_path="AlibabaPuHuiTi-3-45-Light.ttf", # 例如："Arial", "Helvetica", "Times New Roman" | |
| # font_size= out_width // 30 | |
| # ) | |
| org_width, org_height = input_image.size | |
| if max(org_width, org_height) > 1920 or min(org_width, org_height) > 1080: | |
| msg = 'The input image size has exceeded the optimal range. You can consider scaling the input image for better generation effect' | |
| else: | |
| msg = 'Task succeeded' | |
| return [output_image_alpha], msg | |