Add separate VAE loader for SDXL/Illustrious/Pony models
- Use sdxl_vae.safetensors by default instead of checkpoint VAE - Add vae parameter to generate_image and API endpoint - Better quality for modern SDXL-based models Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
+18
-3
@@ -661,8 +661,8 @@ LORA_LOADER_NODE: dict[str, Any] = {
|
||||
},
|
||||
}
|
||||
|
||||
# Default SDXL/Flux compatible workflow template
|
||||
# This is a minimal text-to-image workflow that works with most models
|
||||
# Default SDXL/Illustrious/Pony compatible workflow template
|
||||
# Uses separate VAE loader for better quality with modern models
|
||||
DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
@@ -697,14 +697,21 @@ DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
|
||||
},
|
||||
"8": {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": {"samples": ["3", 0], "vae": ["4", 2]},
|
||||
"inputs": {"samples": ["3", 0], "vae": ["11", 0]},
|
||||
},
|
||||
"9": {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {"filename_prefix": "comfy", "images": ["8", 0]},
|
||||
},
|
||||
"11": {
|
||||
"class_type": "VAELoader",
|
||||
"inputs": {"vae_name": "sdxl_vae.safetensors"},
|
||||
},
|
||||
}
|
||||
|
||||
# Default VAE for SDXL/Illustrious/Pony models
|
||||
DEFAULT_VAE = "sdxl_vae.safetensors"
|
||||
|
||||
|
||||
def _build_workflow(
|
||||
prompt: str,
|
||||
@@ -720,6 +727,7 @@ def _build_workflow(
|
||||
lora_name: str | None = None,
|
||||
lora_strength: float = 1.0,
|
||||
batch_size: int = 1,
|
||||
vae: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a text-to-image workflow from parameters.
|
||||
|
||||
@@ -737,6 +745,7 @@ def _build_workflow(
|
||||
lora_name: LoRA model filename (optional)
|
||||
lora_strength: LoRA strength (default 1.0)
|
||||
batch_size: Number of images to generate in one workflow (default 1)
|
||||
vae: VAE filename (defaults to sdxl_vae.safetensors)
|
||||
|
||||
Returns:
|
||||
ComfyUI workflow dict
|
||||
@@ -766,6 +775,9 @@ def _build_workflow(
|
||||
workflow["6"]["inputs"]["text"] = prompt
|
||||
workflow["7"]["inputs"]["text"] = negative_prompt
|
||||
|
||||
# Set VAE
|
||||
workflow["11"]["inputs"]["vae_name"] = vae or DEFAULT_VAE
|
||||
|
||||
# Inject LoRA loader if specified
|
||||
if lora_name:
|
||||
# Add LoRA loader node (node 10)
|
||||
@@ -806,6 +818,7 @@ def generate_image(
|
||||
lora_name: str | None = None,
|
||||
lora_strength: float = 1.0,
|
||||
batch_size: int = 1,
|
||||
vae: str | None = None,
|
||||
) -> GenerationResult | None:
|
||||
"""Generate an image using a simple text-to-image workflow.
|
||||
|
||||
@@ -827,6 +840,7 @@ def generate_image(
|
||||
lora_name: LoRA model filename (optional)
|
||||
lora_strength: LoRA strength (default 1.0)
|
||||
batch_size: Number of images to generate in one workflow (default 1)
|
||||
vae: VAE filename (defaults to sdxl_vae.safetensors)
|
||||
|
||||
Returns:
|
||||
GenerationResult with image paths, or None if generation failed
|
||||
@@ -860,6 +874,7 @@ def generate_image(
|
||||
lora_name=lora_name,
|
||||
lora_strength=lora_strength,
|
||||
batch_size=batch_size,
|
||||
vae=vae,
|
||||
)
|
||||
|
||||
# Run workflow
|
||||
|
||||
@@ -44,6 +44,7 @@ class GenerateRequest(PydanticBaseModel):
|
||||
seed: int = Field(default=-1, description="Random seed (-1 for random)")
|
||||
sampler: str = Field(default="euler", description="Sampler name")
|
||||
scheduler: str = Field(default="normal", description="Scheduler name")
|
||||
vae: str | None = Field(default=None, description="VAE model name (defaults to sdxl_vae.safetensors)")
|
||||
|
||||
|
||||
class GenerateResponse(PydanticBaseModel):
|
||||
@@ -233,6 +234,7 @@ def comfyui_generate(request: GenerateRequest) -> dict[str, Any]:
|
||||
seed=request.seed,
|
||||
sampler=request.sampler,
|
||||
scheduler=request.scheduler,
|
||||
vae=request.vae,
|
||||
)
|
||||
|
||||
if not result:
|
||||
|
||||
Reference in New Issue
Block a user