fofr / cog-become-image

Turn anyone into another image
https://replicate.com/fofr/become-image
Other
216 stars 25 forks source link

Error by handling base64 webp image #10

Open dupasj opened 2 months ago

dupasj commented 2 months ago

Issue

I've developed a Cloudflare function that uploads an image and then runs your machine learning throught replicate. I've created a route to fetch the uploaded image. However, this route is not available in the development environment, so instead of a public image URL, I'll be sending a buffer of the image content.

When we set a blob/buffer/file in replicate, the input is transformed into base64 here:

totalBytes += buffer.byteLength;
if (totalBytes > MAX_DATA_URI_SIZE) {
  throw new Error(
    `Combined filesize of prediction ${totalBytes} bytes exceeds the 10MB limit for inline encoding. Please provide URLs instead.`
  );
}

const data = bytesToBase64(buffer);
mime = mime ?? "application/octet-stream";

return `data:${mime};base64,${data}`;

However, when a base64 image is sent, I encounter the following error in the replicate logs:

Traceback (most recent call last):
File "/root/.pyenv/versions/3.10.6/lib/python3.10/site-packages/cog/server/worker.py", line 217, in _predict
result = predict(**payload)
File "/src/predict.py", line 186, in predict
filename = self.handle_input_file(image, "image_of_face")
File "/src/predict.py", line 61, in handle_input_file
raise ValueError(f"Unsupported file type: {file_extension}")
ValueError: Unsupported file type: .bin

Upon inspecting your repository, it appears that the file extension is being determined based on the file content here:

file_extension = os.path.splitext(input_file)[1].lower()
if file_extension in [".jpg", ".jpeg"]:
    final_filename = f"{filename}.png"
    image = Image.open(input_file)

    try:
        for orientation in ExifTags.TAGS.keys():
            if ExifTags.TAGS[orientation] == "Orientation":
                break
        exif = dict(image._getexif().items())

        if exif[orientation] == 3:
            image = image.rotate(180, expand=True)
        elif exif[orientation] == 6:
            image = image.rotate(270, expand=True)
        elif exif[orientation] == 8:
            image = image.rotate(90, expand=True)
    except (KeyError, AttributeError):
        # EXIF data does not have orientation
        # Do not rotate
        pass

    image.save(os.path.join(INPUT_DIR, final_filename))
elif file_extension in [".png", ".webp"]:
    final_filename = filename + file_extension
    shutil.copy(input_file, os.path.join(INPUT_DIR, final_filename))
else:
    raise ValueError(f"Unsupported file type: {file_extension}")

return final_filename

Proposed Solution

Additional Context

I've only tested this with webp images (the base64 URL can be opened in my browser). I haven't managed to try base64 with jpeg and png images yet.