add unalign

app.py

@@ -82,11 +82,11 @@ def bytes_to_image(image: bytes) -> Image.Image:
     return image
 
 
-def edit_image(orig_image, edit_direction, edit_power, align, mask, progress=gr.Progress(track_tqdm=True)):
+def edit_image(orig_image, edit_direction, edit_power, align, mask, progress=gr.Progress(track_tqdm=True)): # output_align, output_unalign
     if edit_direction in DIRECTIONS_NAME_SWAP:
         edit_direction = DIRECTIONS_NAME_SWAP[edit_direction]
     if not orig_image:
-        return gr.update(visible=False), gr.update(visible=False), gr.update(value="Need to upload an input image ❗", visible=True)
+        return gr.update(visible=False), gr.update(visible=False), gr.update(value="Need to upload an input image ❗", visible=True), gr.update(visible=False), gr.update(visible=False)
 
     orig_image_bytes = get_bytes(orig_image)
     mask_bytes = get_bytes(mask)
@@ -103,11 +103,16 @@ def edit_image(orig_image, edit_direction, edit_power, align, mask, progress=gr.
         )
 
     if output.image == b"aligner error":
-        return gr.update(visible=False), gr.update(visible=False), gr.update(value="Face aligner can not find face in your image 😢 Try to upload another one", visible=True)
+        return gr.update(visible=False), gr.update(visible=False), gr.update(value="Face aligner can not find face in your image 😢 Try to upload another one", visible=True), gr.update(visible=False), gr.update(visible=False),
 
     output_edited = bytes_to_image(output.image)
     output_inv = bytes_to_image(output.inv_image)
-    return gr.update(value=output_edited, visible=True), gr.update(value=output_inv, visible=True), gr.update(visible=False)
+    if not align:
+        return gr.update(value=output_edited, visible=True), gr.update(value=output_inv, visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+    output_aligned = bytes_to_image(output.aligned)
+    output_unaligned = bytes_to_image(output.unaligned)
+    return gr.update(value=output_edited, visible=True), gr.update(value=output_inv, visible=True), gr.update(visible=False), gr.update(value=output_aligned, visible=True), gr.update(value=output_unaligned, visible=True)
 
 
 def edit_image_clip(orig_image, neutral_prompt, target_prompt, disentanglement, edit_power, align, mask, edit_method, progress=gr.Progress(track_tqdm=True)):
@@ -220,6 +225,9 @@ def get_demo():
                         btn_mask = gr.Button("Generate mask")
 
             with gr.Column():
+                with gr.Row():
+                    output_align = gr.Image(label="Alignet original image", visible=True)
+                    output_unalign = gr.Image(label="Unalinget editing result", visible=True)
                 with gr.Row():
                     output_inv = gr.Image(label="Inversion result", visible=True)
                     output_edit = gr.Image(label="Editing result", visible=True)
@@ -264,12 +272,12 @@ def get_demo():
         btn_predef.click(
             fn=edit_image, 
             inputs=[input_image, predef_editing_direction, predef_editing_power, align, mask],
-            outputs=[output_edit, output_inv, error_message]
+            outputs=[output_edit, output_inv, error_message, output_align, output_unalign]
         )
         btn_clip.click(
             fn=edit_image_clip, 
             inputs=[input_image, neutral_prompt, target_prompt, disentanglement, styleclip_editing_power, align, mask, edit_method],
-            outputs=[output_edit, output_inv, error_message]
+            outputs=[output_edit, output_inv, error_message, output_align, output_unalign,]
         )
         btn_mask.click(
             fn=get_mask, 

inference_pb2.py

@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
 
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0finference.proto\x12\tinference\"r\n\nSFERequest\x12\x12\n\norig_image\x18\x01 \x01(\x0c\x12\x11\n\tdirection\x18\x02 \x01(\t\x12\r\n\x05power\x18\x03 \x01(\x02\x12\x11\n\tuse_cache\x18\x04 \x01(\x08\x12\r\n\x05\x61lign\x18\x05 \x01(\x08\x12\x0c\n\x04mask\x18\x06 \x01(\x0c\"X\n\x0eSFERequestMask\x12\x12\n\norig_image\x18\x01 \x01(\x0c\x12\x10\n\x08trashold\x18\x02 \x01(\x02\x12\x11\n\tuse_cache\x18\x03 \x01(\x08\x12\r\n\x05\x61lign\x18\x04 \x01(\x08\"/\n\x0bSFEResponse\x12\r\n\x05image\x18\x01 \x01(\x0c\x12\x11\n\tinv_image\x18\x02 \x01(\x0c\"\x1f\n\x0fSFEResponseMask\x12\x0c\n\x04mask\x18\x01 \x01(\x0c\x32\x8b\x01\n\nSFEService\x12\x35\n\x04\x65\x64it\x12\x15.inference.SFERequest\x1a\x16.inference.SFEResponse\x12\x46\n\rgenerate_mask\x12\x19.inference.SFERequestMask\x1a\x1a.inference.SFEResponseMaskb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0finference.proto\x12\tinference\"r\n\nSFERequest\x12\x12\n\norig_image\x18\x01 \x01(\x0c\x12\x11\n\tdirection\x18\x02 \x01(\t\x12\r\n\x05power\x18\x03 \x01(\x02\x12\x11\n\tuse_cache\x18\x04 \x01(\x08\x12\r\n\x05\x61lign\x18\x05 \x01(\x08\x12\x0c\n\x04mask\x18\x06 \x01(\x0c\"X\n\x0eSFERequestMask\x12\x12\n\norig_image\x18\x01 \x01(\x0c\x12\x10\n\x08trashold\x18\x02 \x01(\x02\x12\x11\n\tuse_cache\x18\x03 \x01(\x08\x12\r\n\x05\x61lign\x18\x04 \x01(\x08\"S\n\x0bSFEResponse\x12\r\n\x05image\x18\x01 \x01(\x0c\x12\x11\n\tinv_image\x18\x02 \x01(\x0c\x12\x0f\n\x07\x61ligned\x18\x03 \x01(\x0c\x12\x11\n\tunaligned\x18\x04 \x01(\x0c\"\x1f\n\x0fSFEResponseMask\x12\x0c\n\x04mask\x18\x01 \x01(\x0c\x32\x8b\x01\n\nSFEService\x12\x35\n\x04\x65\x64it\x12\x15.inference.SFERequest\x1a\x16.inference.SFEResponse\x12\x46\n\rgenerate_mask\x12\x19.inference.SFERequestMask\x1a\x1a.inference.SFEResponseMaskb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -26,9 +26,9 @@ if not _descriptor._USE_C_DESCRIPTORS:
   _globals['_SFEREQUESTMASK']._serialized_start=146
   _globals['_SFEREQUESTMASK']._serialized_end=234
   _globals['_SFERESPONSE']._serialized_start=236
-  _globals['_SFERESPONSE']._serialized_end=283
-  _globals['_SFERESPONSEMASK']._serialized_start=285
-  _globals['_SFERESPONSEMASK']._serialized_end=316
-  _globals['_SFESERVICE']._serialized_start=319
-  _globals['_SFESERVICE']._serialized_end=458
+  _globals['_SFERESPONSE']._serialized_end=319
+  _globals['_SFERESPONSEMASK']._serialized_start=321
+  _globals['_SFERESPONSEMASK']._serialized_end=352
+  _globals['_SFESERVICE']._serialized_start=355
+  _globals['_SFESERVICE']._serialized_end=494
 # @@protoc_insertion_point(module_scope)