create-image/create_utils.py at main · piktid/create-image · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import os
import sys
import json
import requests
from io import BytesIO
from PIL import Image, ImageFile, ImageFilter
from random import randint

from create_api import create_call, handle_notifications, open_image_from_url


def create_image(PARAM_DICTIONARY, TOKEN_DICTIONARY):

    response_json = create_call(PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
    print(response_json)
    PARAM_DICTIONARY['ID_TASK'] = response_json.get('id_task')

    flag_response, response_notifications = handle_notifications(PARAM_DICTIONARY, TOKEN_DICTIONARY)
    if flag_response is False:
        # Error
        print('Error retrieving the generated images')
        return False

    links = []
    # Get the links from the response
    NUM_GENERATIONS = PARAM_DICTIONARY.get('NUM_GENERATIONS', 1)
    full_links = response_notifications.get("links")
    NUM_GENERATED = len(full_links)

    if NUM_GENERATED < NUM_GENERATIONS:
        # check NSFW key
        print(response_notifications.get("nsfw"))

    for i in range(NUM_GENERATED):
        download_link = ((response_notifications.get("links"))[i]).get("l")
        links.append(download_link)
        print('new image ready for download:', download_link)

    flag_save, final_path = save_generated_imgs(links, PARAM_DICTIONARY, TOKEN_DICTIONARY)
    if flag_save is False:
        print('Error: failed to save the generated image')
        return False, None

    return flag_save, final_path


def save_generated_imgs(links, PARAM_DICTIONARY, TOKEN_DICTIONARY):
    print('Saving the generated image')
    final_path = None
    try:
        options_str = ''
        seed = PARAM_DICTIONARY.get('SEED', 0)

        ar = PARAM_DICTIONARY.get('ASPECT_RATIO')
        if ar is not None:
            options_str = options_str + 'ar' + ar.split(':')[0] + '-' + ar.split(':')[1] + '_'

        prompt = PARAM_DICTIONARY.get('PROMPT', None)
        if prompt is not None:
            options_str = options_str+prompt[0:30]+'_'

        path_output = PARAM_DICTIONARY.get('OUTPUT_PATH')

        if path_output is None:
            path_output = os.path.abspath(os.getcwd())

        img_format = 'png'
        filename = 'img_created'
        image_path = os.path.join(path_output, filename)

        for i, link in enumerate(links):
            final_path = image_path.split('.')[0]+'_'+str(seed)+'_'+str(i)+'_'+options_str+'.'+img_format
            print(f'Final path: {final_path}')

            # save the generated image in the generated folder
            src_img = open_image_from_url(link)
            src_img.save(final_path)

    except Exception as e:
        print(f'Error: {e}')
        return False, None

    return True, final_path