449 lines
17 KiB
Python
449 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any, TypedDict
|
|
|
|
from .model_info import FAMILY_DEFAULT_PRESET, get_family_defaults, get_sampling_preset, is_flux_family
|
|
|
|
|
|
LLS_TASK_CONTEXT_TYPE = "LLS_TASK_CONTEXT"
|
|
|
|
TASK_MODE_CHOICES = ["txt2img", "img2img", "inpaint", "outpaint", "upscale"]
|
|
TASK_CONTROLLER_MODEL_FAMILY_CHOICES = [
|
|
"auto",
|
|
"SD1.5",
|
|
"SDXL",
|
|
"FLUX",
|
|
"SD3",
|
|
"Qwen-Image",
|
|
"Z-Image",
|
|
"Hunyuan",
|
|
]
|
|
WORKFLOW_PRESET_CHOICES = ["simple", "standard", "advanced"]
|
|
QUALITY_PRESET_CHOICES = ["fast", "balanced", "quality"]
|
|
|
|
_RECOMMENDATION_KEYS = {
|
|
"recommended_width",
|
|
"recommended_height",
|
|
"recommended_steps",
|
|
"recommended_cfg",
|
|
"recommended_denoise",
|
|
"recommended_guidance",
|
|
"recommended_sampler",
|
|
"recommended_scheduler",
|
|
"latent_channels",
|
|
"downscale_ratio",
|
|
}
|
|
_TRIGGER_RECOMMENDATION_REFRESH_KEYS = {"model_family", "resolved_model_family", "task_mode", "quality_preset"}
|
|
_QUALITY_TO_SAMPLING_PRESET = {
|
|
"fast": "Fast",
|
|
"balanced": "Balanced",
|
|
"quality": "High Quality",
|
|
}
|
|
_UNKNOWN_FAMILY_DEFAULTS = {
|
|
"recommended_width": 1024,
|
|
"recommended_height": 1024,
|
|
"recommended_steps": 24,
|
|
"recommended_cfg": 4.0,
|
|
"recommended_denoise": 1.0,
|
|
"recommended_guidance": None,
|
|
"recommended_sampler": "euler",
|
|
"recommended_scheduler": "normal",
|
|
"latent_channels": 4,
|
|
"downscale_ratio": 8,
|
|
"text_encoder_type": "unknown",
|
|
"vae_type": "vae",
|
|
}
|
|
|
|
|
|
class TaskContext(TypedDict, total=False):
|
|
task_mode: str
|
|
model_family: str
|
|
resolved_model_family: str
|
|
workflow_preset: str
|
|
quality_preset: str
|
|
enable_upscale: bool
|
|
enable_controlnet: bool
|
|
enable_reference: bool
|
|
use_external_vae: bool
|
|
use_external_text_encoder: bool
|
|
recommended_width: int
|
|
recommended_height: int
|
|
recommended_steps: int
|
|
recommended_cfg: float
|
|
recommended_denoise: float
|
|
recommended_guidance: float | None
|
|
recommended_sampler: str
|
|
recommended_scheduler: str
|
|
checkpoint_name: str
|
|
vae_source: str
|
|
text_encoder_source: str
|
|
latent_channels: int
|
|
downscale_ratio: int
|
|
text_encoder_type: str
|
|
vae_type: str
|
|
supports_img2img: bool
|
|
supports_inpaint: bool
|
|
supports_controlnet: bool
|
|
supports_clip_skip: bool
|
|
supports_flux_guidance: bool
|
|
latent_source: str
|
|
final_width: int
|
|
final_height: int
|
|
batch_size: int
|
|
sampler_name: str
|
|
scheduler: str
|
|
steps: int
|
|
cfg: float
|
|
seed: int
|
|
denoise: float
|
|
guidance: float | None
|
|
decoded: bool
|
|
upscaled: bool
|
|
source: str
|
|
|
|
|
|
def _normalize_bool(value: Any, default: bool = False) -> bool:
|
|
if value in (None, ""):
|
|
return default
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, str):
|
|
return value.strip().lower() in {"1", "true", "yes", "on"}
|
|
return bool(value)
|
|
|
|
|
|
def _normalize_task_mode(value: Any) -> str:
|
|
normalized = str(value or "txt2img").strip().lower()
|
|
return normalized if normalized in TASK_MODE_CHOICES else "txt2img"
|
|
|
|
|
|
def _normalize_workflow_preset(value: Any) -> str:
|
|
normalized = str(value or "simple").strip().lower()
|
|
return normalized if normalized in WORKFLOW_PRESET_CHOICES else "simple"
|
|
|
|
|
|
def _normalize_quality_preset(value: Any) -> str:
|
|
normalized = str(value or "balanced").strip().lower()
|
|
if normalized in QUALITY_PRESET_CHOICES:
|
|
return normalized
|
|
return "balanced"
|
|
|
|
|
|
def _normalize_requested_family(value: Any) -> str:
|
|
normalized = str(value or "auto").strip()
|
|
lowered = normalized.lower()
|
|
if lowered in {"", "auto"}:
|
|
return "auto"
|
|
if normalized in {"SD15", "SD1.5"}:
|
|
return "SD1.5"
|
|
if normalized in {"SDXL", "SDXL_TURBO"}:
|
|
return "SDXL"
|
|
if normalized in {"FLUX", "FLUX_SCHNELL", "FLUX_DEV"}:
|
|
return "FLUX"
|
|
return normalized
|
|
|
|
|
|
def _normalize_resolved_family(value: Any, requested_family: str) -> str:
|
|
normalized = str(value or requested_family or "auto").strip()
|
|
if normalized in {"", "auto"}:
|
|
return requested_family or "auto"
|
|
if normalized in {"SD15", "SD1.5"}:
|
|
return "SD1.5"
|
|
if normalized in {"SDXL", "SDXL_TURBO"}:
|
|
return normalized
|
|
if normalized in {"FLUX", "FLUX_SCHNELL", "FLUX_DEV"}:
|
|
return normalized
|
|
return normalized
|
|
|
|
|
|
def _compat_family(resolved_family: str, requested_family: str) -> str:
|
|
if resolved_family and resolved_family not in {"auto", "FLUX"}:
|
|
return resolved_family
|
|
if requested_family == "FLUX":
|
|
return "FLUX_DEV"
|
|
if requested_family in {"auto", ""}:
|
|
return "SD1.5"
|
|
return requested_family
|
|
|
|
|
|
def _profile_family(family: str) -> str:
|
|
if family in {"auto", ""}:
|
|
return "SD1.5"
|
|
if family in {"SD3", "Qwen-Image", "Z-Image", "Hunyuan"}:
|
|
return "SDXL"
|
|
if family == "FLUX":
|
|
return "FLUX_DEV"
|
|
return family
|
|
|
|
|
|
def _denoise_for_task_mode(task_mode: str) -> float:
|
|
return {
|
|
"txt2img": 1.0,
|
|
"img2img": 0.5,
|
|
"inpaint": 0.6,
|
|
"outpaint": 0.6,
|
|
"upscale": 0.35,
|
|
}.get(task_mode, 1.0)
|
|
|
|
|
|
def _coerce_number(value: Any, caster, default: Any):
|
|
if value in (None, ""):
|
|
return default
|
|
try:
|
|
return caster(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _context_defaults_for_family(requested_family: str, resolved_family: str, task_mode: str, quality_preset: str) -> dict[str, Any]:
|
|
compat_family = _compat_family(resolved_family, requested_family)
|
|
profile_family = _profile_family(compat_family)
|
|
|
|
if profile_family in {"SD1.5", "SDXL", "SDXL_TURBO", "FLUX_SCHNELL", "FLUX_DEV"}:
|
|
family_defaults = get_family_defaults(profile_family)
|
|
preset_name = _QUALITY_TO_SAMPLING_PRESET.get(quality_preset, FAMILY_DEFAULT_PRESET)
|
|
preset = get_sampling_preset({"family": profile_family}, preset_name) or {
|
|
"steps": family_defaults["default_steps"],
|
|
"cfg": family_defaults["default_cfg"],
|
|
"guidance": family_defaults["default_guidance"],
|
|
"sampler_name": family_defaults["default_sampler"],
|
|
"scheduler": family_defaults["default_scheduler"],
|
|
"denoise": family_defaults["default_denoise"],
|
|
}
|
|
defaults = {
|
|
"recommended_width": int(family_defaults["default_width"]),
|
|
"recommended_height": int(family_defaults["default_height"]),
|
|
"recommended_steps": int(preset["steps"]),
|
|
"recommended_cfg": float(preset["cfg"]),
|
|
"recommended_denoise": float(_denoise_for_task_mode(task_mode)),
|
|
"recommended_guidance": None if preset.get("guidance") in ("", None) else float(preset["guidance"]),
|
|
"recommended_sampler": str(preset["sampler_name"]),
|
|
"recommended_scheduler": str(preset["scheduler"]),
|
|
"latent_channels": int(family_defaults["latent_channels"]),
|
|
"downscale_ratio": int(family_defaults["downscale_ratio"]),
|
|
"text_encoder_type": str(family_defaults["text_encoder_type"]),
|
|
"vae_type": "ae" if is_flux_family(profile_family) else "vae",
|
|
}
|
|
return defaults
|
|
|
|
defaults = dict(_UNKNOWN_FAMILY_DEFAULTS)
|
|
defaults["recommended_denoise"] = float(_denoise_for_task_mode(task_mode))
|
|
return defaults
|
|
|
|
|
|
def _context_capabilities(compat_family: str, latent_channels: Any) -> dict[str, bool]:
|
|
is_flux = compat_family in {"FLUX", "FLUX_SCHNELL", "FLUX_DEV"}
|
|
has_latent = latent_channels not in (None, "", 0)
|
|
return {
|
|
"supports_img2img": has_latent,
|
|
"supports_inpaint": has_latent,
|
|
"supports_controlnet": compat_family in {"SD1.5", "SDXL", "SDXL_TURBO"},
|
|
"supports_clip_skip": not is_flux,
|
|
"supports_flux_guidance": is_flux,
|
|
}
|
|
|
|
|
|
def _parse_jsonish_dict(value: Any) -> dict[str, Any]:
|
|
if isinstance(value, dict):
|
|
return dict(value)
|
|
if not isinstance(value, str):
|
|
return {}
|
|
stripped = value.strip()
|
|
if not stripped:
|
|
return {}
|
|
try:
|
|
loaded = json.loads(stripped)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
return loaded if isinstance(loaded, dict) else {}
|
|
|
|
|
|
def _finalize_task_context(raw_context: dict[str, Any]) -> dict[str, Any]:
|
|
context = dict(raw_context)
|
|
requested_family = _normalize_requested_family(
|
|
context.get("model_family") or context.get("requested_model_family") or context.get("family")
|
|
)
|
|
resolved_family = _normalize_resolved_family(context.get("resolved_model_family") or context.get("family"), requested_family)
|
|
if resolved_family == "auto":
|
|
resolved_family = _compat_family(resolved_family, requested_family)
|
|
task_mode = _normalize_task_mode(context.get("task_mode"))
|
|
workflow_preset = _normalize_workflow_preset(context.get("workflow_preset"))
|
|
quality_preset = _normalize_quality_preset(context.get("quality_preset"))
|
|
|
|
defaults = _context_defaults_for_family(requested_family, resolved_family, task_mode, quality_preset)
|
|
compat_family = _compat_family(resolved_family, requested_family)
|
|
capabilities = _context_capabilities(compat_family, context.get("latent_channels", defaults.get("latent_channels")))
|
|
|
|
context["task_mode"] = task_mode
|
|
context["model_family"] = requested_family
|
|
context["resolved_model_family"] = resolved_family
|
|
context["workflow_preset"] = workflow_preset
|
|
context["quality_preset"] = quality_preset
|
|
context["enable_upscale"] = _normalize_bool(context.get("enable_upscale"), True)
|
|
context["enable_controlnet"] = _normalize_bool(context.get("enable_controlnet"), False)
|
|
context["enable_reference"] = _normalize_bool(context.get("enable_reference"), False)
|
|
context["use_external_vae"] = _normalize_bool(context.get("use_external_vae"), False)
|
|
context["use_external_text_encoder"] = _normalize_bool(context.get("use_external_text_encoder"), False)
|
|
|
|
for key, default_value in defaults.items():
|
|
if key in {"recommended_cfg", "recommended_denoise", "recommended_guidance"}:
|
|
context[key] = None if default_value is None else _coerce_number(context.get(key), float, float(default_value))
|
|
continue
|
|
if key in {"recommended_steps", "recommended_width", "recommended_height", "latent_channels", "downscale_ratio"}:
|
|
context[key] = _coerce_number(context.get(key), int, int(default_value))
|
|
continue
|
|
context[key] = context.get(key, default_value)
|
|
|
|
if context["recommended_guidance"] is None and defaults.get("recommended_guidance") is not None:
|
|
context["recommended_guidance"] = float(defaults["recommended_guidance"])
|
|
|
|
context["checkpoint_name"] = str(context.get("checkpoint_name") or context.get("model_name") or "")
|
|
context["vae_source"] = str(context.get("vae_source") or ("external" if context["use_external_vae"] else "auto"))
|
|
context["text_encoder_source"] = str(
|
|
context.get("text_encoder_source") or ("external" if context["use_external_text_encoder"] else "auto")
|
|
)
|
|
context["text_encoder_type"] = str(context.get("text_encoder_type") or defaults["text_encoder_type"])
|
|
context["vae_type"] = str(context.get("vae_type") or defaults["vae_type"])
|
|
|
|
for key, default_value in capabilities.items():
|
|
context[key] = _normalize_bool(context.get(key), default_value)
|
|
|
|
context["family"] = compat_family
|
|
context["default_width"] = int(context["recommended_width"])
|
|
context["default_height"] = int(context["recommended_height"])
|
|
context["default_steps"] = int(context["recommended_steps"])
|
|
context["default_cfg"] = float(context["recommended_cfg"])
|
|
context["default_denoise"] = float(context["recommended_denoise"])
|
|
context["default_guidance"] = context["recommended_guidance"]
|
|
context["default_sampler"] = str(context["recommended_sampler"])
|
|
context["default_scheduler"] = str(context["recommended_scheduler"])
|
|
context["guidance"] = (
|
|
None
|
|
if context.get("guidance") in (None, "")
|
|
else _coerce_number(context.get("guidance"), float, context["recommended_guidance"])
|
|
)
|
|
if context["guidance"] is None:
|
|
context["guidance"] = context["recommended_guidance"]
|
|
context["is_flux"] = _normalize_bool(context.get("is_flux"), is_flux_family(compat_family))
|
|
context["is_turbo"] = _normalize_bool(context.get("is_turbo"), compat_family == "SDXL_TURBO")
|
|
|
|
if "final_width" in context:
|
|
context["final_width"] = _coerce_number(context.get("final_width"), int, context["recommended_width"])
|
|
if "final_height" in context:
|
|
context["final_height"] = _coerce_number(context.get("final_height"), int, context["recommended_height"])
|
|
if "batch_size" in context:
|
|
context["batch_size"] = _coerce_number(context.get("batch_size"), int, 1)
|
|
if "steps" in context:
|
|
context["steps"] = _coerce_number(context.get("steps"), int, context["recommended_steps"])
|
|
if "cfg" in context:
|
|
context["cfg"] = _coerce_number(context.get("cfg"), float, context["recommended_cfg"])
|
|
if "seed" in context:
|
|
context["seed"] = _coerce_number(context.get("seed"), int, -1)
|
|
if "denoise" in context:
|
|
context["denoise"] = _coerce_number(context.get("denoise"), float, context["recommended_denoise"])
|
|
return context
|
|
|
|
|
|
def parse_task_context(task_context: dict[str, Any] | str | None) -> dict[str, Any]:
|
|
raw = _parse_jsonish_dict(task_context)
|
|
if isinstance(task_context, dict):
|
|
raw = dict(task_context)
|
|
|
|
if "resolved_model_family" not in raw and raw.get("family"):
|
|
raw["resolved_model_family"] = raw.get("family")
|
|
if "model_family" not in raw and raw.get("family"):
|
|
raw["model_family"] = raw.get("family")
|
|
|
|
if "recommended_width" not in raw and raw.get("default_width") is not None:
|
|
raw["recommended_width"] = raw.get("default_width")
|
|
if "recommended_height" not in raw and raw.get("default_height") is not None:
|
|
raw["recommended_height"] = raw.get("default_height")
|
|
if "recommended_steps" not in raw and raw.get("default_steps") is not None:
|
|
raw["recommended_steps"] = raw.get("default_steps")
|
|
if "recommended_cfg" not in raw and raw.get("default_cfg") is not None:
|
|
raw["recommended_cfg"] = raw.get("default_cfg")
|
|
if "recommended_denoise" not in raw and raw.get("default_denoise") is not None:
|
|
raw["recommended_denoise"] = raw.get("default_denoise")
|
|
if "recommended_guidance" not in raw and raw.get("default_guidance") is not None:
|
|
raw["recommended_guidance"] = raw.get("default_guidance")
|
|
if "recommended_sampler" not in raw and raw.get("default_sampler") is not None:
|
|
raw["recommended_sampler"] = raw.get("default_sampler")
|
|
if "recommended_scheduler" not in raw and raw.get("default_scheduler") is not None:
|
|
raw["recommended_scheduler"] = raw.get("default_scheduler")
|
|
|
|
return _finalize_task_context(raw)
|
|
|
|
|
|
def create_task_context(
|
|
task_mode: str = "txt2img",
|
|
model_family: str = "auto",
|
|
workflow_preset: str = "simple",
|
|
quality_preset: str = "balanced",
|
|
enable_upscale: bool = False,
|
|
enable_controlnet: bool = False,
|
|
enable_reference: bool = False,
|
|
use_external_vae: bool = False,
|
|
use_external_text_encoder: bool = False,
|
|
**updates: Any,
|
|
) -> dict[str, Any]:
|
|
context = {
|
|
"task_mode": task_mode,
|
|
"model_family": model_family,
|
|
"workflow_preset": workflow_preset,
|
|
"quality_preset": quality_preset,
|
|
"enable_upscale": enable_upscale,
|
|
"enable_controlnet": enable_controlnet,
|
|
"enable_reference": enable_reference,
|
|
"use_external_vae": use_external_vae,
|
|
"use_external_text_encoder": use_external_text_encoder,
|
|
"source": "LLS Task Controller",
|
|
}
|
|
context.update({key: value for key, value in updates.items() if value is not None})
|
|
return _finalize_task_context(context)
|
|
|
|
|
|
def infer_default_params(task_context: dict[str, Any] | str | None) -> dict[str, Any]:
|
|
context = parse_task_context(task_context)
|
|
return {
|
|
key: context.get(key)
|
|
for key in (
|
|
"recommended_width",
|
|
"recommended_height",
|
|
"recommended_steps",
|
|
"recommended_cfg",
|
|
"recommended_denoise",
|
|
"recommended_guidance",
|
|
"recommended_sampler",
|
|
"recommended_scheduler",
|
|
"latent_channels",
|
|
"downscale_ratio",
|
|
)
|
|
}
|
|
|
|
|
|
def update_task_context(task_context: dict[str, Any] | str | None = None, **updates: Any) -> dict[str, Any]:
|
|
context = parse_task_context(task_context)
|
|
refresh_recommendations = any(key in _TRIGGER_RECOMMENDATION_REFRESH_KEYS for key in updates)
|
|
if refresh_recommendations:
|
|
for key in _RECOMMENDATION_KEYS:
|
|
if key not in updates:
|
|
context.pop(key, None)
|
|
for key, value in updates.items():
|
|
if value is not None:
|
|
context[key] = value
|
|
return _finalize_task_context(context)
|
|
|
|
|
|
def merge_task_context(*contexts: dict[str, Any] | str | None) -> dict[str, Any]:
|
|
merged: dict[str, Any] = {}
|
|
for context in contexts:
|
|
if context in (None, ""):
|
|
continue
|
|
merged.update(_parse_jsonish_dict(context) if not isinstance(context, dict) else dict(context))
|
|
return _finalize_task_context(merged)
|
|
|
|
|
|
def task_context_to_json(task_context: dict[str, Any] | str | None) -> str:
|
|
return json.dumps(parse_task_context(task_context), ensure_ascii=True, sort_keys=True)
|