File size: 2,204 Bytes
4adca93 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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()
|