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
|
# Default SDXL/Illustrious/Pony compatible workflow template
|
||||||
# This is a minimal text-to-image workflow that works with most models
|
# Uses separate VAE loader for better quality with modern models
|
||||||
DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
|
DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
|
||||||
"3": {
|
"3": {
|
||||||
"class_type": "KSampler",
|
"class_type": "KSampler",
|
||||||
@@ -697,14 +697,21 @@ DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
|
|||||||
},
|
},
|
||||||
"8": {
|
"8": {
|
||||||
"class_type": "VAEDecode",
|
"class_type": "VAEDecode",
|
||||||
"inputs": {"samples": ["3", 0], "vae": ["4", 2]},
|
"inputs": {"samples": ["3", 0], "vae": ["11", 0]},
|
||||||
},
|
},
|
||||||
"9": {
|
"9": {
|
||||||
"class_type": "SaveImage",
|
"class_type": "SaveImage",
|
||||||
"inputs": {"filename_prefix": "comfy", "images": ["8", 0]},
|
"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(
|
def _build_workflow(
|
||||||
prompt: str,
|
prompt: str,
|
||||||
@@ -720,6 +727,7 @@ def _build_workflow(
|
|||||||
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,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
"""Build a text-to-image workflow from parameters.
|
"""Build a text-to-image workflow from parameters.
|
||||||
|
|
||||||
@@ -737,6 +745,7 @@ def _build_workflow(
|
|||||||
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)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
ComfyUI workflow dict
|
ComfyUI workflow dict
|
||||||
@@ -766,6 +775,9 @@ def _build_workflow(
|
|||||||
workflow["6"]["inputs"]["text"] = prompt
|
workflow["6"]["inputs"]["text"] = prompt
|
||||||
workflow["7"]["inputs"]["text"] = negative_prompt
|
workflow["7"]["inputs"]["text"] = negative_prompt
|
||||||
|
|
||||||
|
# Set VAE
|
||||||
|
workflow["11"]["inputs"]["vae_name"] = vae or DEFAULT_VAE
|
||||||
|
|
||||||
# Inject LoRA loader if specified
|
# Inject LoRA loader if specified
|
||||||
if lora_name:
|
if lora_name:
|
||||||
# Add LoRA loader node (node 10)
|
# Add LoRA loader node (node 10)
|
||||||
@@ -806,6 +818,7 @@ def generate_image(
|
|||||||
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,
|
||||||
) -> GenerationResult | None:
|
) -> GenerationResult | None:
|
||||||
"""Generate an image using a simple text-to-image workflow.
|
"""Generate an image using a simple text-to-image workflow.
|
||||||
|
|
||||||
@@ -827,6 +840,7 @@ def generate_image(
|
|||||||
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)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
GenerationResult with image paths, or None if generation failed
|
GenerationResult with image paths, or None if generation failed
|
||||||
@@ -860,6 +874,7 @@ def generate_image(
|
|||||||
lora_name=lora_name,
|
lora_name=lora_name,
|
||||||
lora_strength=lora_strength,
|
lora_strength=lora_strength,
|
||||||
batch_size=batch_size,
|
batch_size=batch_size,
|
||||||
|
vae=vae,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Run workflow
|
# Run workflow
|
||||||
|
|||||||
@@ -44,6 +44,7 @@ class GenerateRequest(PydanticBaseModel):
|
|||||||
seed: int = Field(default=-1, description="Random seed (-1 for random)")
|
seed: int = Field(default=-1, description="Random seed (-1 for random)")
|
||||||
sampler: str = Field(default="euler", description="Sampler name")
|
sampler: str = Field(default="euler", description="Sampler name")
|
||||||
scheduler: str = Field(default="normal", description="Scheduler 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):
|
class GenerateResponse(PydanticBaseModel):
|
||||||
@@ -233,6 +234,7 @@ def comfyui_generate(request: GenerateRequest) -> dict[str, Any]:
|
|||||||
seed=request.seed,
|
seed=request.seed,
|
||||||
sampler=request.sampler,
|
sampler=request.sampler,
|
||||||
scheduler=request.scheduler,
|
scheduler=request.scheduler,
|
||||||
|
vae=request.vae,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not result:
|
if not result:
|
||||||
|
|||||||
Reference in New Issue
Block a user