util/commons.py

import pprint
import random
import re
from io import BytesIO
from typing import Union

import boto3
import requests

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

from PIL import Image

black_list = {"alphonse mucha": "", "adolphe bouguereau": ""}
pp = pprint.PrettyPrinter(indent=4)
Avatar = [
    {
        "avatarName": "niomi",
        "codename": "1jMGp1kFkG",
        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/7_char_assets/niyomi_1jMGp1kFkG.png",
        "extraPrompt": "1jMGp1kFkG  girl",
    },
    {
        "avatarName": "riya",
        "codename": "vW6AUQtoaY",
        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/12_char_assets/riya_vW6AUQtoaY.png",
        "extraPrompt": "vW6AUQtoaY girl",
    },
    {
        "avatarName": "rajveer",
        "codename": "fSLF0OPkBw",
        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/12_character_assets/rajveer_fSLF0OPkBw.png",
        "extraPrompt": "fSLF0OPkBw guy",
    },
    {
        "avatarName": "bheem",
        "codename": "HL79CB3ODZ",
        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/7_char_assets/scene01001_1.png",
        "extraPrompt": "HL79CB3ODZ boy",
    },
    {
        "avatarName": "chutki",
        "codename": "SJ7JVIS9M7",
        "avatarImage": "https://comic-assets.s3.ap-south-1.amazonaws.com/7_char_assets/14.png",
        "extraPrompt": "SJ7JVIS9M7 girl",
    },
]

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 add_code_names(sentence):
    array_of_objects = Avatar

    for obj in array_of_objects:
        sentence = (
            re.sub(
                r"\b" + obj["avatarName"] + r"\b",
                obj["extraPrompt"],
                sentence,
                flags=re.IGNORECASE,
            )
            + " "
        )
    print(sentence)
    return sentence


def upload_images(images, processName: str, taskId: str):
    imageUrls = []
    for i, image in enumerate(images):
        # img_io = BytesIO()
        # image.save(img_io, "JPEG", quality=70)
        # img_io.seek(0)
        # key = "crecoAI/{}{}_{}.png".format(taskId, processName, i)
        # 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(
            "https://comic-assets.s3.ap-south-1.amazonaws.com/crecoAI/{}{}_{}.png".format(
                taskId, processName, i
            )
        )

    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)
    # s3.upload_fileobj(image, "comic-assets", out_path, ExtraArgs={"ACL": "public-read"})
    # image.close()


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


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://comic-assets.s3.ap-south-1.amazonaws.com/" + key


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