New consistent identity endpoints

nuletizia · Oct 7, 2024 · 2ff5f3b · 2ff5f3b
1 parent 96303af
commit 2ff5f3b
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 39 deletions.
diff --git a/README.md b/README.md
@@ -5,7 +5,7 @@
 </p>
 
 
-# EraseID - v2.3.0
+# EraseID - v2.4.0
 [![Official Website](https://img.shields.io/badge/Official%20Website-piktid.com-blue?style=flat&logo=world&logoColor=white)](https://piktid.com)
 [![Discord Follow](https://dcbadge.vercel.app/api/server/FJU39e9Z4P?style=flat)](https://discord.com/invite/FJU39e9Z4P)
 
@@ -89,22 +89,22 @@ $ python3 main.py --identity_name 'pippo'
 </a>
 
 It is also possible to use the same real identity from a source image into multiple target photos. To avoid malicious uses, only verified trusted users have access to this feature. If you are interested, please contact us via Discord.
-To use a real identity from a source photo 'mydir/myfile.jpg', use the command:
+To use a real identity from a source photo with local path 'mydir/myfile.jpg', use the command:
 
 ```bash
 $ python3 main.py --filepath 'mydir/myfile.jpg' --identity_filepath 'mydir/myfile.jpg' --identity_name 'myidentityname'
 ```
 
 The identity will be stored as 'myidentityname' and it will be used as reference input for the target image 'mydir/myfile.jpg'
 
-If you want to swap also the hair, use the command
+If you want to swap also the hair and use an identity located at 'mylink/myfile.jpg', use the command
 ```bash
-$ python3 main.py --filepath 'mydir/myfile.jpg' --identity_filepath 'mydir/myfile.jpg' --identity_name 'myidentityname' --hair
+$ python3 main.py --filepath 'mydir/myfile.jpg' --identity_url 'mylink/myfile.jpg' --identity_name 'myidentityname' --hair
 ```
 
-It is possible to change the default generation parameters, to do that use the command (you need to be a premium user)
+It is possible to change the default consistent generation parameters, to do that use the command (you need to be a premium user)
 ```bash
-$ python3 main.py --filepath 'mydir/myfile.jpg' --identity_filepath 'mydir/myfile.jpg' --identity_name 'myidentityname' --guidance_scale '1.5' --controlnet_scale '0.1' --prompt_strength '0.5'
+$ python3 main.py --filepath 'mydir/myfile.jpg' --identity_filepath 'mydir/myfile.jpg' --identity_name 'myidentityname' --prompt_strength '0.5' --seed 0
 ```
 
 ## Change facial expression

diff --git a/eraseid_api.py b/eraseid_api.py
@@ -138,34 +138,40 @@ def upload_and_detect_call(src_img, PARAM_DICTIONARY, TOKEN_DICTIONARY):
     return image_address, indices_info, selected_faces_list
 
 
-def upload_reference_face_call(src_img, identity_name, TOKEN_DICTIONARY):
-    # upload the image into PiktID's servers
+def upload_reference_face_call(PARAM_DICTIONARY, TOKEN_DICTIONARY):
+    IDENTITY_NAME = PARAM_DICTIONARY.get('IDENTITY_NAME')
+    face_full_path = PARAM_DICTIONARY.get('IDENTITY_PATH')
+    if face_full_path is None:
+        face_url = PARAM_DICTIONARY.get('IDENTITY_URL')
+        face_response = requests.get(face_url)
+        face_response.raise_for_status()  
+        face_file = BytesIO(face_response.content)
+        face_file.name = 'face.jpg' 
+    else:
+        face_file = open(face_full_path, 'rb')
+
+    # start the generation process given the image parameters
     TOKEN = TOKEN_DICTIONARY.get('access_token', '')
     URL_API = TOKEN_DICTIONARY.get('url_api')
 
-    src_img_B = im_2_buffer(src_img)
-
-    m = MultipartEncoder(
-                        fields={'identity_name': identity_name, 'file': ('file', src_img_B, 'text/plain')}
-                        )
-
-    response = requests.post(URL_API+'/upload_identity',
-                             headers={"Content-Type": m.content_type,
-                                      'Authorization': 'Bearer '+TOKEN},
-                             data=m,
+    response = requests.post(URL_API+'/consistent_identities/upload_face', 
+                             headers={'Authorization': 'Bearer '+TOKEN},
+                             files={'face': face_file},
+                             data={'identity_name': IDENTITY_NAME},
                              )
-    # if the access token is expired
+
     if response.status_code == 401:
         TOKEN_DICTIONARY = refresh_call(TOKEN_DICTIONARY)
         TOKEN = TOKEN_DICTIONARY.get('access_token', '')
         # try with new TOKEN
-        response = requests.post(URL_API+'/upload_identity',
-                                 headers={"Content-Type": m.content_type,
-                                          'Authorization': 'Bearer '+TOKEN},
-                                 data=m,
+        response = requests.post(URL_API+'/consistent_identities/upload_face', 
+                                 headers={'Authorization': 'Bearer '+TOKEN},
+                                 files={'face': face_file},
+                                 data={'identity_name': IDENTITY_NAME},
                                  )
-    print(response.content)
+
     response_json = json.loads(response.text)
+
     return response_json
 
 
@@ -248,14 +254,8 @@ def update_data_generation_call(data, PARAM_DICTIONARY):
 
 def update_data_skin_call(data, PARAM_DICTIONARY, TOKEN_DICTIONARY):
 
-    CUSTOM_PROMPT_FLAG = PARAM_DICTIONARY.get('CUSTOM_PROMPT_FLAG')
     SEED = PARAM_DICTIONARY.get('SEED')
 
-    if CUSTOM_PROMPT_FLAG is True:
-        # overwrite the keyword mechanism
-        extra_data = {'description_type': 'txt'}
-        data.update(extra_data)
-
     OPTIONS_DICT = {}
     if SEED is not None:
         OPTIONS_DICT = {**OPTIONS_DICT, 'seed': SEED}
@@ -296,6 +296,49 @@ def generation_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TOKEN_DIC
     return response_json
 
 
+def consistent_generation_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TOKEN_DICTIONARY):
+
+    IDENTITY_NAME = PARAM_DICTIONARY.get('IDENTITY_NAME')
+
+    SEED = PARAM_DICTIONARY.get('SEED')
+    PROMPT_STRENGTH = PARAM_DICTIONARY.get('PROMPT_STRENGTH')
+
+    OPTIONS_DICT = {}
+
+    if SEED is not None:
+        OPTIONS_DICT = {**OPTIONS_DICT, 'seed': SEED}
+
+    if PROMPT_STRENGTH is not None:
+        OPTIONS_DICT = {**OPTIONS_DICT, 'prompt_strength': PROMPT_STRENGTH}
+
+    OPTIONS = json.dumps(OPTIONS_DICT)
+    extra_options = {'options': OPTIONS}
+
+    data = {'flag_sync': False, 'identity_name': IDENTITY_NAME, 'id_image': image_address, 'id_face': idx_face, 'prompt': prompt}
+    data.update(extra_options)
+    print(f'data to send to generation: {data}')
+
+    # start the generation process given the image parameters
+    TOKEN = TOKEN_DICTIONARY.get('access_token', '')
+    URL_API = TOKEN_DICTIONARY.get('url_api')
+
+    response = requests.post(URL_API+'/consistent_identities/generate',
+                             headers={'Authorization': 'Bearer '+TOKEN},
+                             json=data,
+                             )
+    # if the access token is expired
+    if response.status_code == 401:
+        TOKEN_DICTIONARY = refresh_call(TOKEN_DICTIONARY)
+        TOKEN = TOKEN_DICTIONARY.get('access_token', '')
+        # try with new TOKEN
+        response = requests.post(URL_API+'/consistent_identities/generate',
+                                 headers={'Authorization': 'Bearer '+TOKEN},
+                                 json=data,
+                                 )
+    response_json = json.loads(response.text)
+    return response_json
+
+
 def change_expression_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TOKEN_DICTIONARY):
 
     SEED = PARAM_DICTIONARY.get('SEED')

diff --git a/eraseid_utils.py b/eraseid_utils.py
@@ -1,6 +1,6 @@
 import json
 
-from eraseid_api import upload_and_detect_call, upload_reference_face_call, selection_call, get_identities_call, generation_call, change_expression_call, change_skin_call, handle_notifications_new_generation, handle_notifications_new_skin, get_generated_faces, get_last_generated_face, set_identity_call, replace_call
+from eraseid_api import upload_and_detect_call, upload_reference_face_call, selection_call, get_identities_call, generation_call, consistent_generation_call, change_expression_call, change_skin_call, handle_notifications_new_generation, handle_notifications_new_skin, get_generated_faces, get_last_generated_face, set_identity_call, replace_call
 from cfe_keywords import cfe_dict
 
 
@@ -14,11 +14,12 @@ def find_key_by_value(target_value):
 def process_single_image(input_image, PARAM_DICTIONARY, TOKEN_DICTIONARY):
 
     IDENTITY_NAME = PARAM_DICTIONARY.get('IDENTITY_NAME')
-    IDENTITY_IMAGE = PARAM_DICTIONARY.get('IDENTITY_IMAGE')
+    IDENTITY_PATH = PARAM_DICTIONARY.get('IDENTITY_PATH')
+    IDENTITY_URL = PARAM_DICTIONARY.get('IDENTITY_URL')
 
     # upload the source identity (swap feature), if any
-    if IDENTITY_IMAGE is not None and IDENTITY_NAME is not None:
-        response = upload_reference_face_call(IDENTITY_IMAGE, IDENTITY_NAME, TOKEN_DICTIONARY)
+    if IDENTITY_NAME is not None and (IDENTITY_PATH is not None or IDENTITY_URL is not None):
+        response = upload_reference_face_call(PARAM_DICTIONARY, TOKEN_DICTIONARY)
         print(f'Identity: {IDENTITY_NAME} correctly uploaded to PiktID servers')
 
     CHANGE_ALL_FACES = PARAM_DICTIONARY.get('CHANGE_ALL_FACES')
@@ -121,6 +122,10 @@ def process_single_face(idx_face, count, PARAM_DICTIONARY, TOKEN_DICTIONARY):
         response = change_expression_call(image_address=image_id, idx_face=idx_face, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
         print(f'Cfe response:{response}')
 
+    elif IDENTITY_NAME is not None:
+        print('Swapping faces')
+        response = consistent_generation_call(image_address=image_id, idx_face=idx_face, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
+        print(f'Swap response:{response}')
     else:
         print('Generating new faces')
         response = generation_call(image_address=image_id, idx_face=idx_face, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)

diff --git a/main.py b/main.py
@@ -18,6 +18,7 @@
 
     # Consistent identity parameters
     parser.add_argument('--identity_filepath', help='Input identity file absolute path', type=str, default=None)
+    parser.add_argument('--identity_url', help='Input identity url, use only if no identity path was given', type=str, default=None)
     parser.add_argument('--identity_name', help='Use the face from the stored identities', default=None)
     parser.add_argument('--store_identity', help='Save the generated identity under the name pippo', action='store_true')
 
@@ -46,6 +47,7 @@
 
     # Consistent identity parameters
     IDENTITY_PATH = args.identity_filepath
+    IDENTITY_URL = args.identity_url
     IDENTITY_NAME = args.identity_name  # Default is None, otherwise a string of a stored name
     STORE_IDENTITY_FLAG = args.store_identity  # False if the new identity shall not be saved in the user profile, viceversa True
 
@@ -83,13 +85,15 @@
 
     if IDENTITY_PATH is not None:
         if os.path.exists(IDENTITY_PATH):
-            identity_image = open_image_from_path(IDENTITY_PATH)
             print(f'Using the input identity file located at: {IDENTITY_PATH}')
         else:
-            print('Wrong identity filepath, check again')
-            sys.exit()
+            print('Wrong identity path, check again')
     else:
-        identity_image = None
+        print('Identity path not assigned, check again')
+        if IDENTITY_URL is not None:
+            print(f'Using the input identity located at: {IDENTITY_URL}')
+        else:
+            print('Wrong identity url, check again')
 
     # log in
     TOKEN_DICTIONARY = start_call(EMAIL, PASSWORD)
@@ -100,7 +104,8 @@
             'INPUT_PATH': INPUT_PATH,
             'HAIR_FACTOR': HAIR_FACTOR,
             'CHANGE_ALL_FACES': CHANGE_ALL_FACES,
-            'IDENTITY_IMAGE': identity_image,
+            'IDENTITY_PATH': IDENTITY_PATH,
+            'IDENTITY_URL': IDENTITY_URL,
             'IDENTITY_NAME': IDENTITY_NAME,
             'STORE_IDENTITY_FLAG': STORE_IDENTITY_FLAG,
             'SEED': SEED,