dynamic checkpoint presets with orientation, correct VAEs, and auto-resolved sampler/scheduler/cfg/steps
This commit is contained in:
@@ -0,0 +1,11 @@
|
|||||||
|
{
|
||||||
|
"prompt": "a cat on a windowsill",
|
||||||
|
"width": 1024,
|
||||||
|
"height": 1024,
|
||||||
|
"steps": 20,
|
||||||
|
"cfg": 7.0,
|
||||||
|
"seed": -1,
|
||||||
|
"sampler": "euler",
|
||||||
|
"scheduler": "normal",
|
||||||
|
"lora_strength": 0.8
|
||||||
|
}
|
||||||
+98
-15
@@ -8,6 +8,7 @@ from importlib.metadata import version
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Annotated, Any
|
from typing import Annotated, Any
|
||||||
|
|
||||||
|
import click
|
||||||
import typer
|
import typer
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
from rich.table import Table
|
from rich.table import Table
|
||||||
@@ -769,37 +770,116 @@ def serve(
|
|||||||
# =============================================================================
|
# =============================================================================
|
||||||
|
|
||||||
|
|
||||||
@app.command()
|
@app.command(context_settings={"allow_extra_args": False})
|
||||||
def generate( # noqa: PLR0915
|
def generate( # noqa: PLR0915
|
||||||
prompt: Annotated[str, typer.Argument(help="Positive prompt text")],
|
ctx: typer.Context,
|
||||||
|
prompt: Annotated[str | None, typer.Argument(help="Positive prompt text", show_default=False)] = None,
|
||||||
model: Annotated[str | None, typer.Option("-m", "--model", help="Checkpoint model name")] = None,
|
model: Annotated[str | None, typer.Option("-m", "--model", help="Checkpoint model name")] = None,
|
||||||
width: Annotated[int, typer.Option("-W", "--width", help="Image width")] = 1024,
|
width: Annotated[int | None, typer.Option("-W", "--width", help="Image width (auto from checkpoint)")] = None,
|
||||||
height: Annotated[int, typer.Option("-H", "--height", help="Image height")] = 1024,
|
height: Annotated[int | None, typer.Option("-H", "--height", help="Image height (auto from checkpoint)")] = None,
|
||||||
steps: Annotated[int, typer.Option("--steps", help="Sampling steps")] = 20,
|
steps: Annotated[int | None, typer.Option("--steps", help="Sampling steps (auto from checkpoint)")] = None,
|
||||||
cfg: Annotated[float, typer.Option("--cfg", help="CFG scale")] = 7.0,
|
cfg: Annotated[float | None, typer.Option("--cfg", help="CFG scale (auto from checkpoint)")] = None,
|
||||||
seed: Annotated[int, typer.Option("--seed", "-s", help="Random seed (-1 for random)")] = -1,
|
seed: Annotated[int, typer.Option("--seed", "-s", help="Random seed (-1 for random)")] = -1,
|
||||||
sampler: Annotated[str, typer.Option("--sampler", help="Sampler name")] = "euler",
|
sampler: Annotated[str | None, typer.Option("--sampler", help="Sampler name (auto from checkpoint)")] = None,
|
||||||
scheduler: Annotated[str, typer.Option("--scheduler", help="Scheduler name")] = "normal",
|
scheduler: Annotated[str | None, typer.Option("--scheduler", help="Scheduler name (auto from checkpoint)")] = None,
|
||||||
vae: Annotated[str | None, typer.Option("--vae", help="VAE model name")] = None,
|
vae: Annotated[str | None, typer.Option("--vae", help="VAE model name (auto from checkpoint)")] = None,
|
||||||
|
orientation: Annotated[str, typer.Option("-O", "--orientation", help="Resolution: square, portrait, landscape")] = "square",
|
||||||
lora: Annotated[str | None, typer.Option("-l", "--lora", help="LoRA model name")] = None,
|
lora: Annotated[str | None, typer.Option("-l", "--lora", help="LoRA model name")] = None,
|
||||||
lora_strength: Annotated[float, typer.Option("--lora-strength", help="LoRA strength")] = 0.8,
|
lora_strength: Annotated[float, typer.Option("--lora-strength", help="LoRA strength")] = 0.8,
|
||||||
negative: Annotated[str, typer.Option("-n", "--negative-prompt", help="Negative prompt")] = "",
|
negative: Annotated[str, typer.Option("-n", "--negative-prompt", help="Negative prompt")] = "",
|
||||||
output: Annotated[Path | None, typer.Option("-o", "--output", help="Save path (default: current dir)")] = None,
|
output: Annotated[Path | None, typer.Option("-o", "--output", help="Save path (default: current dir)")] = None,
|
||||||
remote: Annotated[str | None, typer.Option("-r", "--remote", help="Remote server name or URL")] = None,
|
remote: Annotated[str | None, typer.Option("-r", "--remote", help="Remote server name or URL")] = None,
|
||||||
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
|
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
|
||||||
|
json_input: Annotated[
|
||||||
|
str | None, typer.Option("--input", "-I", help="JSON params (keys match CLI options)")
|
||||||
|
] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Generate an image using text-to-image.
|
"""Generate an image using text-to-image.
|
||||||
|
|
||||||
Calls ComfyUI directly when local, or the remote tensors API when --remote is given.
|
Calls ComfyUI directly when local, or the remote tensors API when --remote is given.
|
||||||
|
Accepts --input with a JSON object whose keys match CLI option names. CLI flags override JSON values.
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
tsr generate "a cat on a windowsill"
|
tsr generate "a cat on a windowsill"
|
||||||
tsr generate "portrait photo" -m "flux1-dev-fp8.safetensors" --steps 30
|
tsr generate "portrait photo" -m "flux1-dev-fp8.safetensors" --steps 30
|
||||||
tsr generate "cyberpunk city" -o output.png
|
tsr generate "cyberpunk city" -o output.png
|
||||||
tsr generate "landscape" --remote junkpile
|
tsr generate "landscape" --remote junkpile
|
||||||
|
tsr generate --input '{"prompt": "a mech", "model": "flux1-dev-fp8.safetensors", "steps": 30}'
|
||||||
"""
|
"""
|
||||||
import random as rng # noqa: PLC0415
|
import random as rng # noqa: PLC0415
|
||||||
|
|
||||||
|
# ---- JSON input merging ----
|
||||||
|
if json_input is not None:
|
||||||
|
# Support file paths and raw JSON strings
|
||||||
|
json_path = Path(json_input)
|
||||||
|
if json_path.is_file():
|
||||||
|
json_text = json_path.read_text()
|
||||||
|
elif json_input.lstrip().startswith("{"):
|
||||||
|
json_text = json_input
|
||||||
|
else:
|
||||||
|
console.print(f"[red]Not a JSON string or file:[/red] {json_input}")
|
||||||
|
raise typer.Exit(1)
|
||||||
|
|
||||||
|
try:
|
||||||
|
ji = json.loads(json_text)
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
console.print(f"[red]Invalid JSON input:[/red] {e}")
|
||||||
|
raise typer.Exit(1) from e
|
||||||
|
|
||||||
|
if not isinstance(ji, dict):
|
||||||
|
console.print("[red]JSON input must be an object[/red]")
|
||||||
|
raise typer.Exit(1)
|
||||||
|
|
||||||
|
# Map JSON keys to parameter names (handle aliases)
|
||||||
|
key_map = {"negative_prompt": "negative", "lora_name": "lora"}
|
||||||
|
mapped: dict[str, Any] = {}
|
||||||
|
for k, v in ji.items():
|
||||||
|
mapped[key_map.get(k, k)] = v
|
||||||
|
|
||||||
|
# Determine which CLI params the user explicitly set
|
||||||
|
click_ctx = ctx._context if hasattr(ctx, "_context") else ctx
|
||||||
|
explicit = {
|
||||||
|
p.name
|
||||||
|
for p in click_ctx.command.params
|
||||||
|
if click_ctx.get_parameter_source(p.name) == click.core.ParameterSource.COMMANDLINE
|
||||||
|
} if hasattr(click_ctx, "get_parameter_source") else set()
|
||||||
|
|
||||||
|
# Apply JSON values for anything not explicitly set on CLI
|
||||||
|
if "prompt" in mapped and ("prompt" not in explicit and prompt is None):
|
||||||
|
prompt = mapped["prompt"]
|
||||||
|
if "model" in mapped and "model" not in explicit:
|
||||||
|
model = mapped["model"]
|
||||||
|
if "width" in mapped and "width" not in explicit:
|
||||||
|
width = int(mapped["width"])
|
||||||
|
if "height" in mapped and "height" not in explicit:
|
||||||
|
height = int(mapped["height"])
|
||||||
|
if "steps" in mapped and "steps" not in explicit:
|
||||||
|
steps = int(mapped["steps"])
|
||||||
|
if "cfg" in mapped and "cfg" not in explicit:
|
||||||
|
cfg = float(mapped["cfg"])
|
||||||
|
if "seed" in mapped and "seed" not in explicit:
|
||||||
|
seed = int(mapped["seed"])
|
||||||
|
if "sampler" in mapped and "sampler" not in explicit:
|
||||||
|
sampler = mapped["sampler"]
|
||||||
|
if "scheduler" in mapped and "scheduler" not in explicit:
|
||||||
|
scheduler = mapped["scheduler"]
|
||||||
|
if "vae" in mapped and "vae" not in explicit:
|
||||||
|
vae = mapped["vae"]
|
||||||
|
if "lora" in mapped and "lora" not in explicit:
|
||||||
|
lora = mapped["lora"]
|
||||||
|
if "lora_strength" in mapped and "lora_strength" not in explicit:
|
||||||
|
lora_strength = float(mapped["lora_strength"])
|
||||||
|
if "negative" in mapped and "negative" not in explicit:
|
||||||
|
negative = mapped["negative"]
|
||||||
|
if "output" in mapped and "output" not in explicit:
|
||||||
|
output = Path(mapped["output"])
|
||||||
|
if "remote" in mapped and "remote" not in explicit:
|
||||||
|
remote = mapped["remote"]
|
||||||
|
|
||||||
|
if not prompt:
|
||||||
|
console.print("[red]Prompt is required (as argument or in --input JSON)[/red]")
|
||||||
|
raise typer.Exit(1)
|
||||||
|
|
||||||
from tensors.config import resolve_remote as do_resolve_remote # noqa: PLC0415
|
from tensors.config import resolve_remote as do_resolve_remote # noqa: PLC0415
|
||||||
|
|
||||||
# Resolve remote (explicit flag, or default from config)
|
# Resolve remote (explicit flag, or default from config)
|
||||||
@@ -883,6 +963,7 @@ def generate( # noqa: PLR0915
|
|||||||
lora_name=lora,
|
lora_name=lora,
|
||||||
lora_strength=lora_strength,
|
lora_strength=lora_strength,
|
||||||
vae=vae,
|
vae=vae,
|
||||||
|
orientation=orientation,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not result_local:
|
if not result_local:
|
||||||
@@ -1576,13 +1657,14 @@ def comfy_generate( # noqa: PLR0915
|
|||||||
url: Annotated[str | None, typer.Option("--url", "-u", help="ComfyUI server URL")] = None,
|
url: Annotated[str | None, typer.Option("--url", "-u", help="ComfyUI server URL")] = None,
|
||||||
negative: Annotated[str, typer.Option("-n", "--negative", help="Negative prompt")] = "",
|
negative: Annotated[str, typer.Option("-n", "--negative", help="Negative prompt")] = "",
|
||||||
model: Annotated[str | None, typer.Option("-m", "--model", help="Checkpoint model name")] = None,
|
model: Annotated[str | None, typer.Option("-m", "--model", help="Checkpoint model name")] = None,
|
||||||
width: Annotated[int, typer.Option("-W", "--width", help="Image width")] = 1024,
|
width: Annotated[int | None, typer.Option("-W", "--width", help="Image width (auto from checkpoint)")] = None,
|
||||||
height: Annotated[int, typer.Option("-H", "--height", help="Image height")] = 1024,
|
height: Annotated[int | None, typer.Option("-H", "--height", help="Image height (auto from checkpoint)")] = None,
|
||||||
steps: Annotated[int, typer.Option("--steps", help="Sampling steps")] = 20,
|
steps: Annotated[int | None, typer.Option("--steps", help="Sampling steps (auto from checkpoint)")] = None,
|
||||||
cfg: Annotated[float, typer.Option("--cfg", help="CFG scale")] = 7.0,
|
cfg: Annotated[float | None, typer.Option("--cfg", help="CFG scale (auto from checkpoint)")] = None,
|
||||||
seed: Annotated[int, typer.Option("--seed", "-s", help="Random seed (-1 for random)")] = -1,
|
seed: Annotated[int, typer.Option("--seed", "-s", help="Random seed (-1 for random)")] = -1,
|
||||||
sampler: Annotated[str, typer.Option("--sampler", help="Sampler name")] = "euler",
|
sampler: Annotated[str | None, typer.Option("--sampler", help="Sampler name (auto from checkpoint)")] = None,
|
||||||
scheduler: Annotated[str, typer.Option("--scheduler", help="Scheduler name")] = "normal",
|
scheduler: Annotated[str | None, typer.Option("--scheduler", help="Scheduler name (auto from checkpoint)")] = None,
|
||||||
|
orientation: Annotated[str, typer.Option("-O", "--orientation", help="Resolution: square, portrait, landscape")] = "square",
|
||||||
output: Annotated[Path | None, typer.Option("-o", "--output", help="Output file path")] = None,
|
output: Annotated[Path | None, typer.Option("-o", "--output", help="Output file path")] = None,
|
||||||
count: Annotated[int, typer.Option("-c", "--count", help="Number of images to generate")] = 1,
|
count: Annotated[int, typer.Option("-c", "--count", help="Number of images to generate")] = 1,
|
||||||
lora: Annotated[str | None, typer.Option("-l", "--lora", help="LoRA model name")] = None,
|
lora: Annotated[str | None, typer.Option("-l", "--lora", help="LoRA model name")] = None,
|
||||||
@@ -1686,6 +1768,7 @@ def comfy_generate( # noqa: PLR0915
|
|||||||
lora_name=lora,
|
lora_name=lora,
|
||||||
lora_strength=lora_strength,
|
lora_strength=lora_strength,
|
||||||
batch_size=count,
|
batch_size=count,
|
||||||
|
orientation=orientation,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not result:
|
if not result:
|
||||||
|
|||||||
+64
-40
@@ -709,47 +709,65 @@ DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
# Default VAE for SDXL/Illustrious/Pony models
|
|
||||||
DEFAULT_VAE = "sdxl_vae.safetensors"
|
|
||||||
|
|
||||||
|
|
||||||
def _build_workflow(
|
def _build_workflow(
|
||||||
prompt: str,
|
prompt: str,
|
||||||
negative_prompt: str = "",
|
negative_prompt: str = "",
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
width: int = 1024,
|
width: int | None = None,
|
||||||
height: int = 1024,
|
height: int | None = None,
|
||||||
steps: int = 20,
|
steps: int | None = None,
|
||||||
cfg: float = 7.0,
|
cfg: float | None = None,
|
||||||
seed: int = -1,
|
seed: int = -1,
|
||||||
sampler: str = "euler",
|
sampler: str | None = None,
|
||||||
scheduler: str = "normal",
|
scheduler: str | None = None,
|
||||||
lora_name: str | None = None,
|
lora_name: str | None = None,
|
||||||
lora_strength: float = 1.0,
|
lora_strength: float = 1.0,
|
||||||
batch_size: int = 1,
|
batch_size: int = 1,
|
||||||
vae: str | None = None,
|
vae: str | None = None,
|
||||||
|
orientation: str = "square",
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
"""Build a text-to-image workflow from parameters.
|
"""Build a text-to-image workflow from parameters.
|
||||||
|
|
||||||
|
Parameters set to None are auto-resolved from the checkpoint's family preset
|
||||||
|
via config.get_model_generation_defaults(). User-provided values always win.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
prompt: Positive prompt text
|
prompt: Positive prompt text
|
||||||
negative_prompt: Negative prompt text
|
negative_prompt: Negative prompt text
|
||||||
model: Checkpoint filename (if None, uses first available)
|
model: Checkpoint filename (if None, uses first available)
|
||||||
width: Image width
|
width: Image width (None = use preset for orientation)
|
||||||
height: Image height
|
height: Image height (None = use preset for orientation)
|
||||||
steps: Number of sampling steps
|
steps: Number of sampling steps (None = use preset)
|
||||||
cfg: CFG scale
|
cfg: CFG scale (None = use preset)
|
||||||
seed: Random seed (-1 for random)
|
seed: Random seed (-1 for random)
|
||||||
sampler: Sampler name
|
sampler: Sampler name (None = use preset)
|
||||||
scheduler: Scheduler name
|
scheduler: Scheduler name (None = use preset)
|
||||||
lora_name: LoRA model filename (optional)
|
lora_name: LoRA model filename (optional)
|
||||||
lora_strength: LoRA strength (default 1.0)
|
lora_strength: LoRA strength (default 1.0)
|
||||||
batch_size: Number of images to generate in one workflow (default 1)
|
batch_size: Number of images to generate in one workflow (default 1)
|
||||||
vae: VAE filename (defaults to sdxl_vae.safetensors)
|
vae: VAE filename (None = use preset)
|
||||||
|
orientation: Resolution orientation: "square", "portrait", or "landscape"
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
ComfyUI workflow dict
|
ComfyUI workflow dict
|
||||||
"""
|
"""
|
||||||
|
from tensors.config import get_model_generation_defaults, resolve_orientation # noqa: PLC0415
|
||||||
|
|
||||||
|
# Get preset defaults for this checkpoint family
|
||||||
|
defaults = get_model_generation_defaults(model or "") if model else get_model_generation_defaults("")
|
||||||
|
|
||||||
|
# Resolve orientation-based resolution
|
||||||
|
res_w, res_h = resolve_orientation(defaults.get("family"), orientation)
|
||||||
|
|
||||||
|
# Merge: user overrides > preset defaults
|
||||||
|
resolved_sampler = sampler if sampler is not None else defaults.get("sampler", "euler")
|
||||||
|
resolved_scheduler = scheduler if scheduler is not None else defaults.get("scheduler", "normal")
|
||||||
|
resolved_cfg = cfg if cfg is not None else defaults.get("cfg", 7.0)
|
||||||
|
resolved_steps = steps if steps is not None else defaults.get("steps", 20)
|
||||||
|
resolved_width = width if width is not None else res_w
|
||||||
|
resolved_height = height if height is not None else res_h
|
||||||
|
resolved_vae = vae if vae is not None else defaults.get("vae")
|
||||||
|
|
||||||
workflow = copy.deepcopy(DEFAULT_WORKFLOW_TEMPLATE)
|
workflow = copy.deepcopy(DEFAULT_WORKFLOW_TEMPLATE)
|
||||||
|
|
||||||
# Set seed (random if -1)
|
# Set seed (random if -1)
|
||||||
@@ -757,30 +775,30 @@ def _build_workflow(
|
|||||||
|
|
||||||
# Update KSampler settings
|
# Update KSampler settings
|
||||||
workflow["3"]["inputs"]["seed"] = actual_seed
|
workflow["3"]["inputs"]["seed"] = actual_seed
|
||||||
workflow["3"]["inputs"]["steps"] = steps
|
workflow["3"]["inputs"]["steps"] = resolved_steps
|
||||||
workflow["3"]["inputs"]["cfg"] = cfg
|
workflow["3"]["inputs"]["cfg"] = resolved_cfg
|
||||||
workflow["3"]["inputs"]["sampler_name"] = sampler
|
workflow["3"]["inputs"]["sampler_name"] = resolved_sampler
|
||||||
workflow["3"]["inputs"]["scheduler"] = scheduler
|
workflow["3"]["inputs"]["scheduler"] = resolved_scheduler
|
||||||
|
|
||||||
# Set model
|
# Set model
|
||||||
if model:
|
if model:
|
||||||
workflow["4"]["inputs"]["ckpt_name"] = model
|
workflow["4"]["inputs"]["ckpt_name"] = model
|
||||||
|
|
||||||
# Set dimensions and batch size
|
# Set dimensions and batch size
|
||||||
workflow["5"]["inputs"]["width"] = width
|
workflow["5"]["inputs"]["width"] = resolved_width
|
||||||
workflow["5"]["inputs"]["height"] = height
|
workflow["5"]["inputs"]["height"] = resolved_height
|
||||||
workflow["5"]["inputs"]["batch_size"] = batch_size
|
workflow["5"]["inputs"]["batch_size"] = batch_size
|
||||||
|
|
||||||
# Set prompts
|
# Set prompts
|
||||||
workflow["6"]["inputs"]["text"] = prompt
|
workflow["6"]["inputs"]["text"] = prompt
|
||||||
workflow["7"]["inputs"]["text"] = negative_prompt
|
workflow["7"]["inputs"]["text"] = negative_prompt
|
||||||
|
|
||||||
# Set VAE - use external VAE if specified, otherwise use checkpoint's built-in VAE
|
# Set VAE - use preset VAE if available, otherwise use checkpoint's built-in
|
||||||
if vae:
|
if resolved_vae:
|
||||||
# Use external VAE loader (node 11)
|
# Use external VAE loader (node 11)
|
||||||
workflow["11"]["inputs"]["vae_name"] = vae
|
workflow["11"]["inputs"]["vae_name"] = resolved_vae
|
||||||
else:
|
else:
|
||||||
# Use VAE from checkpoint (node 4, output index 2) - works for SD 1.5 models
|
# Use VAE from checkpoint (node 4, output index 2) - fallback for unknown models
|
||||||
# Remove VAELoader node and connect VAEDecode directly to checkpoint
|
# Remove VAELoader node and connect VAEDecode directly to checkpoint
|
||||||
del workflow["11"]
|
del workflow["11"]
|
||||||
workflow["8"]["inputs"]["vae"] = ["4", 2]
|
workflow["8"]["inputs"]["vae"] = ["4", 2]
|
||||||
@@ -812,13 +830,13 @@ def generate_image(
|
|||||||
url: str | None = None,
|
url: str | None = None,
|
||||||
negative_prompt: str = "",
|
negative_prompt: str = "",
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
width: int = 1024,
|
width: int | None = None,
|
||||||
height: int = 1024,
|
height: int | None = None,
|
||||||
steps: int = 20,
|
steps: int | None = None,
|
||||||
cfg: float = 7.0,
|
cfg: float | None = None,
|
||||||
seed: int = -1,
|
seed: int = -1,
|
||||||
sampler: str = "euler",
|
sampler: str | None = None,
|
||||||
scheduler: str = "normal",
|
scheduler: str | None = None,
|
||||||
console: Console | None = None,
|
console: Console | None = None,
|
||||||
on_progress: ProgressCallback | None = None,
|
on_progress: ProgressCallback | None = None,
|
||||||
timeout: float = 600.0,
|
timeout: float = 600.0,
|
||||||
@@ -826,28 +844,33 @@ def generate_image(
|
|||||||
lora_strength: float = 1.0,
|
lora_strength: float = 1.0,
|
||||||
batch_size: int = 1,
|
batch_size: int = 1,
|
||||||
vae: str | None = None,
|
vae: str | None = None,
|
||||||
|
orientation: str = "square",
|
||||||
) -> GenerationResult | None:
|
) -> GenerationResult | None:
|
||||||
"""Generate an image using a simple text-to-image workflow.
|
"""Generate an image using a simple text-to-image workflow.
|
||||||
|
|
||||||
|
Parameters set to None are auto-resolved from the checkpoint's family preset.
|
||||||
|
User-provided values always override preset defaults.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
prompt: Positive prompt text
|
prompt: Positive prompt text
|
||||||
url: ComfyUI base URL
|
url: ComfyUI base URL
|
||||||
negative_prompt: Negative prompt text
|
negative_prompt: Negative prompt text
|
||||||
model: Checkpoint filename (if None, must be pre-loaded in ComfyUI)
|
model: Checkpoint filename (if None, must be pre-loaded in ComfyUI)
|
||||||
width: Image width
|
width: Image width (None = use preset for orientation)
|
||||||
height: Image height
|
height: Image height (None = use preset for orientation)
|
||||||
steps: Number of sampling steps
|
steps: Number of sampling steps (None = use preset)
|
||||||
cfg: CFG scale
|
cfg: CFG scale (None = use preset)
|
||||||
seed: Random seed (-1 for random)
|
seed: Random seed (-1 for random)
|
||||||
sampler: Sampler name (euler, dpm_2, etc.)
|
sampler: Sampler name (None = use preset)
|
||||||
scheduler: Scheduler name (normal, karras, etc.)
|
scheduler: Scheduler name (None = use preset)
|
||||||
console: Rich console for progress output
|
console: Rich console for progress output
|
||||||
on_progress: Optional callback for progress updates
|
on_progress: Optional callback for progress updates
|
||||||
timeout: Maximum wait time in seconds
|
timeout: Maximum wait time in seconds
|
||||||
lora_name: LoRA model filename (optional)
|
lora_name: LoRA model filename (optional)
|
||||||
lora_strength: LoRA strength (default 1.0)
|
lora_strength: LoRA strength (default 1.0)
|
||||||
batch_size: Number of images to generate in one workflow (default 1)
|
batch_size: Number of images to generate in one workflow (default 1)
|
||||||
vae: VAE filename (defaults to sdxl_vae.safetensors)
|
vae: VAE filename (None = use preset)
|
||||||
|
orientation: Resolution orientation: "square", "portrait", or "landscape"
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
GenerationResult with image paths, or None if generation failed
|
GenerationResult with image paths, or None if generation failed
|
||||||
@@ -882,6 +905,7 @@ def generate_image(
|
|||||||
lora_strength=lora_strength,
|
lora_strength=lora_strength,
|
||||||
batch_size=batch_size,
|
batch_size=batch_size,
|
||||||
vae=vae,
|
vae=vae,
|
||||||
|
orientation=orientation,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Run workflow
|
# Run workflow
|
||||||
|
|||||||
+63
-19
@@ -512,29 +512,35 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
|
|||||||
"negative_prompt": "score_5, score_4, ugly, deformed, blurry, bad anatomy, bad hands, missing fingers",
|
"negative_prompt": "score_5, score_4, ugly, deformed, blurry, bad anatomy, bad hands, missing fingers",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
"cfg": 7.0,
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
|
"cfg": 6.5,
|
||||||
"clip_skip": 2,
|
"clip_skip": 2,
|
||||||
"sampler": "euler_ancestral",
|
"sampler": "euler_ancestral",
|
||||||
"scheduler": "normal",
|
"scheduler": "normal",
|
||||||
"steps": 25,
|
"steps": 25,
|
||||||
"vae": "sdxl_vae.safetensors",
|
"vae": "ponyStandardVAE_v10.safetensors",
|
||||||
},
|
},
|
||||||
"illustrious": {
|
"illustrious": {
|
||||||
"quality_prefix": "masterpiece, best quality, highres",
|
"quality_prefix": "masterpiece, best quality, highres",
|
||||||
"negative_prompt": "worst quality, bad quality, low quality, lowres, bad anatomy, bad hands, jpeg artifacts, watermark",
|
"negative_prompt": "worst quality, bad quality, low quality, lowres, bad anatomy, bad hands, jpeg artifacts, watermark",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
"cfg": 6.0,
|
"cfg": 6.0,
|
||||||
"sampler": "euler_ancestral",
|
"sampler": "euler_ancestral",
|
||||||
"scheduler": "normal",
|
"scheduler": "normal",
|
||||||
"steps": 25,
|
"steps": 25,
|
||||||
"vae": "sdxl_vae.safetensors",
|
"vae": "illustriousXLV20_v10.safetensors",
|
||||||
},
|
},
|
||||||
"sdxl": {
|
"sdxl": {
|
||||||
"quality_prefix": "",
|
"quality_prefix": "",
|
||||||
"negative_prompt": "ugly, deformed, bad anatomy, bad hands, extra fingers, missing fingers, blurry, watermark",
|
"negative_prompt": "ugly, deformed, bad anatomy, bad hands, extra fingers, missing fingers, blurry, watermark",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
"cfg": 7.0,
|
"cfg": 7.0,
|
||||||
"sampler": "dpmpp_2m",
|
"sampler": "dpmpp_2m",
|
||||||
"scheduler": "karras",
|
"scheduler": "karras",
|
||||||
@@ -546,21 +552,25 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
|
|||||||
"negative_prompt": "ugly, deformed, bad anatomy, bad hands, extra fingers, missing fingers, blurry, watermark",
|
"negative_prompt": "ugly, deformed, bad anatomy, bad hands, extra fingers, missing fingers, blurry, watermark",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
"cfg": 2.0,
|
"cfg": 2.0,
|
||||||
"sampler": "euler",
|
"sampler": "euler",
|
||||||
"scheduler": "sgm_uniform",
|
"scheduler": "sgm_uniform",
|
||||||
"steps": 8, # Lightning models use fewer steps
|
"steps": 8,
|
||||||
"vae": "sdxl_vae.safetensors",
|
"vae": "sdxl_vae.safetensors",
|
||||||
},
|
},
|
||||||
"sdxl_turbo": {
|
"sdxl_turbo": {
|
||||||
"quality_prefix": "",
|
"quality_prefix": "",
|
||||||
"negative_prompt": "", # Turbo models work best without negative prompts
|
"negative_prompt": "",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
"cfg": 1.0, # Very low CFG for turbo
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
|
"cfg": 1.0,
|
||||||
"sampler": "euler_ancestral",
|
"sampler": "euler_ancestral",
|
||||||
"scheduler": "normal",
|
"scheduler": "normal",
|
||||||
"steps": 4, # Turbo models use very few steps
|
"steps": 4,
|
||||||
"vae": "sdxl_vae.safetensors",
|
"vae": "sdxl_vae.safetensors",
|
||||||
},
|
},
|
||||||
"sd15": {
|
"sd15": {
|
||||||
@@ -571,28 +581,34 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
|
|||||||
),
|
),
|
||||||
"width": 512,
|
"width": 512,
|
||||||
"height": 512,
|
"height": 512,
|
||||||
|
"portrait": (512, 768),
|
||||||
|
"landscape": (768, 512),
|
||||||
"cfg": 7.0,
|
"cfg": 7.0,
|
||||||
"sampler": "dpmpp_2m",
|
"sampler": "euler_ancestral",
|
||||||
"scheduler": "karras",
|
"scheduler": "normal",
|
||||||
"steps": 20,
|
"steps": 25,
|
||||||
"vae": None, # Use checkpoint's built-in VAE
|
"vae": "vae-ft-mse-840000-ema-pruned.safetensors",
|
||||||
},
|
},
|
||||||
"sd15_lcm": {
|
"sd15_lcm": {
|
||||||
"quality_prefix": "masterpiece, best quality",
|
"quality_prefix": "masterpiece, best quality",
|
||||||
"negative_prompt": "", # LCM works best with minimal negative
|
"negative_prompt": "",
|
||||||
"width": 512,
|
"width": 512,
|
||||||
"height": 512,
|
"height": 512,
|
||||||
|
"portrait": (512, 768),
|
||||||
|
"landscape": (768, 512),
|
||||||
"cfg": 1.5,
|
"cfg": 1.5,
|
||||||
"sampler": "lcm",
|
"sampler": "lcm",
|
||||||
"scheduler": "normal",
|
"scheduler": "normal",
|
||||||
"steps": 6,
|
"steps": 6,
|
||||||
"vae": None, # Use checkpoint's built-in VAE
|
"vae": "vae-ft-mse-840000-ema-pruned.safetensors",
|
||||||
},
|
},
|
||||||
"flux": {
|
"flux": {
|
||||||
"quality_prefix": "",
|
"quality_prefix": "",
|
||||||
"negative_prompt": "", # Flux doesn't use negative prompts effectively
|
"negative_prompt": "",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
"cfg": 3.5,
|
"cfg": 3.5,
|
||||||
"sampler": "euler",
|
"sampler": "euler",
|
||||||
"scheduler": "simple",
|
"scheduler": "simple",
|
||||||
@@ -604,21 +620,25 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
|
|||||||
"negative_prompt": "",
|
"negative_prompt": "",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
"cfg": 1.0, # Schnell uses low CFG
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
|
"cfg": 1.0,
|
||||||
"sampler": "euler",
|
"sampler": "euler",
|
||||||
"scheduler": "simple",
|
"scheduler": "simple",
|
||||||
"steps": 4, # Schnell is a distilled model, very few steps
|
"steps": 4,
|
||||||
"vae": "ae.safetensors",
|
"vae": "ae.safetensors",
|
||||||
},
|
},
|
||||||
"zimage": {
|
"zimage": {
|
||||||
"quality_prefix": "",
|
"quality_prefix": "",
|
||||||
"negative_prompt": "", # Turbo models work best without negative prompts
|
"negative_prompt": "",
|
||||||
"width": 1024,
|
"width": 1024,
|
||||||
"height": 1024,
|
"height": 1024,
|
||||||
"cfg": 1.0, # Very low CFG for turbo
|
"portrait": (832, 1216),
|
||||||
|
"landscape": (1216, 832),
|
||||||
|
"cfg": 1.0,
|
||||||
"sampler": "euler",
|
"sampler": "euler",
|
||||||
"scheduler": "simple",
|
"scheduler": "simple",
|
||||||
"steps": 4, # ZImageTurbo is a distilled model
|
"steps": 4,
|
||||||
"vae": "ae.safetensors",
|
"vae": "ae.safetensors",
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
@@ -733,6 +753,30 @@ def get_model_generation_defaults(model_name: str, base_model: str | None = None
|
|||||||
return defaults
|
return defaults
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_orientation(family: str | None, orientation: str = "square") -> tuple[int, int]:
|
||||||
|
"""Get width/height for a model family and orientation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
family: Model family key (e.g. "pony", "sd15", "sdxl") or None for default
|
||||||
|
orientation: One of "square", "portrait", "landscape"
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(width, height) tuple
|
||||||
|
"""
|
||||||
|
defaults = MODEL_FAMILY_DEFAULTS.get(family or "sdxl", MODEL_FAMILY_DEFAULTS["sdxl"])
|
||||||
|
w: int = defaults["width"]
|
||||||
|
h: int = defaults["height"]
|
||||||
|
fallback = (w, h)
|
||||||
|
|
||||||
|
if orientation == "portrait":
|
||||||
|
pair: tuple[int, int] = defaults.get("portrait", fallback)
|
||||||
|
return pair
|
||||||
|
if orientation == "landscape":
|
||||||
|
pair = defaults.get("landscape", fallback)
|
||||||
|
return pair
|
||||||
|
return fallback
|
||||||
|
|
||||||
|
|
||||||
def get_comfyui_url() -> str:
|
def get_comfyui_url() -> str:
|
||||||
"""Get the ComfyUI server URL.
|
"""Get the ComfyUI server URL.
|
||||||
|
|
||||||
|
|||||||
@@ -356,7 +356,7 @@ class TestModelFamilyDetection:
|
|||||||
assert defaults["sampler"] == "euler_ancestral"
|
assert defaults["sampler"] == "euler_ancestral"
|
||||||
assert defaults["scheduler"] == "normal"
|
assert defaults["scheduler"] == "normal"
|
||||||
assert defaults["steps"] == 25
|
assert defaults["steps"] == 25
|
||||||
assert defaults["cfg"] == 7.0
|
assert defaults["cfg"] == 6.5
|
||||||
|
|
||||||
def test_get_model_generation_defaults_flux(self) -> None:
|
def test_get_model_generation_defaults_flux(self) -> None:
|
||||||
"""Test getting generation defaults for Flux models."""
|
"""Test getting generation defaults for Flux models."""
|
||||||
|
|||||||
Reference in New Issue
Block a user