1 files changed, 6 insertions, 16 deletions

diff --git a/src/png.c b/src/png.c
index 8eee7c7..23ff489 100755
--- a/src/png.c
+++ b/src/png.c
@@ -99,31 +99,21 @@ void write_png(FILE *file, void *image, int width, int height, flam3_img_comment
 	       PNG_COMPRESSION_TYPE_BASE,
 	       PNG_FILTER_TYPE_BASE);
 	       
-  /* Swap the bytes if we're doing 16bpc and on little-endian platform */       
-  if (2==bpc && testbe != htons(testbe)) {
-    unsigned short *im = (unsigned short *)image;
-    for (i=0; i<height*width*4; i++) {
-       im[i] = htons(im[i]);
-    }
-  }
-
   if (pngcom_enable==1)
 	  png_set_text(png_ptr, info_ptr, text, FLAM3_PNG_COM);
 
   png_write_info(png_ptr, info_ptr);
+
+  /* Must set this after the write_info */
+  if (2==bpc && testbe != htons(testbe)) {
+     png_set_swap(png_ptr);
+  }
+
   png_write_image(png_ptr, (png_bytepp) rows);
   png_write_end(png_ptr, info_ptr);
   png_destroy_write_struct(&png_ptr, &info_ptr);
   free(rows);
 
-  /* Swap back the bytes in case we're doing strips */       
-  if (2==bpc && testbe != htons(testbe)) {
-    unsigned short *im = (unsigned short *)image;
-    for (i=0; i<height*width*4; i++) {
-       im[i] = htons(im[i]);
-    }
-  }
-
 }
 
 #define SIG_CHECK_SIZE 8