File size: 2,400 Bytes
37112ef |
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 |
from typing import List, Optional, Dict, Any
from pydantic import BaseModel, field_validator
from PIL import Image
from config import Config as appConfig
class ControlNetReq(BaseModel):
controlnets: List[str] # ["canny", "tile", "depth", "scribble"]
control_images: List[Image.Image]
controlnet_conditioning_scale: List[float]
class Config:
arbitrary_types_allowed=True
class BaseReq(BaseModel):
model: str = ""
prompt: str = ""
negative_prompt: Optional[str] = None
fast_generation: Optional[bool] = True
loras: Optional[list] = []
embeddings: Optional[list] = None
resize_mode: Optional[str] = "resize_and_fill" # resize_only, crop_and_resize, resize_and_fill
scheduler: Optional[str] = "euler_fl"
height: int = 1024
width: int = 1024
num_images_per_prompt: int = 1
num_inference_steps: int = 8
clip_skip: Optional[int] = None
guidance_scale: float = 3.5
seed: Optional[int] = 0
refiner: bool = False
vae: bool = True
controlnet_config: Optional[ControlNetReq] = None
custom_addons: Optional[Dict[Any, Any]] = None
class Config:
arbitrary_types_allowed=True
@field_validator('model', 'negative_prompt', 'embeddings', 'clip_skip', 'controlnet_config')
def check_model(cls, values):
for m in appConfig.IMAGES_MODELS:
if m.get('repo_id') == values.get('model'):
loader = m.get('loader')
if loader == "flux" and values.get('negative_prompt'):
raise ValueError("Negative prompt is not supported for Flux models.")
if loader == "flux" and values.get('embeddings'):
raise ValueError("Embeddings are not supported for Flux models.")
if loader == "flux" and values.get('clip_skip'):
raise ValueError("Clip skip is not supported for Flux models.")
if loader == "flux" and values.get('controlnet_config'):
if "scribble" in values.get('controlnet_config').controlnets:
raise ValueError("Scribble is not supported for Flux models.")
return values
class BaseImg2ImgReq(BaseReq):
image: Image.Image
strength: float = 1.0
class Config:
arbitrary_types_allowed=True
class BaseInpaintReq(BaseImg2ImgReq):
mask_image: Image.Image
class Config:
arbitrary_types_allowed=True
|