Update api.py

Author: zhchaoxing

modules/api/api.py

@@ -233,7 +233,7 @@ def encode_pil_to_base64(image):
             # })
 
             # Convert dict to bytes
-            exif_bytes = piexif.dump(set_img_exif_dict()) # Doesn't work here even though the exif is confirmed to have been added here
+            exif_bytes = piexif.dump(set_img_exif_dict())
 
             if opts.samples_format.lower() in ("jpg", "jpeg"):
                 image.save(output_bytes, format="JPEG", exif=exif_bytes, quality=opts.jpeg_quality)
@@ -249,7 +249,7 @@ def encode_pil_to_base64(image):
 
     return base64.b64encode(bytes_data)
 
-def export_pil_to_bytes(image, quality):
+def export_pil_to_bytes(image, quality, image_id=""):
     with io.BytesIO() as output_bytes:
 
         if opts.samples_format.lower() == 'png':
@@ -262,10 +262,13 @@ def export_pil_to_bytes(image, quality):
             image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=quality if quality else opts.jpeg_quality)
 
         elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
-            parameters = image.info.get('parameters', None)
-            exif_bytes = piexif.dump({
-                "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
-            })
+            # parameters = image.info.get('parameters', None)
+            # exif_bytes = piexif.dump({
+            #     "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
+            # })
+            # Convert dict to bytes
+            exif_bytes = piexif.dump(set_img_exif_dict(image_id=image_id))
+
             if opts.samples_format.lower() in ("jpg", "jpeg"):
                 image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=quality if quality else opts.jpeg_quality)
             else:
@@ -953,25 +956,25 @@ def post_invocations(self, b64images, quality):
                 image_id = datetime.datetime.now().strftime(f"%Y%m%d%H%M%S-{uuid.uuid4()}")
                 suffix = opts.samples_format.lower()
                 img = decode_to_image(b64image)
-                if suffix in ("jpg", "jpeg"):
-                    # print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} ***&&& post_invocations_before")
-                    # print(img.getexif())  # Print the EXIF data
-
-                    # Save the image with the EXIF data
-                    exif_bytes = piexif.dump(set_img_exif_dict(image_id=image_id))
-                    with io.BytesIO() as output_bytes:
-                        img.save(output_bytes, format="JPEG", exif=exif_bytes, quality=opts.jpeg_quality)
-                        output_bytes.seek(0)  # Go back to the start of the BytesIO object
-                        img_after = Image.open(output_bytes)
-                        print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} GenID@{user_input_data['generation_id']} ***&&& post_invocations_after")
-                        print(img_after.getexif())  # Print the EXIF data
-                    
-                bytes_data = export_pil_to_bytes(img, quality)
+                # if suffix in ("jpg", "jpeg"):
+                #     # print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} ***&&& post_invocations_before")
+                #     # print(img.getexif())  # Print the EXIF data
+                #
+                #     # Save the image with the EXIF data
+                #     exif_bytes = piexif.dump(set_img_exif_dict(image_id=image_id))
+                #     with io.BytesIO() as output_bytes:
+                #         img.save(output_bytes, format="JPEG", exif=exif_bytes, quality=opts.jpeg_quality)
+                #         output_bytes.seek(0)  # Go back to the start of the BytesIO object
+                #         img_after = Image.open(output_bytes)
+                #         print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} GenID@{user_input_data['generation_id']} ***&&& post_invocations_after")
+                #         print(img_after.getexif())  # Print the EXIF data
+                print(f"log@{datetime.datetime.now().strftime(f'%Y%m%d%H%M%S')} GenID@{user_input_data['generation_id']} image_id@{image_id}")
+                bytes_data = export_pil_to_bytes(img, quality, image_id=image_id)
                 # bytes_data = export_pil_to_bytes(decode_to_image(b64image), quality)
                 shared.s3_client.put_object(
                     Body=bytes_data,
                     Bucket=bucket,
-                    Key=f'{key}/{image_id}.{suffix}',
+                    Key=f'{key}/{image_id}.{suffix}' #,
                     # ExtraArgs={
                     #     "ContentType": f"image/{suffix}"
                     # }