app.py

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()