MuseV/musev/auto_prompt/attributes/human.py

from copy import deepcopy
import numpy as np
import random
import json

from .attributes import (
 MultiAttr2Text,
 AttriributeIsText,
 AttributeIsTextAndName,
 PresetMultiAttr2Text,
)
from .style import Style
from .render import Render
from . import AttrRegister


__all__ = [
 "Age",
 "Sex",
 "Singing",
 "Country",
 "Lighting",
 "Headwear",
 "Eyes",
 "Irises",
 "Hair",
 "Skin",
 "Face",
 "Smile",
 "Expression",
 "Clothes",
 "Nose",
 "Mouth",
 "Beard",
 "Necklace",
 "KeyWords",
 "InsightFace",
 "Caption",
 "Env",
 "Decoration",
 "Festival",
 "SpringHeadwear",
 "SpringClothes",
 "Animal",
]


@AttrRegister.register
class Sex(AttriributeIsText):
 name = "sex"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Headwear(AttriributeIsText):
 name = "headwear"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Expression(AttriributeIsText):
 name = "expression"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class KeyWords(AttriributeIsText):
 name = "keywords"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Singing(AttriributeIsText):
 def __init__(self, name: str = "singing") -> None:
 super().__init__(name)


@AttrRegister.register
class Country(AttriributeIsText):
 name = "country"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Clothes(AttriributeIsText):
 name = "clothes"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Age(AttributeIsTextAndName):
 name = "age"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)

 def __call__(self, attributes: str) -> str:
 if not isinstance(attributes, str):
 attributes = str(attributes)
 attributes = attributes.split(",")
 text = ", ".join(
 ["{}-year-old".format(attr) if attr != "" else "" for attr in attributes]
 )
 return text


@AttrRegister.register
class Eyes(AttributeIsTextAndName):
 name = "eyes"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Hair(AttributeIsTextAndName):
 name = "hair"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Background(AttributeIsTextAndName):
 name = "background"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Skin(AttributeIsTextAndName):
 name = "skin"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Face(AttributeIsTextAndName):
 name = "face"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Smile(AttributeIsTextAndName):
 name = "smile"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Nose(AttributeIsTextAndName):
 name = "nose"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Mouth(AttributeIsTextAndName):
 name = "mouth"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Beard(AttriributeIsText):
 name = "beard"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Necklace(AttributeIsTextAndName):
 name = "necklace"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Irises(AttributeIsTextAndName):
 name = "irises"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


@AttrRegister.register
class Lighting(AttributeIsTextAndName):
 name = "lighting"

 def __init__(self, name: str = None) -> None:
 super().__init__(name)


PresetPortraitAttributes = [
 Age,
 Sex,
 Singing,
 Country,
 Lighting,
 Headwear,
 Eyes,
 Irises,
 Hair,
 Skin,
 Face,
 Smile,
 Expression,
 Clothes,
 Nose,
 Mouth,
 Beard,
 Necklace,
 Style,
 KeyWords,
 Render,
]


class PortraitMultiAttr2Text(PresetMultiAttr2Text):
 preset_attributes = PresetPortraitAttributes

 def __init__(self, funcs: list = None, use_preset=True, name="portrait") -> None:
 super().__init__(funcs, use_preset, name)


