Spaces:
Running
Running
| import os | |
| import glob | |
| import gradio as gr | |
| import base64 | |
| import cv2 | |
| import numpy as np | |
| import oss2 | |
| import time | |
| from ai_service_python_sdk.client.api.ai_service_aigc_images_api import AIGCImagesApi | |
| from ai_service_python_sdk.client.api.ai_service_job_api import AiServiceJobApi | |
| from ai_service_python_sdk.client.api_client import ApiClient | |
| from ai_service_python_sdk.test import appId, host, token | |
| host = os.getenv("PAI_REC_HOST") | |
| appId = os.getenv("PAI_REC_APP_ID") | |
| token = os.getenv("PAI_REC_TOKEN") | |
| access_key_id = os.getenv('OSS_ACCESS_KEY_ID') | |
| access_key_secret = os.getenv('OSS_ACCESS_KEY_SECRET') | |
| bucket_name = os.getenv('OSS_BUCKET') | |
| endpoint = os.getenv('OSS_ENDPOINT') | |
| def upload_file(files, current_files): | |
| file_paths = [file_d['name'] for file_d in current_files] + [file.name for file in files] | |
| return file_paths | |
| def decode_image_from_base64jpeg(base64_image): | |
| image_bytes = base64.b64decode(base64_image) | |
| np_arr = np.frombuffer(image_bytes, np.uint8) | |
| image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| return image | |
| def upload(image_path, number): | |
| bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name) | |
| file_name = image_path.split('/')[-1] | |
| ext = file_name.split('.')[-1] | |
| file_name = str(number) + '.' + ext | |
| timestamp = str(time.time()).split('.')[0] | |
| bucket_folder = 'aigc-data/easyphoto_demo_data/' + timestamp + '_user_lora/' | |
| oss_file_path = bucket_folder + file_name | |
| bucket.put_object_from_file(oss_file_path, image_path) | |
| file_url = 'https://' + bucket_name + '.' + endpoint + '/' + bucket_folder + file_name | |
| return file_url | |
| def upload_template(image_path): | |
| bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name) | |
| file_name = image_path.split('/')[-1] | |
| timestamp = str(time.time()).split('.')[0] | |
| bucket_folder = 'aigc-data/easyphoto_demo_data/' + timestamp + '_user_template/' | |
| oss_file_path = bucket_folder + file_name | |
| bucket.put_object_from_file(oss_file_path, image_path) | |
| file_url = 'https://' + bucket_name + '.' + endpoint + '/' + bucket_folder + file_name | |
| return file_url | |
| def easyphoto_train(instance_images): | |
| images = [] | |
| if instance_images is None or len(instance_images)==0: | |
| output = 'Status: no image updated!' | |
| return output, [], [] | |
| for number, image in enumerate(instance_images): | |
| image_path = image['name'] | |
| image_url = upload(image_path, number) | |
| images.append(image_url) | |
| client = ApiClient(host, appId, token) | |
| api = AIGCImagesApi(client) | |
| response = api.aigc_images_train(images, 'photog_train_freetier', None) | |
| message = response.message | |
| model_id = response.data['model_id'] | |
| job_id = response.data['job_id'] | |
| if message == 'success': | |
| state = 'training job submitted.' | |
| output = 'Status: ' + state | |
| print("job id: " + str(job_id)) | |
| print("model id: " + str(model_id)) | |
| return output, job_id, model_id | |
| else: | |
| output = 'Status: submitting training job failed!' | |
| return output, [], [] | |
| def easyphoto_check(job_id): | |
| client = ApiClient(host, appId, token) | |
| api = AiServiceJobApi(client) | |
| if job_id is None: | |
| output = 'Status: checking training status failed! No job id.' | |
| else: | |
| try: | |
| job_id = int(str(job_id).strip()) | |
| response = api.get_async_job_with_id(job_id) | |
| message = response.data['job']['message'] | |
| output = 'Status: ' + message | |
| except: | |
| output = 'Status: checking training status failed!' | |
| return output | |
| def easyphoto_infer(model_id, selected_template_images, additional_prompt, seed, before_face_fusion_ratio, after_face_fusion_ratio, first_diffusion_steps, first_denoising_strength, second_diffusion_steps, second_denoising_strength, crop_face_preprocess, apply_face_fusion_before, apply_face_fusion_after, color_shift_middle, color_shift_last, background_restore): | |
| image_urls = [] | |
| if len(selected_template_images) == 0: | |
| output_info = 'Status: no templete selected!' | |
| return output_info, [] | |
| selected_template_images = eval(selected_template_images) | |
| for image in selected_template_images: | |
| image_url = upload_template(image) | |
| image_urls.append(image_url) | |
| client = ApiClient(host, appId, token) | |
| api = AIGCImagesApi(client) | |
| outputs = [] | |
| output_info = None | |
| if model_id is None: | |
| output_info = 'Status: no model id provided!' | |
| return output_info, [] | |
| model_id = str(model_id).strip() | |
| print('model id: ' + model_id) | |
| if job_id is None: | |
| output_info = 'Status: no job id provided, please do model training first!' | |
| return output_info, [] | |
| job_id = str(job_id).strip() | |
| print('job id: ' + job_id) | |
| check_client = ApiClient(host, appId, token) | |
| check_api = AiServiceJobApi(check_client) | |
| try: | |
| job_id = int(str(job_id).strip()) | |
| response = check_api.get_async_job_with_id(job_id) | |
| message = response.data['job']['message'] | |
| if not message == 'success': | |
| output = 'Status: model training incomplete!' | |
| return output, [] | |
| except: | |
| output = 'Status: checking training status failed, please do model training first!' | |
| return output, [] | |
| for image_url in image_urls: | |
| try: | |
| params = { | |
| "additional_prompt": additional_prompt, | |
| "seed": seed, | |
| "before_face_fusion_ratio": before_face_fusion_ratio, | |
| "after_face_fusion_ratio": after_face_fusion_ratio, | |
| "first_diffusion_steps": first_diffusion_steps, | |
| "first_denoising_strength": first_denoising_strength, | |
| "second_diffusion_steps": second_diffusion_steps, | |
| "second_denoising_strength": second_denoising_strength, | |
| "crop_face_preprocess": crop_face_preprocess, | |
| "apply_face_fusion_before": apply_face_fusion_before, | |
| "apply_face_fusion_after": apply_face_fusion_after, | |
| "color_shift_middle": color_shift_middle, | |
| "color_shift_last": color_shift_last, | |
| "background_restore": background_restore | |
| } | |
| response = api.aigc_images_create(model_id, image_url, 'photog_infer_freetier', params) | |
| except: | |
| output_info = 'Status: calling eas service failed!' | |
| return output_info, [] | |
| data = response.data | |
| message = response.message | |
| if message == 'success': | |
| image = data['image'] | |
| image = decode_image_from_base64jpeg(image) | |
| outputs.append(image) | |
| output_info = 'Status: generating image succesfully!' | |
| else: | |
| output_info = 'Status: generating image failed!' | |
| return output_info, [] | |
| return output_info, outputs | |
| with gr.Blocks() as easyphoto_demo: | |
| model_id = gr.Textbox(visible=False) | |
| job_id = gr.Textbox(visible=False) | |
| with gr.TabItem('Training'): | |
| with gr.Blocks(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| instance_images = gr.Gallery().style(columns=[4], rows=[2], object_fit="contain", height="auto") | |
| with gr.Row(): | |
| upload_button = gr.UploadButton( | |
| "Upload Photos", file_types=["image"], file_count="multiple" | |
| ) | |
| clear_button = gr.Button("Clear Photos") | |
| clear_button.click(fn=lambda: [], inputs=None, outputs=instance_images) | |
| upload_button.upload(upload_file, inputs=[upload_button, instance_images], outputs=instance_images, queue=False) | |
| gr.Markdown( | |
| ''' | |
| Training steps: | |
| 1. Please upload 5-20 half-body photos or head and shoulder photos, ensuring that the facial proportions are not too small. | |
| 2. Click the training button below to submit the training task. It will take approximately 15 minutes, and you can check the status of your training task. Please refrain from clicking the submit training task button multiple times! | |
| 3. Once the model training is completed, the task status will display success. Switch to inference mode and generate photos based on the template. | |
| 4. If you experience lag during uploading, please resize the uploaded images to a size below 1.5MB if possible. | |
| 5. During the training or inference process, please do not refresh or close the window. | |
| ''' | |
| ) | |
| with gr.Row(): | |
| run_button = gr.Button('Submit My Training Job') | |
| check_button = gr.Button('Check My Training Job Status') | |
| output_message = gr.Textbox(value="", label="Status", interactive=False) | |
| run_button.click(fn=easyphoto_train, | |
| inputs=[instance_images], | |
| outputs=[output_message, job_id, model_id]) | |
| check_button.click(fn=easyphoto_check, | |
| inputs=[job_id], | |
| outputs=[output_message]) | |
| with gr.TabItem('Inference'): | |
| templates = glob.glob(r'./*.jpg') | |
| preset_template = list(templates) | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| template_gallery_list = [(i, i) for i in preset_template] | |
| gallery = gr.Gallery(template_gallery_list).style(columns=[4], rows=[2], object_fit="contain", height="auto") | |
| def select_function(evt: gr.SelectData): | |
| return [preset_template[evt.index]] | |
| selected_template_images = gr.Text(show_label=False, visible=False, placeholder="Selected") | |
| gallery.select(select_function, None, selected_template_images) | |
| with gr.Accordion("Advanced Options", open=False): | |
| additional_prompt = gr.Textbox( | |
| label="Additional Prompt", | |
| lines=3, | |
| value='masterpiece, beauty', | |
| interactive=True | |
| ) | |
| seed = gr.Textbox( | |
| label="Seed", | |
| value=12345, | |
| ) | |
| with gr.Row(): | |
| before_face_fusion_ratio = gr.Slider( | |
| minimum=0.2, maximum=0.8, value=0.50, | |
| step=0.05, label='Face Fusion Ratio Before' | |
| ) | |
| after_face_fusion_ratio = gr.Slider( | |
| minimum=0.2, maximum=0.8, value=0.50, | |
| step=0.05, label='Face Fusion Ratio After' | |
| ) | |
| with gr.Row(): | |
| first_diffusion_steps = gr.Slider( | |
| minimum=15, maximum=50, value=50, | |
| step=1, label='First Diffusion steps' | |
| ) | |
| first_denoising_strength = gr.Slider( | |
| minimum=0.30, maximum=0.60, value=0.45, | |
| step=0.05, label='First Diffusion denoising strength' | |
| ) | |
| with gr.Row(): | |
| second_diffusion_steps = gr.Slider( | |
| minimum=15, maximum=50, value=20, | |
| step=1, label='Second Diffusion steps' | |
| ) | |
| second_denoising_strength = gr.Slider( | |
| minimum=0.20, maximum=0.40, value=0.30, | |
| step=0.05, label='Second Diffusion denoising strength' | |
| ) | |
| with gr.Row(): | |
| crop_face_preprocess = gr.Checkbox( | |
| label="Crop Face Preprocess", | |
| value=True | |
| ) | |
| apply_face_fusion_before = gr.Checkbox( | |
| label="Apply Face Fusion Before", | |
| value=True | |
| ) | |
| apply_face_fusion_after = gr.Checkbox( | |
| label="Apply Face Fusion After", | |
| value=True | |
| ) | |
| with gr.Row(): | |
| color_shift_middle = gr.Checkbox( | |
| label="Apply color shift first", | |
| value=True | |
| ) | |
| color_shift_last = gr.Checkbox( | |
| label="Apply color shift last", | |
| value=True | |
| ) | |
| background_restore = gr.Checkbox( | |
| label="Background Restore", | |
| value=False | |
| ) | |
| with gr.Box(): | |
| gr.Markdown( | |
| ''' | |
| Parameters: | |
| 1. **Face Fusion Ratio Before** represents the proportion of the first facial fusion, which is higher and more similar to the training object. | |
| 2. **Face Fusion Ratio After** represents the proportion of the second facial fusion, which is higher and more similar to the training object. | |
| 3. **Crop Face Preprocess** represents whether to crop the image before generation, which can adapt to images with smaller faces. | |
| 4. **Apply Face Fusion Before** represents whether to perform the first facial fusion. | |
| 5. **Apply Face Fusion After** represents whether to perform the second facial fusion. | |
| ''' | |
| ) | |
| with gr.Column(): | |
| gr.Markdown('Generated Results') | |
| output_images = gr.Gallery( | |
| label='Output', | |
| show_label=False | |
| ).style(columns=[4], rows=[2], object_fit="contain", height="auto") | |
| display_button = gr.Button('Start Generation') | |
| infer_progress = gr.Textbox( | |
| label="Generation Progress", | |
| value="", | |
| interactive=False | |
| ) | |
| display_button.click( | |
| fn=easyphoto_infer, | |
| inputs=[model_id, selected_template_images, additional_prompt, seed, before_face_fusion_ratio, after_face_fusion_ratio, first_diffusion_steps, first_denoising_strength, second_diffusion_steps, second_denoising_strength, crop_face_preprocess, apply_face_fusion_before, apply_face_fusion_after, color_shift_middle, color_shift_last, background_restore], | |
| outputs=[infer_progress, output_images] | |
| ) | |
| gr.Markdown( | |
| """ | |
| Useful Links | |
| EasyPhoto GitHub: https://github.com/aigc-apps/sd-webui-EasyPhoto | |
| Alibaba Cloud Freetier: https://help.aliyun.com/document_detail/2567864.html | |
| PAI-DSW Gallery: https://gallery.pai-ml.com/#/preview/deepLearning/cv/stable_diffusion_easyphoto | |
| """) | |
| easyphoto_demo.launch(share=False).queue() |