added Stable Diffusion 3 support alongside ComfyUI configuration

this commit adds four environment variables:

- COMFYUI_CFG_SCALE
- COMFYUI_SAMPLER
- COMFYUI_SCHEDULER
- COMFYUI_SD3 (merely setting this at all will enable SD3 mode)
This commit is contained in:
John Karabudak
2024-06-16 03:19:04 -02:30
parent 9f7ef209fa
commit ea074fa9bf
3 changed files with 60 additions and 1 deletions

View File

@@ -190,7 +190,10 @@ class ImageGenerationPayload(BaseModel):
width: int
height: int
n: int = 1
cfg_scale: Optional[float] = None
sampler: Optional[str] = None
scheduler: Optional[str] = None
sd3: Optional[bool] = None
def comfyui_generate_image(
model: str, payload: ImageGenerationPayload, client_id, base_url
@@ -199,6 +202,18 @@ def comfyui_generate_image(
comfyui_prompt = json.loads(COMFYUI_DEFAULT_PROMPT)
if payload.cfg_scale:
comfyui_prompt["3"]["inputs"]["cfg"] = payload.cfg_scale
if payload.sampler:
comfyui_prompt["3"]["inputs"]["sampler"] = payload.sampler
if payload.scheduler:
comfyui_prompt["3"]["inputs"]["scheduler"] = payload.scheduler
if payload.sd3:
comfyui_prompt["5"]["class_type"] = "EmptySD3LatentImage"
comfyui_prompt["4"]["inputs"]["ckpt_name"] = model
comfyui_prompt["5"]["inputs"]["batch_size"] = payload.n
comfyui_prompt["5"]["inputs"]["width"] = payload.width