aladac/tensors
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add LoRA support and model family quality defaults to comfy generate

Browse Source
This commit is contained in:
Adam Ladachowski
2026-02-19 04:30:32 +01:00
parent 91c5c1e0a7
commit 1ed40a3142
5 changed files with 248 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.coverage
BIN
View File
Binary file not shown.
tensors/cli.py
+72 -8
View File
@@ -21,6 +21,7 @@ from tensors.api import (
)
from tensors.config import (
CONFIG_FILE,
MODEL_FAMILY_DEFAULTS,
BaseModel,
CommercialUse,
ModelType,
@@ -28,6 +29,7 @@ from tensors.config import (
Period,
Provider,
SortOrder,
detect_model_family,
get_default_output_path,
get_model_paths,
load_api_key,
@@ -1118,7 +1120,7 @@ def comfy_history(
@comfy_app.command("generate")
def comfy_generate(
def comfy_generate( # noqa: PLR0915
prompt: Annotated[str, typer.Argument(help="Positive prompt text")],
url: Annotated[str | None, typer.Option("--url", "-u", help="ComfyUI server URL")] = None,
negative: Annotated[str, typer.Option("-n", "--negative", help="Negative prompt")] = "",
@@ -1132,6 +1134,10 @@ def comfy_generate(
scheduler: Annotated[str, typer.Option("--scheduler", help="Scheduler name")] = "normal",
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,
lora: Annotated[str | None, typer.Option("-l", "--lora", help="LoRA model name")] = None,
lora_strength: Annotated[float, typer.Option("--lora-strength", help="LoRA strength")] = 1.0,
no_quality: Annotated[bool, typer.Option("--no-quality", help="Disable auto quality tags")] = False,
no_negative: Annotated[bool, typer.Option("--no-negative", help="Disable auto negative prompt")] = False,
json_output: Annotated[bool, typer.Option("--json", "-j", help="Output as JSON")] = False,
) -> None:
"""Generate an image with a simple text-to-image workflow.
@@ -1141,6 +1147,8 @@ def comfy_generate(
tsr comfy generate "portrait photo" -n "blurry, bad quality" --steps 30
tsr comfy generate "landscape" -m "flux1-dev-fp8.safetensors" -W 1024 -H 768
tsr comfy generate "cyberpunk city" --count 4 -o batch.png
tsr comfy generate "girl" --lora spumcostyle.safetensors --lora-strength 0.8
tsr comfy generate "raw prompt" --no-quality --no-negative
"""
import random # noqa: PLC0415
@@ -1152,6 +1160,64 @@ def comfy_generate(
# Determine base seed for batch
base_seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
# Detect model family and apply defaults
family_defaults: dict[str, Any] = {}
model_family: str | None = None
if model:
# Try to get base_model from database
base_model_str: str | None = None
try:
with Database() as db:
db.init_schema()
base_model_str = db.get_base_model_by_filename(model)
except Exception:
pass
model_family = detect_model_family(model, base_model_str)
if model_family:
family_defaults = MODEL_FAMILY_DEFAULTS.get(model_family, {})
if not json_output:
console.print(f"[dim]Detected model family: {model_family}[/dim]")
# Build enhanced prompt with quality prefix and LoRA trigger words
enhanced_prompt = prompt
prompt_parts: list[str] = []
# Add LoRA trigger words if using LoRA
if lora:
try:
with Database() as db:
db.init_schema()
trigger_words = db.get_trigger_words_by_filename(lora)
if trigger_words:
prompt_parts.extend(trigger_words)
if not json_output:
console.print(f"[dim]LoRA trigger words: {', '.join(trigger_words)}[/dim]")
except Exception:
pass
# Add quality prefix based on model family
if not no_quality and family_defaults.get("quality_prefix"):
prompt_parts.append(family_defaults["quality_prefix"])
# Add user prompt
prompt_parts.append(prompt)
enhanced_prompt = ", ".join(prompt_parts) if len(prompt_parts) > 1 else prompt
# Build enhanced negative prompt
enhanced_negative = negative
if not no_negative and family_defaults.get("negative_prompt"):
family_negative = family_defaults["negative_prompt"]
enhanced_negative = f"{negative}, {family_negative}" if negative else family_negative
if not json_output and (enhanced_prompt != prompt or enhanced_negative != negative):
if enhanced_prompt != prompt:
truncated = enhanced_prompt[:100] + "..." if len(enhanced_prompt) > 100 else enhanced_prompt # noqa: PLR2004
console.print(f"[dim]Enhanced prompt: {truncated}[/dim]")
if enhanced_negative != negative:
truncated = enhanced_negative[:80] + "..." if len(enhanced_negative) > 80 else enhanced_negative # noqa: PLR2004
console.print(f"[dim]Enhanced negative: {truncated}[/dim]")
for i in range(count):
current_seed = base_seed + i if seed >= 0 else -1 # Increment seed or use random each time
@@ -1159,9 +1225,9 @@ def comfy_generate(
console.print(f"\n[cyan]Generating image {i + 1}/{count}...[/cyan]")
result = generate_image(
prompt=prompt,
prompt=enhanced_prompt,
url=url,
negative_prompt=negative,
negative_prompt=enhanced_negative,
model=model,
width=width,
height=height,
@@ -1171,6 +1237,8 @@ def comfy_generate(
sampler=sampler,
scheduler=scheduler,
console=console if not json_output else None,
lora_name=lora,
lora_strength=lora_strength,
)
if not result:
@@ -1195,11 +1263,7 @@ def comfy_generate(
img_path = result.images[0]
img_data = get_image(str(img_path), url=url)
if img_data:
if count == 1:
save_path = output
else:
# Add index suffix for batch: output.png -> output_001.png
save_path = output.parent / f"{output.stem}_{i + 1:03d}{output.suffix}"
save_path = output if count == 1 else output.parent / f"{output.stem}_{i + 1:03d}{output.suffix}"
save_path.write_bytes(img_data)
saved_path = save_path
all_saved.append(save_path)
tensors/comfyui.py
+41
View File
@@ -526,6 +526,18 @@ def run_workflow(
# Simple Text-to-Image Generation
# ============================================================================
# LoRA loader node template (inserted between checkpoint and sampler)
LORA_LOADER_NODE: dict[str, Any] = {
"class_type": "LoraLoader",
"inputs": {
"lora_name": "",
"strength_model": 1.0,
"strength_clip": 1.0,
"model": ["4", 0], # From checkpoint
"clip": ["4", 1], # From checkpoint
},
}
# Default SDXL/Flux compatible workflow template
# This is a minimal text-to-image workflow that works with most models
DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
@@ -582,6 +594,8 @@ def _build_workflow(
seed: int = -1,
sampler: str = "euler",
scheduler: str = "normal",
lora_name: str | None = None,
lora_strength: float = 1.0,
) -> dict[str, Any]:
"""Build a text-to-image workflow from parameters.
@@ -596,6 +610,8 @@ def _build_workflow(
seed: Random seed (-1 for random)
sampler: Sampler name
scheduler: Scheduler name
lora_name: LoRA model filename (optional)
lora_strength: LoRA strength (default 1.0)
Returns:
ComfyUI workflow dict
@@ -624,6 +640,25 @@ def _build_workflow(
workflow["6"]["inputs"]["text"] = prompt
workflow["7"]["inputs"]["text"] = negative_prompt
# Inject LoRA loader if specified
if lora_name:
# Add LoRA loader node (node 10)
lora_node = copy.deepcopy(LORA_LOADER_NODE)
lora_node["inputs"]["lora_name"] = lora_name
lora_node["inputs"]["strength_model"] = lora_strength
lora_node["inputs"]["strength_clip"] = lora_strength
# LoRA takes model/clip from checkpoint (node 4)
lora_node["inputs"]["model"] = ["4", 0]
lora_node["inputs"]["clip"] = ["4", 1]
workflow["10"] = lora_node
# Reroute KSampler model input from checkpoint (4) to LoRA (10)
workflow["3"]["inputs"]["model"] = ["10", 0]
# Reroute CLIP text encoders from checkpoint (4) to LoRA (10)
workflow["6"]["inputs"]["clip"] = ["10", 1]
workflow["7"]["inputs"]["clip"] = ["10", 1]
return workflow
@@ -642,6 +677,8 @@ def generate_image(
console: Console | None = None,
on_progress: ProgressCallback | None = None,
timeout: float = 600.0,
lora_name: str | None = None,
lora_strength: float = 1.0,
) -> GenerationResult | None:
"""Generate an image using a simple text-to-image workflow.
@@ -660,6 +697,8 @@ def generate_image(
console: Rich console for progress output
on_progress: Optional callback for progress updates
timeout: Maximum wait time in seconds
lora_name: LoRA model filename (optional)
lora_strength: LoRA strength (default 1.0)
Returns:
GenerationResult with image paths, or None if generation failed
@@ -690,6 +729,8 @@ def generate_image(
seed=seed,
sampler=sampler,
scheduler=scheduler,
lora_name=lora_name,
lora_strength=lora_strength,
)
# Run workflow
tensors/config.py
+87
View File
@@ -502,6 +502,93 @@ COMFYUI_DEFAULT_CFG = 7.0
COMFYUI_DEFAULT_SAMPLER = "euler"
COMFYUI_DEFAULT_SCHEDULER = "normal"
# ============================================================================
# Model Family Defaults (Quality Tags, Negative Prompts, etc.)
# ============================================================================
MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
"pony": {
"quality_prefix": "score_9, score_8_up, score_7_up",
"negative_prompt": "score_5, score_4, ugly, deformed, blurry, bad anatomy, bad hands, missing fingers",
"width": 1024,
"height": 1024,
"cfg": 7.0,
"clip_skip": 2,
},
"illustrious": {
"quality_prefix": "masterpiece, best quality, highres",
"negative_prompt": "worst quality, bad quality, low quality, lowres, bad anatomy, bad hands, jpeg artifacts, watermark",
"width": 1024,
"height": 1024,
"cfg": 6.0,
},
"sdxl": {
"quality_prefix": "",
"negative_prompt": "ugly, deformed, bad anatomy, bad hands, extra fingers, missing fingers, blurry, watermark",
"width": 1024,
"height": 1024,
"cfg": 7.0,
},
"sd15": {
"quality_prefix": "masterpiece, best quality",
"negative_prompt": (
"(worst quality:2), (low quality:2), bad anatomy, bad hands, extra fingers, "
"missing fingers, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, watermark"
),
"width": 512,
"height": 512,
"cfg": 7.0,
},
"flux": {
"quality_prefix": "",
"negative_prompt": "", # Flux doesn't use negative prompts effectively
"width": 1024,
"height": 1024,
"cfg": 3.5,
},
}
def detect_model_family(model_name: str, base_model: str | None = None) -> str | None: # noqa: PLR0911
"""Detect model family from filename or CivitAI base_model field.
Args:
model_name: Filename of the model (e.g., "ponyDiffusionV6XL.safetensors")
base_model: Optional CivitAI base_model field (e.g., "Pony", "SDXL 1.0")
Returns:
Model family key (pony, illustrious, sdxl, sd15, flux) or None if unknown
"""
name_lower = model_name.lower()
base_lower = (base_model or "").lower()
# Check base_model field first (most reliable from CivitAI)
if base_lower:
if "pony" in base_lower:
return "pony"
if "illustrious" in base_lower:
return "illustrious"
if "flux" in base_lower:
return "flux"
if "sd 1.5" in base_lower or "sd 1.4" in base_lower:
return "sd15"
if "sdxl" in base_lower:
return "sdxl"
# Fall back to filename heuristics
if "pony" in name_lower:
return "pony"
if "illustrious" in name_lower or "noob" in name_lower:
return "illustrious"
if "flux" in name_lower:
return "flux"
if any(x in name_lower for x in ["sd15", "sd1.5", "sd_1.5", "dreamshaper", "realistic", "deliberate", "anything"]):
return "sd15"
if any(x in name_lower for x in ["sdxl", "xl_"]):
return "sdxl"
return None
def get_comfyui_url() -> str:
"""Get the ComfyUI server URL.
tensors/db.py
+48
View File
@@ -755,6 +755,54 @@ class Database:
).all()
return [w.word for w in words]
def get_trigger_words_by_filename(self, filename: str) -> list[str]:
"""Get trigger words for a LoRA by matching filename in version_files.
Args:
filename: The filename to search for (e.g., "spumcostyle.safetensors")
Returns:
List of trigger/trained words from CivitAI metadata
"""
with self.session() as session:
# Find version file by filename match
vf = session.exec(select(VersionFile).where(VersionFile.name == filename)).first()
if not vf:
# Try partial match (without extension)
base_name = filename.rsplit(".", 1)[0] if "." in filename else filename
vf = session.exec(select(VersionFile).where(col(VersionFile.name).contains(base_name))).first()
if not vf or not vf.version_id:
return []
words = session.exec(
select(TrainedWord).where(TrainedWord.version_id == vf.version_id).order_by(col(TrainedWord.position))
).all()
return [w.word for w in words]
def get_base_model_by_filename(self, filename: str) -> str | None:
"""Get base_model for a checkpoint/LoRA by filename lookup.
Args:
filename: The filename to search for
Returns:
Base model string (e.g., "Pony", "SDXL 1.0") or None
"""
with self.session() as session:
# Find version file by filename match
vf = session.exec(select(VersionFile).where(VersionFile.name == filename)).first()
if not vf:
# Try partial match (without extension)
base_name = filename.rsplit(".", 1)[0] if "." in filename else filename
vf = session.exec(select(VersionFile).where(col(VersionFile.name).contains(base_name))).first()
if not vf or not vf.version_id:
return None
mv = session.get(ModelVersion, vf.version_id)
return mv.base_model if mv else None
# =========================================================================
# Statistics
# =========================================================================