@AttrRegister.register
class InsightFace(AttriributeIsText):
 name = "insight_face"
 face_render_dict = {
 "boy": "handsome,elegant",
 "girl": "gorgeous,kawaii,colorful",
 }
 key_words = "delicate face,beautiful eyes"

 def __call__(self, attributes: str) -> str:
 """将insight faces 检测的结果转化成prompt
 convert the results of insight faces detection to prompt
 Args:
 face_list (_type_): _description_

 Returns:
 _type_: _description_
 """
 attributes = json.loads(attributes)
 face_list = attributes["info"]
 if len(face_list) == 0:
 return ""

 if attributes["image_type"] == "body":
 for face in face_list:
 if "black" in face and face["black"]:
 return "african,dark skin"
 return ""

 gender_dict = {"girl": 0, "boy": 0}
 face_render_list = []
 black = False

 for face in face_list:
 if face["ratio"] < 0.02:
 continue

 if face["gender"] == 0:
 gender_dict["girl"] += 1
 face_render_list.append(self.face_render_dict["girl"])
 else:
 gender_dict["boy"] += 1
 face_render_list.append(self.face_render_dict["boy"])

 if "black" in face and face["black"]:
 black = True

 if len(face_render_list) == 0:
 return ""
 elif len(face_render_list) == 1:
 solo = True
 else:
 solo = False

 gender = ""
 for g, num in gender_dict.items():
 if num > 0:
 if gender:
 gender += ", "
 gender += "{}{}".format(num, g)
 if num > 1:
 gender += "s"

 face_render_list = ",".join(face_render_list)
 face_render_list = face_render_list.split(",")
 face_render = list(set(face_render_list))
 face_render.sort(key=face_render_list.index)
 face_render = ",".join(face_render)
 if gender_dict["girl"] == 0:
 face_render = "male focus," + face_render

 insightface_prompt = "{},{},{}".format(gender, face_render, self.key_words)

 if solo:
 insightface_prompt += ",solo"
 if black:
 insightface_prompt = "african,dark skin," + insightface_prompt

 return insightface_prompt


@AttrRegister.register
class Caption(AttriributeIsText):
 name = "caption"


@AttrRegister.register
class Env(AttriributeIsText):
 name = "env"
 envs_list = [
 "east asian architecture",
 "fireworks",
 "snow, snowflakes",
 "snowing, snowflakes",
 ]

 def __call__(self, attributes: str = None) -> str:
 if attributes != "" and attributes != " " and attributes is not None:
 return attributes
 else:
 return random.choice(self.envs_list)


@AttrRegister.register
class Decoration(AttriributeIsText):
 name = "decoration"

 def __init__(self, name: str = None) -> None:
 self.decoration_list = [
 "chinese knot",
 "flowers",
 "food",
 "lanterns",
 "red envelop",
 ]
 super().__init__(name)

 def __call__(self, attributes: str = None) -> str:
 if attributes != "" and attributes != " " and attributes is not None:
 return attributes
 else:
 return random.choice(self.decoration_list)


@AttrRegister.register
class Festival(AttriributeIsText):
 name = "festival"
 festival_list = ["new year"]

 def __init__(self, name: str = None) -> None:
 super().__init__(name)

 def __call__(self, attributes: str = None) -> str:
 if attributes != "" and attributes != " " and attributes is not None:
 return attributes
 else:
 return random.choice(self.festival_list)


@AttrRegister.register
class SpringHeadwear(AttriributeIsText):
 name = "spring_headwear"
 headwear_list = ["rabbit ears", "rabbit ears, fur hat"]

 def __call__(self, attributes: str = None) -> str:
 if attributes != "" and attributes != " " and attributes is not None:
 return attributes
 else:
 return random.choice(self.headwear_list)


@AttrRegister.register
class SpringClothes(AttriributeIsText):
 name = "spring_clothes"
 clothes_list = [
 "mittens,chinese clothes",
 "mittens,fur trim",
 "mittens,red scarf",
 "mittens,winter clothes",
 ]

 def __call__(self, attributes: str = None) -> str:
 if attributes != "" and attributes != " " and attributes is not None:
 return attributes
 else:
 return random.choice(self.clothes_list)


@AttrRegister.register
class Animal(AttriributeIsText):
 name = "animal"
 animal_list = ["rabbit", "holding rabbits"]

 def __call__(self, attributes: str = None) -> str:
 if attributes != "" and attributes != " " and attributes is not None:
 return attributes
 else:
 return random.choice(self.animal_list)