feat: Add dry run mode as UI parameter in advanced node

- Add dry_run boolean parameter to TripleKSamplerWan22LightningAdvanced INPUT_TYPES
- Replace ENABLE_DRY_RUN constant with user-configurable UI parameter
- Update _run_sampling_stage method to accept dry_run parameter
- Pass dry_run parameter through all sampling stage calls
- Remove ENABLE_DRY_RUN from constants.py and imports
- Update README.md with comprehensive feature documentation:
  - Add dry run mode section with usage instructions
  - Update model switching strategies to reflect all 5 options
  - Add KJNodes compatibility documentation
  - Update code structure to show inheritance hierarchy
  - Enhance version history with v0.3.1 changes
- Maintain simple node simplicity (no dry run parameter exposed)

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
VraethrDalkr
2025-09-14 11:17:24 -04:00
co-authored by Claude
parent ce8a1b1b62
commit 2dc0f8bb57
3 changed files with 84 additions and 22 deletions
+65 -11
View File
@@ -29,10 +29,12 @@ Main triple-stage sampler with optimized ease-of-use:
Advanced variant with complete configurability:
- **Full Parameter Control**: Every aspect of the sampling process is configurable
- **5 Switching Strategies**: Manual step, manual boundary, T2V boundary, I2V boundary, 50% midpoint
- **Dynamic UI**: Parameters show/hide based on selected strategy for cleaner interface
- **Auto-calculation Options**: Both base_steps and switch_step can be auto-computed (-1) or manually set
- **Dry Run Mode**: Test configurations without actual sampling execution
- **Lightning-start Aware**: Auto-calculation accounts for when Lightning processing begins
- **Model Switching**: Supports step-based, sigma boundary, and manual switching strategies
- **Professional Features**: Advanced users get ultimate flexibility
- **Professional Features**: Advanced users get ultimate flexibility and testing capabilities
## Installation
@@ -98,8 +100,13 @@ The example workflow includes:
- `lightning_steps`: Total lightning steps (typically 8)
4. Choose model switching strategy:
- **Midpoint**: Simple 50/50 split between lightning models
- **Boundary**: Sigma-based switching (0.875 for T2V, 0.900 for I2V)
- **"50% of steps"**: Simple 50/50 split between lightning models
- **"Manual switch step"**: Precise step-based control
- **"T2V boundary"**: Auto-optimized for text-to-video (0.875)
- **"I2V boundary"**: Auto-optimized for image-to-video (0.900)
- **"Manual boundary"**: Custom sigma-based switching
5. Optional: Enable **Dry Run** mode (Advanced node) to test configurations without sampling
### Parameter Guidelines
@@ -110,11 +117,24 @@ The example workflow includes:
## Model Switching Strategies
### Midpoint Strategy
The advanced node offers 5 switching strategies with dynamic UI that shows only relevant parameters:
### 1. "50% of steps" (Auto-Midpoint)
Simple approach that switches at 50% of lightning steps. Reliable and straightforward.
### Sigma Boundary Strategy
More sophisticated approach that analyzes sigma schedules to determine optimal switching point. Better quality but requires tuning boundary parameter for different use cases.
### 2. "Manual switch step"
Allows precise control over the switching step. Shows switch_step parameter for manual configuration.
### 3. "T2V boundary" (Auto-Boundary)
Optimized for text-to-video models, automatically uses boundary value of 0.875 for sigma-based switching.
### 4. "I2V boundary" (Auto-Boundary)
Optimized for image-to-video models, automatically uses boundary value of 0.900 for sigma-based switching.
### 5. "Manual boundary"
Full control over sigma boundary value. Shows switch_boundary parameter for manual configuration.
**Dynamic UI**: Only relevant parameters are shown based on the selected strategy, keeping the interface clean and focused.
## Configuration
@@ -124,6 +144,9 @@ The package includes a `constants.py` file with user-configurable parameters:
# Quality threshold for automatic base step calculation
MIN_TOTAL_STEPS = 20
# KJNodes compatibility fix toggle
ENABLE_KJNODES_COMPATIBILITY_FIX = True
# Default sigma boundaries for different model types
DEFAULT_BOUNDARY_T2V = 0.875 # Text-to-video models
DEFAULT_BOUNDARY_I2V = 0.900 # Image-to-video models
@@ -135,6 +158,31 @@ LOG_LEVEL = "INFO"
Users can modify these values to tune sampling behavior without changing the core implementation.
## Advanced Features
### Dry Run Mode (Advanced Node Only)
The advanced node includes a **Dry Run** boolean parameter for testing configurations:
- **Purpose**: Test parameter combinations without expensive sampling operations
- **Validation**: Performs complete parameter validation and logging
- **Performance**: Instant feedback for workflow testing and debugging
- **Usage**: Enable the "dry_run" checkbox in the advanced node interface
### KJNodes Compatibility
Automatic compatibility with KJNodes when installed:
- **Auto-Detection**: Automatically detects KJNodes installation
- **Parameter Stripping**: Safely handles transformer_options parameter conflicts
- **Dual-Layer Fix**: Uses both model proxy and function patching for reliability
- **Configurable**: Can be disabled by setting `ENABLE_KJNODES_COMPATIBILITY_FIX = False`
### Clean Architecture
Built with professional software engineering principles:
- **Inheritance Hierarchy**: Base → Advanced → Simple node structure eliminates code duplication
- **Dynamic UI**: Interface adapts to selected strategy for optimal user experience
- **Comprehensive Validation**: Edge cases and parameter conflicts are handled gracefully
## Logging
The nodes provide detailed logging including:
@@ -150,9 +198,13 @@ The nodes provide detailed logging including:
- Standard ComfyUI dependencies
### Code Structure
- `__init__.py`: Package initialization and node registration
- `triple_ksampler_wan22.py`: Main implementation with both node classes
- `__init__.py`: Package initialization, node registration, and web directory setup
- `triple_ksampler_wan22.py`: Main implementation with clean inheritance hierarchy:
- `TripleKSamplerWan22Base`: Shared functionality and core sampling logic
- `TripleKSamplerWan22LightningAdvanced`: Advanced node with full parameter control
- `TripleKSamplerWan22Lightning`: Simple node with streamlined interface
- `constants.py`: Configuration constants and user-tunable parameters
- `web/triple_ksampler_ui.js`: JavaScript for dynamic UI parameter visibility
### Testing
To test the nodes are properly loaded:
@@ -163,10 +215,12 @@ To test the nodes are properly loaded:
## Version History
### v0.3.1 (Current)
- **NEW FEATURE**: Dry Run Mode - Boolean parameter in advanced node for testing configurations without sampling
- **ARCHITECTURE**: Major inheritance refactor - Base → Advanced → Simple eliminates ~400 lines of code duplication
- **COMPATIBILITY**: Enhanced KJNodes compatibility with configurable dual-layer fix system
- **DEVELOPMENT**: Comprehensive test suite with VSCode integration and type stub support
- **BREAKING**: Standardized error handling with consistent 'fail fast' approach
- **Enhanced Error Messages**: All configuration conflicts now raise ValueError with actionable guidance
- **Documentation Consistency**: Fixed simple node docstring, updated error message examples with suggestions
- **Development Infrastructure**: Implemented dev-only testing strategy with comprehensive validation tests
- **Code Quality**: Simplified line break usage, enhanced parameter validation with edge case coverage
### v0.3.0
-6
View File
@@ -12,12 +12,6 @@ behavior without changing the core implementation.
# base_steps = lightning_start * ceil(MIN_TOTAL_STEPS / lightning_steps)
MIN_TOTAL_STEPS = 20
# Dry run mode for testing and debugging
# When enabled, performs all calculations and logging but skips actual sampling
# operations. Returns mock latent data. Useful for testing parameter logic
# without running expensive sampling. Should remain False for production use.
ENABLE_DRY_RUN = False
# KJNodes compatibility fix toggle
# When enabled, applies a temporary monkey-patch to fix KJNodes transformer_options
# compatibility issue. Disable this once KJNodes is updated to support the parameter.
+19 -5
View File
@@ -23,7 +23,7 @@ import nodes
import torch
from comfy_extras.nodes_model_advanced import ModelSamplingSD3
from .constants import MIN_TOTAL_STEPS, LOGGER_PREFIX, DEFAULT_BOUNDARY_T2V, DEFAULT_BOUNDARY_I2V, ENABLE_DRY_RUN, ENABLE_KJNODES_COMPATIBILITY_FIX
from .constants import MIN_TOTAL_STEPS, LOGGER_PREFIX, DEFAULT_BOUNDARY_T2V, DEFAULT_BOUNDARY_I2V, ENABLE_KJNODES_COMPATIBILITY_FIX
# Configure module logger
logger = logging.getLogger(__name__)
@@ -203,6 +203,7 @@ class TripleKSamplerWan22Base:
end_at_step: int,
add_noise: bool,
return_with_leftover_noise: bool,
dry_run: bool = False,
stage_name: str = "Sampler",
stage_info: str = None
) -> Tuple[Dict[str, torch.Tensor]]:
@@ -223,6 +224,7 @@ class TripleKSamplerWan22Base:
end_at_step: Ending step (exclusive).
add_noise: Whether to add initial noise.
return_with_leftover_noise: Whether to return with remaining noise.
dry_run: Enable dry run mode for testing configurations without sampling.
stage_name: Stage identifier for logging.
stage_info: Optional stage information to log right before sampling.
@@ -248,7 +250,7 @@ class TripleKSamplerWan22Base:
)
logger.info("%s: %s - %s", stage_name, stage_type, stage_info)
if ENABLE_DRY_RUN:
if dry_run:
return (latent,)
# Apply KJNodes compatibility fix if enabled
@@ -500,6 +502,13 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
)
}
),
"dry_run": (
"BOOLEAN",
{
"default": False,
"tooltip": "Enable dry run mode for testing configurations without actual sampling."
}
),
}
}
@@ -532,7 +541,8 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
scheduler: str,
switch_strategy: str,
switch_boundary: float = 0.875,
switch_step: int = -1
switch_step: int = -1,
dry_run: bool = False
) -> Tuple[Dict[str, torch.Tensor]]:
"""
Execute triple-stage cascade sampling with comprehensive logging.
@@ -556,6 +566,7 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
switch_strategy: Strategy for lightning model switching.
switch_boundary: Sigma boundary for manual boundary strategy (optional).
switch_step: Manual switch step for manual switch step strategy (optional).
dry_run: Enable dry run mode for testing configurations without sampling (optional).
Returns:
Tuple containing final latent dictionary.
@@ -577,7 +588,7 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
logger.info(" switch_boundary=%.3f", switch_boundary)
elif switch_strategy == "Manual switch step":
logger.info(" switch_step=%d", switch_step)
logger.info("Constants: MIN_TOTAL_STEPS=%d, DRY_RUN=%s", MIN_TOTAL_STEPS, ENABLE_DRY_RUN)
logger.info("Constants: MIN_TOTAL_STEPS=%d, DRY_RUN=%s", MIN_TOTAL_STEPS, dry_run)
bare_logger.info("")
# Validate parameters
@@ -787,6 +798,7 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
end_at_step=base_steps,
add_noise=stage1_add_noise,
return_with_leftover_noise=True,
dry_run=dry_run,
stage_name="Stage 1",
stage_info=stage1_info
)
@@ -816,6 +828,7 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
end_at_step=switch_step_final,
add_noise=stage2_add_noise,
return_with_leftover_noise=True,
dry_run=dry_run,
stage_name="Stage 2",
stage_info=stage2_info
)
@@ -839,6 +852,7 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
end_at_step=lightning_steps,
add_noise=stage3_add_noise,
return_with_leftover_noise=False,
dry_run=dry_run,
stage_name="Stage 3",
stage_info=stage3_info
)
@@ -847,7 +861,7 @@ class TripleKSamplerWan22LightningAdvanced(TripleKSamplerWan22Base):
bare_logger.info("")
# Log dry run summary if enabled
if ENABLE_DRY_RUN:
if dry_run:
logger.info("[DRY RUN] Complete - All calculations performed, no actual sampling executed")
return stage3_result