data/task.py

from enum import Enum
from typing import Union

import numpy as np


class TaskType(Enum):
    TEXT_TO_IMAGE = "GENERATE_AI_IMAGE"
    IMAGE_TO_IMAGE = "IMAGE_TO_IMAGE"
    POSE = "POSE"
    CANNY = "CANNY"
    REMOVE_BG = "REMOVE_BG"
    INPAINT = "INPAINT"
    UPSCALE_IMAGE = "UPSCALE_IMAGE"


class ModelType(Enum):
    REAL = 10000
    ANIME = 10001
    COMIC = 10002


class Task:
    def __init__(self, data):
        self.__data = data
        if data.get("seed", -1) == None or self.get_seed() == -1:
            self.__data["seed"] = np.random.randint(0, np.iinfo(np.int64).max)

    def get_taskId(self) -> str:
        return self.__data.get("task_id")

    def get_sourceId(self) -> str:
        return self.__data.get("source_id")

    def get_imageUrl(self) -> str:
        return self.__data.get("imageUrl")

    def get_prompt(self) -> str:
        return self.__data.get("prompt")

    def get_userId(self) -> str:
        return self.__data.get("userId", "")

    def get_email(self) -> str:
        return self.__data.get("email", "")

    def get_style(self) -> str:
        return self.__data.get("style", None)

    def get_iteration(self) -> float:
        return float(self.__data.get("iteration", 3.0))

    def get_modelType(self) -> ModelType:
        id = int(self.__data.get("model_id", 10000))
        return ModelType(id)

    def get_width(self) -> int:
        return int(self.__data.get("width", 512))

    def get_height(self) -> int:
        return int(self.__data.get("height", 512))

    def get_seed(self) -> int:
        return int(self.__data.get("seed", -1))

    def get_steps(self) -> int:
        return int(self.__data.get("steps", "75"))

    def get_type(self) -> Union[TaskType, None]:
        try:
            return TaskType(self.__data.get("task_type"))
        except ValueError:
            return None

    def get_maskImageUrl(self) -> str:
        return self.__data.get("maskImageUrl")

    def get_negative_prompt(self) -> str:
        return self.__data.get("negative_prompt", "")

    def is_prompt_engineering(self) -> bool:
        return self.__data.get("auto_mode", True)

    def get_raw(self) -> dict:
        return self.__data.copy()