internals/util/commons.py

import base64
import json
import os
import pprint
import random
import re
from io import BytesIO
from pathlib import Path
from typing import Any, Optional, Union

import boto3
import requests

from internals.util.config import api_endpoint, api_headers

s3 = boto3.client("s3")
import io
import urllib.request

from PIL import Image

black_list = {"alphonse mucha": "", "adolphe bouguereau": ""}
pp = pprint.PrettyPrinter(indent=4)

webhook_url = (
    "https://hooks.slack.com/services/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"
)
error_webhook = (
    "https://hooks.slack.com/services/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"
)

characterSheets = [
    "character+sheets/1.1.png",
    "character+sheets/10.1.png",
    "character+sheets/11.1.png",
    "character+sheets/12.1.png",
    "character+sheets/13.1.png",
    "character+sheets/14.1.png",
    "character+sheets/16.1.png",
    "character+sheets/17.1.png",
    "character+sheets/18.1.png",
    "character+sheets/19.1.png",
    "character+sheets/2.1.png",
    "character+sheets/20.1.png",
    "character+sheets/21.1.png",
    "character+sheets/22.1.png",
    "character+sheets/23.1.png",
    "character+sheets/24.1.png",
    "character+sheets/25.1.png",
    "character+sheets/26.1.png",
    "character+sheets/27.1.png",
    "character+sheets/28.1.png",
    "character+sheets/29.1.png",
    "character+sheets/3.1.png",
    "character+sheets/30.1.png",
    "character+sheets/31.1.png",
    "character+sheets/32.1.png",
    "character+sheets/33.1.png",
    "character+sheets/34.1.png",
    "character+sheets/35.1.png",
    "character+sheets/36.1.png",
    "character+sheets/38.1.png",
    "character+sheets/39.1.png",
    "character+sheets/4.1.png",
    "character+sheets/40.1.png",
    "character+sheets/42.1.png",
    "character+sheets/43.1.png",
    "character+sheets/44.1.png",
    "character+sheets/45.1.png",
    "character+sheets/46.1.png",
    "character+sheets/47.1.png",
    "character+sheets/48.1.png",
    "character+sheets/49.1.png",
    "character+sheets/5.1.png",
    "character+sheets/50.1.png",
    "character+sheets/51.1.png",
    "character+sheets/52.1.png",
    "character+sheets/53.1.png",
    "character+sheets/54.1.png",
    "character+sheets/55.1.png",
    "character+sheets/56.1.png",
    "character+sheets/57.1.png",
    "character+sheets/58.1.png",
    "character+sheets/59.1.png",
    "character+sheets/60.1.png",
    "character+sheets/61.1.png",
    "character+sheets/62.1.png",
    "character+sheets/63.1.png",
    "character+sheets/64.1.png",
    "character+sheets/65.1.png",
    "character+sheets/66.1.png",
    "character+sheets/7.1.png",
    "character+sheets/8.1.png",
    "character+sheets/9.1.png",
]


def upload_images(images, processName: str, taskId: str):
    imageUrls = []
    for i, image in enumerate(images):
        img_io = BytesIO()
        image.save(img_io, format="PNG")
        img_io.seek(0)
        key = "crecoAI/{}{}_{}.png".format(taskId, processName, i)
        res = requests.post(
            api_endpoint()
            + "/autodraft-content/v1.0/upload/crecoai-assets-2?fileName="
            + "{}{}_{}.png".format(taskId, processName, i),
            headers=api_headers(),
            files={"file": ("image.png", img_io, "image/png")},
        )
        imgUrl = res.json()["imageUrl"]
        # t = s3.put_object(
        #     Bucket="comic-assets", Key=key, Body=img_io.getvalue(), ACL="public-read"
        # )
        # print("uploading done to s3", key, t)
        imageUrls.append(imgUrl)

    print({"promptImages": imageUrls})

    return imageUrls


def upload_image(image: Union[Image.Image, BytesIO], out_path):
    if type(image) is Image.Image:
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        image = buffer

    image.seek(0)
    print(
        api_endpoint()
        + "/autodraft-content/v1.0/upload/crecoai-assets-2?fileName="
        + str(out_path).replace("crecoAI/", ""),
    )
    res = requests.post(
        api_endpoint()
        + "/autodraft-content/v1.0/upload/crecoai-assets-2?fileName="
        + str(out_path).replace("crecoAI/", ""),
        headers=api_headers(),
        files={"file": ("image.png", image, "image/png")},
    )
    print(res.text)
    imgUrl = res.json()["imageUrl"]
    # s3.upload_fileobj(image, "comic-assets", out_path, ExtraArgs={"ACL": "public-read"})
    image.close()

    image_url = imgUrl
    print({"promptImages": image_url})

    return image_url


def download_image(url, mode="RGB") -> Image.Image:
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert(mode)


def download_file(url, out_path: Path):
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(out_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)


def base64_to_image(base64_string):
    imgdata = base64.b64decode(base64_string)
    return Image.open(io.BytesIO(imgdata)).convert("RGB")


def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue())
    return img_str.decode("ascii")


def pickPoses():
    random_images = random.sample(characterSheets, 4)
    poses = []
    prefix = "https://comic-assets.s3.ap-south-1.amazonaws.com/"

    # Use list comprehension to add prefix to all elements in the array
    random_images_with_prefix = [prefix + img for img in random_images]

    print(random_images_with_prefix)
    for imageUrl in random_images_with_prefix:
        # Download and resize the image
        init_image = download_image(imageUrl).resize((512, 512))

        # Open the pose image
        imageUrlPose = imageUrl
        # print(imageUrl)
        input_image_bytes = read_url(imageUrlPose)
        # print(input_image_bytes)
        pose_image = Image.open(io.BytesIO(input_image_bytes)).convert("RGB")
        # print(pose_image)
        pose_image = pose_image.resize((512, 512))
        # print(pose_image)
        # Append the result to the poses array
        poses.append(pose_image)

    return poses


def construct_default_s3_url(key):
    return "https://assets.autodraft.in/" + key


def safe_index(array, index) -> Optional[Any]:
    if not array:
        return None
    if index < 0:
        return None
    if index >= len(array):
        return None
    return array[index]


def read_url(url: str):
    with urllib.request.urlopen(url) as u:
        return u.read()


def disable_safety_checker(pipe):
    def dummy(images, **kwargs):
        return images, False

    pipe.safety_checker = None