diff --git a/README.md b/README.md index be10a46..fe69371 100644 --- a/README.md +++ b/README.md @@ -1,165 +1,331 @@ -# 🐓 ComfyUI Bawk Nodes +
+

🐓 ComfyUI Bawk Nodes v2.0.0

-A collection of ComfyUI nodes **focused primarily on FLUX model workflows**. While some nodes may work with other diffusion models, all development and testing is done specifically with FLUX architectures to ensure optimal performance and compatibility. +**A collection of FLUX-optimized ComfyUI nodes for efficient AI image generation.** -## ⚡ FLUX-First Design Philosophy +
-Our nodes are built with FLUX models as the primary target: -- **Optimized for FLUX**: All nodes designed around FLUX model architecture -- **FLUX-Tested**: Extensive testing with FLUX Dev, FLUX Schnell, and FLUX variants -- **Other Models**: May work with SDXL/SD1.5 but **not officially supported or tested** +![Image Description](previews/efficient-workflow-preview.png) -> **Note**: While these nodes might function with other diffusion models, we only guarantee compatibility and provide support for FLUX models. Use with other models at your own discretion. +--- -## Current Nodes +## 🎯 **What's New in v2.0.0** -### 🚀 Diffusion Model Loader (Advanced) -A powerful diffusion model loader specifically optimized for FLUX models with separate component loading. +**Major Rewrite**: Complete FLUX-first redesign with modular architecture and workflow consolidation. -**FLUX-Optimized Features:** -- **FLUX Model Support**: Loads models from `diffusion_models` directory (FLUX format) -- **Separate Component Loading**: Independent VAE and dual CLIP text encoder support -- **FLUX Weight Types**: Support for FP8, FP16, BF16, and FP32 optimized for FLUX -- **FLUX CLIP Types**: Proper `CLIPType.FLUX` handling for T5 + CLIP-L encoders -- **Model String Output**: Returns clean model name for workflow identification -- **FLUX Validation**: Input validation designed around FLUX model requirements +- 🎲 **Enhanced Wildcard Encoder** with 6 LoRA slots +- 🐓 **All-in-One BawkSampler** with integrated VAE decoding +- 📁 **Modular Architecture** for better maintainability +- ⚡ **Streamlined Workflows** - fewer nodes, more power +- 💾 **Enhanced Image Saver** with prompt saving -## 📦 Installation +--- + +## 🚀 **Node Collection Overview** + +| Node | Description | Category | +|------|-------------|----------| +| 🚀 **Diffusion Model Loader** | Advanced FLUX-optimized model loading | Loaders | +| 🎲 **FLUX Wildcard Encoder** | Text encoding + 6 LoRA slots + wildcards | Conditioning | +| 🐓 **Bawk Sampler** | All-in-one latent generation, sampling & VAE decode | Sampling | +| 💾 **FLUX Image Saver** | Organized saving with metadata & prompt files | Image | +| 📝 **FLUX Prompt Saver** | Standalone prompt archiving | Text | + +--- + +## 🔥 **Complete FLUX Workflow** + +**Before BawkNodes (5+ nodes):** +``` +CheckpointLoader → LoraLoader → CLIPTextEncode → EmptyLatent → KSampler → VAEDecode → SaveImage +``` + +**After BawkNodes (3 nodes):** +``` +🚀 DiffusionModelLoader → 🎲 FluxWildcardEncode → 🐓 BawkSampler → 💾 FluxImageSaver +``` + +**60% fewer nodes, 100% of the power!** + +--- + +## 📦 **Installation** ### Method 1: ComfyUI Manager (Recommended) 1. Open ComfyUI Manager -2. Search for "Bawk Nodes" or "ComfyUI-BawkNodes" +2. Search for "Bawk Nodes" 3. Click Install +4. Restart ComfyUI ### Method 2: Manual Installation -1. Navigate to your ComfyUI custom nodes directory: - ```bash - cd ComfyUI/custom_nodes/ - ``` - -2. Clone this repository: - ```bash - git clone https://github.com/juddisjudd/ComfyUI-BawkNodes.git - ``` - -3. Restart ComfyUI - -## 🎯 Usage - -### Diffusion Model Loader (Advanced) -1. Add the "🚀 Diffusion Model Loader (Advanced)" node to your workflow -2. Select your diffusion model from the dropdown -3. Choose your VAE (or use "baked VAE" for none) -4. Select your text encoders (CLIP models) -5. Choose your preferred weight data type -6. Connect the outputs to your workflow - -### Advanced Configuration - -#### Weight Data Types -- **default**: Automatic selection based on hardware -- **fp8_e4m3fn**: 8-bit floating point (requires modern GPUs) -- **fp8_e4m3fn_fast**: Optimized 8-bit variant -- **fp8_e5m2**: Alternative 8-bit format -- **fp16**: 16-bit floating point (most common) -- **bf16**: Brain floating point 16-bit -- **fp32**: Full precision 32-bit - -#### Separate VAE Loading -- Select "baked VAE" to use the VAE included in your model -- Choose a specific VAE file to override the model's VAE - -#### Dual Text Encoders (FLUX Models) -- **clip_name1**: First text encoder (typically T5 for FLUX) -- **clip_name2**: Second text encoder (typically CLIP-L for FLUX) - -## 🔧 Node Inputs - -| Input | Type | Default | Description | -|-------|------|---------|-------------| -| `model_name` | STRING | - | Checkpoint file name (required) | -| `weight_dtype` | COMBO | "default" | Weight data type | -| `vae_name` | COMBO | "baked VAE" | VAE model (optional) | -| `clip_name1` | COMBO | "none" | First text encoder (optional) | -| `clip_name2` | COMBO | "none" | Second text encoder (optional) | - -## 📤 Node Outputs - -| Output | Type | Description | -|--------|------|-------------| -| `MODEL` | MODEL | Loaded diffusion model | -| `VAE` | VAE | Variational autoencoder | -| `CLIP` | CLIP | Text encoder(s) | -| `MODEL_STRING` | STRING | Model information summary | - - -## 🔮 Planned Nodes (FLUX-Focused) - -***TBD*** -*All future nodes will maintain our FLUX-first design philosophy* - -## Example Workflows - -### Basic FLUX Workflow -``` -Diffusion Model Loader (Advanced) -├── model_name: "flux1-dev.safetensors" -├── weight_dtype: "fp8_e4m3fn" -├── vae_name: "ae.safetensors" -├── clip_name1: "t5xxl_fp16.safetensors" -└── clip_name2: "clip_l.safetensors" +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/juddisjudd/ComfyUI-BawkNodes.git +# Restart ComfyUI ``` -## ⚡ FLUX Performance Tips +--- -1. **Hardware Optimization**: - - Use FP8 data types on RTX 4000+ series GPUs for maximum VRAM efficiency - - FLUX models benefit significantly from modern GPU architectures - - Ensure adequate VRAM (12GB+ recommended for FLUX Dev) +## 🎲 **Node Details** -2. **FLUX-Specific Settings**: - - Use `fp8_e4m3fn` for best quality/memory balance - - Use `fp8_e4m3fn_fast` for maximum speed - - T5 + CLIP-L combination provides optimal text understanding +### 🚀 **Diffusion Model Loader (Advanced)** -3. **Model Organization**: - - Keep FLUX models in `ComfyUI/models/diffusion_models/` - - Use `ae.safetensors` VAE for all FLUX variants - - Separate text encoders allow better memory management +**FLUX-optimized model loading with advanced features.** -## Troubleshooting +**Features:** +- Multiple model formats (FLUX, SDXL, SD1.5) +- Flexible weight data types (fp8, fp16, bf16, fp32) +- Separate VAE and CLIP loading +- Multiple directory support -### FLUX-Specific Issues +**Inputs:** +- `model_name` - Model from diffusion_models folder +- `vae_name` - VAE or "baked VAE" +- `clip_name1/2` - CLIP models for FLUX +- `weight_dtype` - Precision optimization -**"Model not found"** -- Ensure FLUX models are in `ComfyUI/models/diffusion_models/` (NOT checkpoints!) -- VAE files go in `ComfyUI/models/vae/` -- Text encoders go in `ComfyUI/models/text_encoders/` or `ComfyUI/models/clip/` +**Outputs:** `MODEL`, `VAE`, `CLIP`, `MODEL_STRING` -**Matrix multiplication errors with FLUX samplers** -- This usually means a checkpoint was loaded instead of a diffusion model -- Ensure your FLUX model is in `diffusion_models` directory -- Our loader is specifically designed to prevent this issue +--- -**"Insufficient VRAM" with FLUX** -- Try `fp8_e4m3fn` or `fp8_e4m3fn_fast` weight types -- FLUX models are large - consider using smaller variants for lower VRAM -- Ensure no other models are loaded in memory +### 🎲 **FLUX Wildcard Encoder** -**Empty dropdowns** -- The node only shows files that actually exist in the correct directories -- Check that your FLUX files are in the proper locations -- This loader ONLY shows diffusion models (FLUX format) +**Enhanced text encoder with 6 LoRA slots and wildcard support.** -**Not compatible with other models** -- Remember: These nodes are designed specifically for FLUX -- Other diffusion models may not work correctly -- Use standard ComfyUI loaders for non-FLUX models +**Features:** +- **Wildcard Processing**: `{option1|option2|option3}` syntax +- **6 LoRA Slots**: Individual enable/disable toggles +- **Fuzzy LoRA Matching**: Flexible file resolution +- **FLUX Optimization**: 16-channel conditioning +**Inputs:** +- `model`, `clip` - From model loader +- `prompt` - Text with wildcard support +- `wildcard_seed` - Seed for consistent wildcard selection +- `lora_X_on` - Enable/disable each LoRA (X = 1-6) +- `lora_X_name` - LoRA selection dropdown +- `lora_X_strength` - Strength adjustment (-10.0 to +10.0) -## Support +**Outputs:** `MODEL`, `CLIP`, `CONDITIONING`, `PROMPT_OUT` -- **Issues**: [GitHub Issues](https://github.com/juddisjudd/ComfyUI-BawkNodes/issues) +**Example Prompt with Wildcards:** +``` +A {beautiful|stunning|gorgeous} {cat|dog|bird} in a {forest|garden|meadow}, +{photorealistic|artistic|stylized} style +``` + +--- + +### 🐓 **Bawk Sampler (All-in-One)** + +**Complete latent generation, sampling, and VAE decoding in one node.** + +**Features:** +- **Smart Resolution Presets**: Pre-configured FLUX-optimized resolutions +- **Custom Resolution Support**: Manual width/height with 64px alignment +- **Advanced FLUX Sampling**: All FLUX-specific parameters +- **Integrated VAE Decoding**: Direct image output +- **Batch Generation**: Up to 64 images at once + +**Key Inputs:** +- `model`, `conditioning`, `vae` - From previous nodes +- `resolution` - Smart presets or custom +- `batch_size` - Number of images (default: 4) +- `sampler` - Sampling method (default: euler) +- `scheduler` - Noise schedule (default: beta) +- `steps` - Sampling steps (default: 30) +- `guidance` - FLUX guidance scale (default: 3.5) +- `max_shift` - FLUX max shift (default: 0.5) +- `base_shift` - FLUX base shift (default: 0.3) + +**Resolution Presets:** +- `FHD 16:9 - 1920x1080` (default) +- `Medium Square - 1024x1024` +- `Portrait 9:16 - 1080x1920` +- `Ultra-wide - 1792x768` +- And many more... + +**Outputs:** `IMAGE`, `LATENT` + +--- + +### 💾 **FLUX Image Saver** + +**Organized image saving with metadata and prompt archiving.** + +**Features:** +- **Smart Folder Organization**: `[MODEL]-DD-MM-YYYY` structure +- **Multiple Formats**: PNG, JPG, WebP support +- **Metadata Embedding**: PNG metadata support +- **Prompt File Saving**: Separate `.txt` files with processed prompts +- **JSON Metadata**: Complete generation parameters + +**Inputs:** +- `images` - From BawkSampler +- `model_string` - From model loader +- `processed_prompt` - From wildcard encoder +- `save_prompt` - Enable prompt file saving (default: True) +- `format` - Image format (PNG/JPG/WebP) +- `quality` - Compression quality (1-100) + +**File Output Example:** +``` +ComfyUI/output/[FLUX_Model]-01-08-2025/ +├── flux_image_01-08-2025_14-30-15_001.png +├── flux_image_01-08-2025_14-30-15_002.png +├── flux_image_01-08-2025_14-30-15_prompt.txt +├── flux_image_01-08-2025_14-30-15_001_metadata.json +└── flux_image_01-08-2025_14-30-15_002_metadata.json +``` + +--- + +### 📝 **FLUX Prompt Saver** + +**Standalone prompt and parameter archiving.** + +**Features:** +- **JSON Format**: Structured data storage +- **Complete Parameters**: All generation settings +- **Organized Storage**: Matches image saver folder structure +- **Workflow Integration**: Links with other BawkNodes + +--- + +## 🛠️ **Advanced Usage** + +### **Wildcard Examples** + +**Basic Wildcards:** +``` +A {red|blue|green} car in the {city|countryside} +``` + +**Nested Concepts:** +``` +{A majestic|An elegant|A powerful} {dragon|phoenix|griffin} +{soaring through|perched upon|emerging from} {clouds|mountains|flames} +``` + +**Style Variations:** +``` +Portrait of a woman, {photorealistic|oil painting|digital art|watercolor} style, +{studio lighting|natural lighting|dramatic lighting} +``` + +### **LoRA Management** + +**Best Practices:** +1. **Enable LoRAs individually** for precise control +2. **Use strength between 0.5-1.5** for most LoRAs +3. **Combine complementary LoRAs** (style + subject) +4. **Test different combinations** for unique results + +**Example LoRA Setup:** +- LoRA 1: `realistic_skin_v2.safetensors` (0.8) +- LoRA 2: `dramatic_lighting.safetensors` (0.6) +- LoRA 3: `detail_enhancer.safetensors` (0.4) + +### **Resolution Guidelines** + +**Recommended Presets:** +- **Square**: `Medium Square - 1024x1024` +- **Landscape**: `FHD 16:9 - 1920x1080` +- **Portrait**: `Portrait 9:16 - 1080x1920` +- **Widescreen**: `Ultra-wide - 1792x768` + +**Custom Resolution Rules:** +- Must be multiples of 64 pixels +- Keep total pixel count reasonable (<4MP for speed) +- Consider VRAM limitations for large batches + +--- + +## 🔧 **Configuration** + +### **Model Setup** + +1. **FLUX Models**: Place in `models/diffusion_models/` +2. **VAE Files**: Place in `models/vae/` +3. **CLIP Models**: Place in `models/text_encoders/` +4. **LoRA Files**: Place in `models/loras/` + +### **Recommended Settings** + +**For Speed:** +- Resolution: `Medium Square - 1024x1024` +- Batch Size: `4` +- Steps: `20-25` +- Sampler: `euler` + +**For Quality:** +- Resolution: `FHD 16:9 - 1920x1080` +- Batch Size: `1-2` +- Steps: `30-40` +- Sampler: `dpmpp_2m` + +**For Experimentation:** +- Use wildcards with high variation +- Enable multiple LoRAs +- Try different guidance scales (2.0-5.0) + +--- + +## **Troubleshooting** + +### **Common Issues** + +**Node Not Appearing:** +```bash +# Check ComfyUI console for errors +# Ensure all files are in correct directories +# Restart ComfyUI completely +``` + +**LoRA Not Loading:** +- Check file is in `models/loras/` +- Verify file isn't corrupted +- Check console for specific error messages + +**Memory Issues:** +- Reduce batch size +- Use lower resolution +- Enable `fp8` weight dtype in loader + +**Generation Errors:** +- Verify all connections are correct +- Check that VAE is connected to BawkSampler +- Ensure CLIP and MODEL are from same loader + +### **Performance Optimization** + +**VRAM Usage:** +- Use `fp8_e4m3fn_fast` for weight dtype +- Reduce batch size for large images +- Close other GPU applications + +**Speed Improvements:** +- Use `euler` sampler with `beta` scheduler +- Reduce step count (20-30 is often sufficient) +- Use medium resolution presets + +--- + +## 📄 **License** + +GPL-3.0 license - see [LICENSE](LICENSE) file for details. + +--- + +## 🙏 **Acknowledgments** +- **[rgthree](https://github.com/rgthree/rgthree-comfy)** - Inspiration for dynamic UI patterns + +--- + +## 🔗 **Links** + +- **Comfy Registry**: [Comfy-Registry](https://registry.comfy.org/publishers/judd/nodes/comfyui-bawknodes) +- **Issues**: [Report Bugs](https://github.com/juddisjudd/ComfyUI-BawkNodes/issues) --- diff --git a/__init__.py b/__init__.py index b32e2d6..40021bb 100644 --- a/__init__.py +++ b/__init__.py @@ -3,29 +3,43 @@ ComfyUI Bawk Nodes - Node Registration File: __init__.py """ -from .nodes import DiffusionModelLoader +from .nodes import ( + DiffusionModelLoader, + FluxImageSaver, + FluxPromptSaver, + FluxWildcardEncode, + BawkSampler +) -# Node class mappings for ComfyUI discovery NODE_CLASS_MAPPINGS = { "DiffusionModelLoader": DiffusionModelLoader, + "FluxImageSaver": FluxImageSaver, + "FluxPromptSaver": FluxPromptSaver, + "FluxWildcardEncode": FluxWildcardEncode, + "BawkSampler": BawkSampler, } -# Display names in the ComfyUI interface NODE_DISPLAY_NAME_MAPPINGS = { "DiffusionModelLoader": "🚀 Diffusion Model Loader (Advanced)", + "FluxImageSaver": "💾 FLUX Image Saver", + "FluxPromptSaver": "📝 FLUX Prompt Saver", + "FluxWildcardEncode": "🎲 FLUX Wildcard Encoder", + "BawkSampler": "🐓 Bawk Sampler (All-in-One)", } -# Export for ComfyUI __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] -# Package metadata -__version__ = "1.0.0" +__version__ = "2.0.0" __author__ = "Bawk Nodes" -__description__ = "A collection of useful ComfyUI nodes for enhanced workflows" +__description__ = "A complete collection of FLUX-optimized ComfyUI nodes for enhanced workflows" -# Optional: Print loading message print(f"🐓 ComfyUI Bawk Nodes v{__version__} loaded successfully!") +print(" 🎉 Major Update - Complete FLUX Workflow Suite!") print(" Current nodes:") print(" • 🚀 Diffusion Model Loader (Advanced) - FLUX-optimized model loading") -print(" • More nodes coming soon...") +print(" • 🎲 FLUX Wildcard Encoder - Text encoding with wildcard support and 6 LoRA slots") +print(" • 💾 FLUX Image Saver - Organized image saving with metadata") +print(" • 📝 FLUX Prompt Saver - Save prompts and generation parameters") +print(" • 🐓 Bawk Sampler (All-in-One) - Combined latent optimizer and sampler") +print(" 📦 Modular architecture for easy maintenance and debugging") print(" Visit: https://github.com/juddisjudd/ComfyUI-BawkNodes") \ No newline at end of file diff --git a/nodes.py b/nodes.py index 8e792f5..123351e 100644 --- a/nodes.py +++ b/nodes.py @@ -1,339 +1,18 @@ """ -ComfyUI Bawk Nodes - Diffusion Model Loader +ComfyUI Bawk Nodes - Main nodes file (clean imports only) File: nodes.py """ -import os -import torch -import comfy.sd -import comfy.model_management as mm -import comfy.utils -import folder_paths -from typing import Optional, Tuple, Dict, Any, Union -from .modules.validation import ValidationMixin +from .nodes.diffusion_model_loader import DiffusionModelLoader +from .nodes.flux_wildcard_encode import FluxWildcardEncode +from .nodes.flux_image_saver import FluxImageSaver +from .nodes.flux_prompt_saver import FluxPromptSaver +from .nodes.bawk_sampler import BawkSampler - -class DiffusionModelLoader(ValidationMixin): - """ - Advanced Diffusion Model Loader with support for: - - Multiple model formats (FLUX, SDXL, SD1.5) - - Flexible weight data types (fp8, fp16, bf16, fp32) - - Separate VAE and CLIP loading - - Multiple directory support - """ - - @classmethod - def INPUT_TYPES(cls): - # Match the exact format of the working DoomFluxLoader - return { - "required": { - "model_name": (folder_paths.get_filename_list("diffusion_models"),), - "vae_name": (folder_paths.get_filename_list("vae") + ["baked VAE"],), - "clip_name1": (folder_paths.get_filename_list("text_encoders") + ["none"],), - "clip_name2": (folder_paths.get_filename_list("text_encoders") + ["none"],), - "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],), - } - } - - RETURN_TYPES = ("MODEL", "VAE", "CLIP", "STRING") - RETURN_NAMES = ("MODEL", "VAE", "CLIP", "MODEL_STRING") - FUNCTION = "load_diffusion_model" - CATEGORY = "BawkNodes/loaders" - DESCRIPTION = "Advanced diffusion model loader for FLUX and other modern architectures" - - @classmethod - def VALIDATE_INPUTS(cls, **kwargs): - """Simple validation matching the working loader pattern""" - errors = [] - - # Validate model exists in diffusion_models directory - model_name = kwargs.get("model_name") - if model_name: - try: - folder_paths.get_full_path_or_raise("diffusion_models", model_name) - except: - errors.append(f"Model '{model_name}' not found in diffusion_models directory") - - # Validate VAE if specified - vae_name = kwargs.get("vae_name", "baked VAE") - if vae_name and vae_name != "baked VAE": - try: - folder_paths.get_full_path_or_raise("vae", vae_name) - except: - errors.append(f"VAE '{vae_name}' not found in vae directory") - - # Validate CLIP models if specified - clip1 = kwargs.get("clip_name1", "none") - if clip1 and clip1 != "none": - try: - folder_paths.get_full_path_or_raise("text_encoders", clip1) - except: - errors.append(f"CLIP model '{clip1}' not found in text_encoders directory") - - clip2 = kwargs.get("clip_name2", "none") - if clip2 and clip2 != "none": - try: - folder_paths.get_full_path_or_raise("text_encoders", clip2) - except: - errors.append(f"CLIP model '{clip2}' not found in text_encoders directory") - - return "; ".join(errors) if errors else True - - def load_diffusion_model( - self, - model_name: str, - weight_dtype: str = "default", - vae_name: str = "baked VAE", - clip_name1: str = "none", - clip_name2: str = "none" - ) -> Tuple[Any, Any, Any, str]: - """ - Main loading function for diffusion models (based on working DoomFluxLoader) - """ - try: - # Handle "none" and "baked VAE" values like the working loader - vae_name = vae_name if vae_name and vae_name != "baked VAE" else None - clip_name1 = clip_name1 if clip_name1 and clip_name1 != "none" else None - clip_name2 = clip_name2 if clip_name2 and clip_name2 != "none" else None - weight_dtype = weight_dtype if weight_dtype else "default" - - # Set up model options based on weight_dtype (like DoomFluxLoader) - model_options = {} - if weight_dtype == "fp8_e4m3fn": - model_options["dtype"] = torch.float8_e4m3fn - elif weight_dtype == "fp8_e4m3fn_fast": - model_options["dtype"] = torch.float8_e4m3fn - model_options["fp8_optimizations"] = True - elif weight_dtype == "fp8_e5m2": - model_options["dtype"] = torch.float8_e5m2 - - # Load diffusion model using the exact same method as DoomFluxLoader - print(f"[DiffusionModelLoader] Loading diffusion model: {model_name}") - model_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name) - model = comfy.sd.load_diffusion_model(model_path, model_options=model_options) - - if model is None: - raise RuntimeError("Failed to load diffusion model - output was None") - - # Load VAE (exactly like DoomFluxLoader) - vae = None - if vae_name: - print(f"[DiffusionModelLoader] Loading VAE: {vae_name}") - vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) - vae_sd = comfy.utils.load_torch_file(vae_path) - vae = comfy.sd.VAE(sd=vae_sd) - - # Load CLIP models (exactly like DoomFluxLoader) - clip_paths = [] - if clip_name1: - clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name1)) - if clip_name2: - clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name2)) - - clip = None - if clip_paths: - print(f"[DiffusionModelLoader] Loading CLIP models") - clip = comfy.sd.load_clip( - ckpt_paths=clip_paths, - embedding_directory=folder_paths.get_folder_paths("embeddings"), - clip_type=comfy.sd.CLIPType.FLUX - ) - - # Generate simple model name string - model_string = os.path.splitext(model_name)[0] - - print(f"[DiffusionModelLoader] Successfully loaded diffusion model: {model_string}") - return (model, vae, clip, model_string) - - except Exception as e: - error_msg = f"Failed to load diffusion model '{model_name}': {str(e)}" - print(f"[DiffusionModelLoader] Error: {error_msg}") - raise RuntimeError(error_msg) - - def _create_cache_key(self, model_name: str, weight_dtype: str, vae_name: str, - clip_name1: str, clip_name2: str) -> str: - """Create unique cache key for this configuration""" - return f"{model_name}|{weight_dtype}|{vae_name}|{clip_name1}|{clip_name2}" - - def _validate_model_path(self, model_name: str) -> str: - """Validate model exists and is accessible, checking multiple directories""" - model_path = None - - # Try checkpoints directory first - model_path = folder_paths.get_full_path("checkpoints", model_name) - - # If not found, try diffusion_models directory - if model_path is None: - model_path = folder_paths.get_full_path("diffusion_models", model_name) - - if model_path is None: - raise FileNotFoundError(f"Model not found in checkpoints or diffusion_models directories: {model_name}") - - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model file does not exist: {model_path}") - - # Basic file size validation - file_size = os.path.getsize(model_path) - if file_size < 1024 * 1024: # Less than 1MB - raise ValueError(f"Model file appears corrupted (too small): {model_path}") - - return model_path - - def _select_dtype(self, weight_dtype: str, device: torch.device) -> torch.dtype: - """Select appropriate data type based on user preference and hardware""" - if weight_dtype == "default": - return self._auto_select_dtype(device) - - dtype_map = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, - } - - # Handle FP8 types if supported - if weight_dtype.startswith("fp8_"): - if self._supports_fp8(device): - if weight_dtype == "fp8_e4m3fn": - return torch.float8_e4m3fn - elif weight_dtype == "fp8_e5m2": - return torch.float8_e5m2 - else: - print(f"FP8 not supported on device {device}, falling back to fp16") - return torch.float16 - - selected_dtype = dtype_map.get(weight_dtype, torch.float32) - - # Validate dtype compatibility - if not self.validate_dtype_compatibility(weight_dtype, device): - print(f"Dtype {weight_dtype} not compatible with {device}, using auto selection") - return self._auto_select_dtype(device) - - return selected_dtype - - def _auto_select_dtype(self, device: torch.device) -> torch.dtype: - """Automatically select best dtype for device""" - if self._supports_fp8(device): - free_memory = mm.get_free_memory(device) - if free_memory < 6 * 1024**3: # Less than 6GB VRAM - return torch.float8_e4m3fn - - if mm.should_use_fp16(device): - return torch.float16 - - if device.type == "cuda" and torch.cuda.is_bf16_supported(): - return torch.bfloat16 - - return torch.float32 - - def _supports_fp8(self, device: torch.device) -> bool: - """Check if device supports FP8 computation""" - if device.type != "cuda": - return False - - try: - compute_capability = torch.cuda.get_device_capability(device) - return compute_capability >= (8, 9) # Ada Lovelace or newer - except: - return False - - def _load_checkpoint_components( - self, model_path: str, dtype: torch.dtype, device: torch.device - ) -> Tuple[Any, Any, Any]: - """Load main checkpoint components using ComfyUI's standard API""" - try: - # Use ComfyUI's memory management - mm.soft_empty_cache() # Clear cache before loading - - # Use the same API as CheckpointLoaderSimple - out = comfy.sd.load_checkpoint_guess_config( - model_path, - output_vae=True, - output_clip=True, - embedding_directory=folder_paths.get_folder_paths("embeddings") - ) - - if out is None or len(out) < 3: - raise RuntimeError("Failed to load checkpoint components") - - return out[:3] # model, clip, vae - - except Exception as e: - if "Could not detect model type" in str(e): - raise RuntimeError( - f"Unsupported model format. This may be a FLUX model that requires " - f"separate component loading. Try using individual loaders for " - f"UNet/CLIP/VAE components." - ) - raise RuntimeError(f"Error loading checkpoint: {str(e)}") - - def _load_separate_vae(self, vae_name: str): - """Load separate VAE component""" - vae_path = folder_paths.get_full_path("vae", vae_name) - if vae_path is None: - raise FileNotFoundError(f"VAE not found: {vae_name}") - - try: - # Use ComfyUI's standard VAE loading - vae_sd = comfy.utils.load_torch_file(vae_path, safe_load=True) - vae = comfy.sd.VAE(sd=vae_sd) - return vae - except Exception as e: - raise RuntimeError(f"Failed to load VAE '{vae_name}': {str(e)}") - - def _load_custom_clip(self, clip_name1: str, clip_name2: str): - """Load custom CLIP models (for FLUX dual text encoders)""" - clip_paths = [] - - if clip_name1 != "none": - # Try clip directory first, then text_encoders - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - if clip_path1 is None: - clip_path1 = folder_paths.get_full_path("text_encoders", clip_name1) - - if clip_path1 is None: - raise FileNotFoundError(f"CLIP model not found in clip or text_encoders directories: {clip_name1}") - clip_paths.append(clip_path1) - - if clip_name2 != "none": - # Try clip directory first, then text_encoders - clip_path2 = folder_paths.get_full_path("clip", clip_name2) - if clip_path2 is None: - clip_path2 = folder_paths.get_full_path("text_encoders", clip_name2) - - if clip_path2 is None: - raise FileNotFoundError(f"CLIP model not found in clip or text_encoders directories: {clip_name2}") - clip_paths.append(clip_path2) - - if not clip_paths: - return None - - try: - # Use ComfyUI's standard CLIP loading - clip = comfy.sd.load_clip( - ckpt_paths=clip_paths, - embedding_directory=folder_paths.get_folder_paths("embeddings") - ) - return clip - except Exception as e: - raise RuntimeError(f"Failed to load CLIP models: {str(e)}") - - def _generate_model_info( - self, model_name: str, dtype: torch.dtype, device: torch.device, - vae_name: str, clip_name1: str, clip_name2: str - ) -> str: - """Generate model name string (just the base model name without extension)""" - # Extract base name without extension - base_name = os.path.splitext(model_name)[0] - return base_name - - def get_model_string_preview(self, model_name: str, weight_dtype: str = "default", - vae_name: str = "baked VAE", clip_name1: str = "none", - clip_name2: str = "none") -> str: - """ - Preview what the model string output will look like. - Returns just the base model name without extension. - """ - try: - return os.path.splitext(model_name)[0] - except Exception as e: - return f"Error: {str(e)}" \ No newline at end of file +__all__ = [ + "DiffusionModelLoader", + "FluxWildcardEncode", + "FluxImageSaver", + "FluxPromptSaver", + "BawkSampler" +] \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..a89a298 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,18 @@ +""" +Nodes package initialization +File: nodes/__init__.py +""" + +from .diffusion_model_loader import DiffusionModelLoader +from .flux_wildcard_encode import FluxWildcardEncode +from .flux_image_saver import FluxImageSaver +from .flux_prompt_saver import FluxPromptSaver +from .bawk_sampler import BawkSampler + +__all__ = [ + "DiffusionModelLoader", + "FluxWildcardEncode", + "FluxImageSaver", + "FluxPromptSaver", + "BawkSampler" +] \ No newline at end of file diff --git a/nodes/bawk_sampler.py b/nodes/bawk_sampler.py new file mode 100644 index 0000000..8bf2071 --- /dev/null +++ b/nodes/bawk_sampler.py @@ -0,0 +1,337 @@ +""" +BawkSampler - Combined Latent Generator and Advanced Sampler for FLUX +File: nodes/bawk_sampler.py +""" + +import torch +import math +import time +import logging +import comfy.samplers +import comfy.model_management as mm +from comfy.utils import ProgressBar +from comfy_extras.nodes_custom_sampler import Noise_RandomNoise, BasicScheduler, BasicGuider, SamplerCustomAdvanced +from comfy_extras.nodes_latent import LatentBatch +from comfy_extras.nodes_model_advanced import ModelSamplingFlux, ModelSamplingAuraFlow +from typing import Tuple, Dict, Any, List + + +def round_to_nearest_multiple(value, multiple): + """Rounds a value to the nearest multiple of 'multiple'.""" + if multiple <= 0: + return value + return int(round(value / multiple) * multiple) + + +def parse_string_to_list(input_string): + """Parse comma/newline separated string to list of numbers""" + try: + if not input_string: + return [] + items = input_string.replace('\n', ',').split(',') + result = [] + for item in items: + item = item.strip() + if not item: + continue + try: + num = float(item) + if num.is_integer(): + num = int(num) + result.append(num) + except ValueError: + continue + return result + except: + return [] + + +def conditioning_set_values(conditioning, values): + """Set conditioning values like guidance scale""" + c = [] + for t in conditioning: + n = [t[0], t[1].copy()] + for k, v in values.items(): + if k == "guidance": + n[1]['guidance_scale'] = v + c.append(tuple(n)) + return c + + +class BawkSampler: + """ + Combined latent generator and advanced sampler for FLUX models. + Generates optimized latent dimensions, then performs advanced sampling. + Designed to work with FluxWildcardEncode for prompt handling. + """ + + @classmethod + def INPUT_TYPES(cls): + # Combined resolution presets with aspect ratios built-in + # Format: "Name - WxH (aspect ratio)" + resolution_presets = [ + # 16:9 aspect ratio options + "HD 16:9 - 1024x576", + "FHD 16:9 - 1920x1080", + "2K 16:9 - 2048x1152", + + # 1:1 square options + "Small Square - 768x768", + "Medium Square - 1024x1024", + "Large Square - 1536x1536", + "XL Square - 2048x2048", + + # 3:2 photo aspect ratio + "Photo 3:2 - 1152x768", + "Photo 3:2 HD - 1728x1152", + + # 4:3 classic aspect ratio + "Classic 4:3 - 1024x768", + "Classic 4:3 HD - 1536x1152", + + # 9:16 portrait/mobile + "Portrait 9:16 - 576x1024", + "Mobile 9:16 - 1080x1920", + + # 2:3 portrait photo + "Portrait Photo - 768x1152", + + # 21:9 ultra-wide + "Ultra-wide - 1792x768", + + # Custom option + "Custom Resolution", + ] + + return { + "required": { + # Core inputs + "model": ("MODEL",), + "conditioning": ("CONDITIONING",), + "vae": ("VAE",), # Added VAE input for decoding + + # Latent generation + "resolution": (resolution_presets, { + "default": "FHD 16:9 - 1920x1080", + "tooltip": "Select resolution preset or custom option" + }), + "batch_size": ("INT", { + "default": 4, "min": 1, "max": 64, "step": 1, + "tooltip": "Number of images to generate" + }), + + # Sampling parameters + "seed": ("INT", { + "default": 0, "min": 0, "max": 0xffffffffffffffff, + "tooltip": "Random seed for generation" + }), + "sampler": (comfy.samplers.KSampler.SAMPLERS, { + "default": "euler", + "tooltip": "Sampling method for generation" + }), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, { + "default": "beta", + "tooltip": "Noise scheduling method" + }), + "steps": ("INT", { + "default": 30, "min": 1, "max": 100, "step": 1, + "tooltip": "Number of sampling steps" + }), + "guidance": ("FLOAT", { + "default": 3.5, "min": 0.0, "max": 20.0, "step": 0.1, + "tooltip": "Guidance scale for FLUX" + }), + "max_shift": ("FLOAT", { + "default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01, + "tooltip": "Max shift parameter for FLUX" + }), + "base_shift": ("FLOAT", { + "default": 0.3, "min": 0.0, "max": 2.0, "step": 0.01, + "tooltip": "Base shift parameter for FLUX" + }), + "denoise": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Denoise strength" + }), + + # Custom resolution toggle + "use_custom_resolution": ("BOOLEAN", { + "default": False, + "tooltip": "Enable custom width/height instead of presets" + }), + }, + "optional": { + # Custom resolution (only used when use_custom_resolution=True) + "custom_width": ("INT", { + "default": 1920, "min": 64, "max": 4096, "step": 64, + "tooltip": "Custom width (must be multiple of 64). Only used when custom resolution is enabled." + }), + "custom_height": ("INT", { + "default": 1080, "min": 64, "max": 4096, "step": 64, + "tooltip": "Custom height (must be multiple of 64). Only used when custom resolution is enabled." + }), + } + } + + RETURN_TYPES = ("IMAGE", "LATENT") + RETURN_NAMES = ("images", "latent") + FUNCTION = "generate_sample_and_decode" + CATEGORY = "BawkNodes/sampling" + DESCRIPTION = "Combined latent generator, sampler, and VAE decoder optimized for FLUX models" + + def generate_sample_and_decode( + self, + model, conditioning, vae, + resolution="FHD 16:9 - 1920x1080", batch_size=4, + seed=0, sampler="euler", scheduler="beta", steps=30, + guidance=3.5, max_shift=0.5, base_shift=0.3, denoise=1.0, + use_custom_resolution=False, custom_width=1920, custom_height=1080 + ): + """ + Generate optimized latent, perform FLUX sampling, and decode to images + """ + try: + # Step 1: Generate optimized latent + latent = self._generate_optimized_latent( + resolution, batch_size, use_custom_resolution, custom_width, custom_height + ) + + # Step 2: Perform sampling + sampled_latent = self._perform_flux_sampling( + model, conditioning, latent, seed, sampler, scheduler, + steps, guidance, max_shift, base_shift, denoise + ) + + # Step 3: Decode latent to images using VAE + decoded_images = self._decode_latent_to_images(vae, sampled_latent) + + print(f"[BawkSampler] Successfully generated, sampled, and decoded {batch_size} images") + + return (decoded_images, sampled_latent) + + except Exception as e: + error_msg = f"BawkSampler failed: {str(e)}" + print(f"[BawkSampler] Error: {error_msg}") + raise RuntimeError(error_msg) + + def _generate_optimized_latent( + self, resolution, batch_size, use_custom_resolution, custom_width, custom_height + ) -> Dict: + """Generate optimized empty latent for FLUX""" + + # Determine target dimensions + if use_custom_resolution: + # Use custom dimensions (ensure they're multiples of 64) + target_width = round_to_nearest_multiple(custom_width, 64) + target_height = round_to_nearest_multiple(custom_height, 64) + print(f"[BawkSampler] Using custom resolution: {target_width}x{target_height}") + else: + # Parse dimensions from preset string + target_width, target_height = self._parse_resolution_preset(resolution) + print(f"[BawkSampler] Using preset resolution: {resolution}") + + # FLUX-specific parameters + vae_scale_factor = 8 + latent_channels = 16 # FLUX uses 16 channels + + # Calculate latent dimensions + latent_width = target_width // vae_scale_factor + latent_height = target_height // vae_scale_factor + + # Generate latent tensor + try: + latent_tensor = torch.zeros([batch_size, latent_channels, latent_height, latent_width]) + latent = {"samples": latent_tensor} + except Exception as e: + raise RuntimeError(f"Error creating FLUX latent tensor [{batch_size}, {latent_channels}, {latent_height}, {latent_width}]: {e}") + + print(f"[BawkSampler] Generated FLUX latent: {latent_width}x{latent_height} " + f"(pixel: {target_width}x{target_height}, batch: {batch_size})") + + return latent + + def _parse_resolution_preset(self, resolution_preset): + """Parse width and height from resolution preset string""" + try: + # Extract dimensions from string like "FHD 16:9 - 1920x1080" + # Split by " - " and take the part after it + parts = resolution_preset.split(" - ") + if len(parts) < 2: + raise ValueError(f"Invalid resolution format: {resolution_preset}") + + dimensions = parts[1] # e.g., "1920x1080" + width_str, height_str = dimensions.split("x") + + target_width = int(width_str) + target_height = int(height_str) + + # Ensure dimensions are multiples of 64 for FLUX compatibility + target_width = round_to_nearest_multiple(target_width, 64) + target_height = round_to_nearest_multiple(target_height, 64) + + return target_width, target_height + + except Exception as e: + print(f"[BawkSampler] Warning: Could not parse resolution '{resolution_preset}': {e}") + # Fallback to default FHD resolution + return 1920, 1080 + + def _perform_flux_sampling( + self, model, conditioning, latent_image, seed, sampler, scheduler, + steps, guidance, max_shift, base_shift, denoise + ): + """Perform FLUX-optimized sampling""" + + # Single parameter sampling (no multi-parameter sweeps for cleaner UX) + noise_seed = seed + + # Initialize FLUX sampling components + basicscheduler = BasicScheduler() + basicguider = BasicGuider() + samplercustomadvanced = SamplerCustomAdvanced() + modelsamplingflux = ModelSamplingFlux() + randnoise = Noise_RandomNoise(noise_seed) + + # Get dimensions for model sampling + width = latent_image["samples"].shape[3] * 8 + height = latent_image["samples"].shape[2] * 8 + + # Apply FLUX model sampling + work_model = modelsamplingflux.patch(model, max_shift, base_shift, width, height)[0] + + # Set guidance in conditioning + cond = conditioning_set_values(conditioning, {"guidance": guidance}) + guider = basicguider.get_guider(work_model, cond)[0] + + # Create sampler object + samplerobj = comfy.samplers.sampler_object(sampler) + + # Generate sigmas + sigmas = basicscheduler.get_sigmas(work_model, scheduler, steps, denoise)[0] + + logging.info(f"FLUX Sampling: seed={noise_seed}, {sampler}_{scheduler}, " + f"steps={steps}, guidance={guidance}, max_shift={max_shift}, base_shift={base_shift}") + + # Perform FLUX sampling + latent = samplercustomadvanced.sample( + randnoise, guider, samplerobj, sigmas, latent_image + )[1] + + print(f"[BawkSampler] Completed FLUX sampling with {sampler}_{scheduler}") + return latent + + def _decode_latent_to_images(self, vae, latent): + """Decode latent to images using the provided VAE""" + try: + print(f"[BawkSampler] Decoding latent to images...") + + # Decode latent using VAE + decoded_images = vae.decode(latent["samples"]) + + print(f"[BawkSampler] Successfully decoded latent to images, shape: {decoded_images.shape}") + return decoded_images + + except Exception as e: + error_msg = f"Failed to decode latent: {str(e)}" + print(f"[BawkSampler] Error: {error_msg}") + raise RuntimeError(error_msg) \ No newline at end of file diff --git a/nodes/diffusion_model_loader.py b/nodes/diffusion_model_loader.py new file mode 100644 index 0000000..7474e3a --- /dev/null +++ b/nodes/diffusion_model_loader.py @@ -0,0 +1,156 @@ +""" +DiffusionModelLoader Node +File: nodes/diffusion_model_loader.py +""" + +import os +import torch +import comfy.sd +import comfy.utils +import folder_paths +from typing import Tuple, Any + +# Try importing validation from parent modules +try: + from ..modules.validation import ValidationMixin +except ImportError: + # Fallback if modules structure is different + class ValidationMixin: + pass + + +class DiffusionModelLoader(ValidationMixin): + """ + Advanced Diffusion Model Loader with support for: + - Multiple model formats (FLUX, SDXL, SD1.5) + - Flexible weight data types (fp8, fp16, bf16, fp32) + - Separate VAE and CLIP loading + - Multiple directory support + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_name": (folder_paths.get_filename_list("diffusion_models"),), + "vae_name": (folder_paths.get_filename_list("vae") + ["baked VAE"],), + "clip_name1": (folder_paths.get_filename_list("text_encoders") + ["none"],), + "clip_name2": (folder_paths.get_filename_list("text_encoders") + ["none"],), + "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],), + } + } + + RETURN_TYPES = ("MODEL", "VAE", "CLIP", "STRING") + RETURN_NAMES = ("MODEL", "VAE", "CLIP", "MODEL_STRING") + FUNCTION = "load_diffusion_model" + CATEGORY = "BawkNodes/loaders" + DESCRIPTION = "Advanced diffusion model loader for FLUX and other modern architectures" + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): + """Simple validation matching the working loader pattern""" + errors = [] + + # Validate model exists in diffusion_models directory + model_name = kwargs.get("model_name") + if model_name: + try: + folder_paths.get_full_path_or_raise("diffusion_models", model_name) + except: + errors.append(f"Model '{model_name}' not found in diffusion_models directory") + + # Validate VAE if specified + vae_name = kwargs.get("vae_name", "baked VAE") + if vae_name and vae_name != "baked VAE": + try: + folder_paths.get_full_path_or_raise("vae", vae_name) + except: + errors.append(f"VAE '{vae_name}' not found in vae directory") + + # Validate CLIP models if specified + clip1 = kwargs.get("clip_name1", "none") + if clip1 and clip1 != "none": + try: + folder_paths.get_full_path_or_raise("text_encoders", clip1) + except: + errors.append(f"CLIP model '{clip1}' not found in text_encoders directory") + + clip2 = kwargs.get("clip_name2", "none") + if clip2 and clip2 != "none": + try: + folder_paths.get_full_path_or_raise("text_encoders", clip2) + except: + errors.append(f"CLIP model '{clip2}' not found in text_encoders directory") + + return "; ".join(errors) if errors else True + + def load_diffusion_model( + self, + model_name: str, + weight_dtype: str = "default", + vae_name: str = "baked VAE", + clip_name1: str = "none", + clip_name2: str = "none" + ) -> Tuple[Any, Any, Any, str]: + """ + Main loading function for diffusion models + """ + try: + # Handle "none" and "baked VAE" values + vae_name = vae_name if vae_name and vae_name != "baked VAE" else None + clip_name1 = clip_name1 if clip_name1 and clip_name1 != "none" else None + clip_name2 = clip_name2 if clip_name2 and clip_name2 != "none" else None + weight_dtype = weight_dtype if weight_dtype else "default" + + # Set up model options based on weight_dtype + model_options = {} + if weight_dtype == "fp8_e4m3fn": + model_options["dtype"] = torch.float8_e4m3fn + elif weight_dtype == "fp8_e4m3fn_fast": + model_options["dtype"] = torch.float8_e4m3fn + model_options["fp8_optimizations"] = True + elif weight_dtype == "fp8_e5m2": + model_options["dtype"] = torch.float8_e5m2 + + # Load diffusion model + print(f"[DiffusionModelLoader] Loading diffusion model: {model_name}") + model_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name) + model = comfy.sd.load_diffusion_model(model_path, model_options=model_options) + + if model is None: + raise RuntimeError("Failed to load diffusion model - output was None") + + # Load VAE + vae = None + if vae_name: + print(f"[DiffusionModelLoader] Loading VAE: {vae_name}") + vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) + vae_sd = comfy.utils.load_torch_file(vae_path) + vae = comfy.sd.VAE(sd=vae_sd) + + # Load CLIP models + clip_paths = [] + if clip_name1: + clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name1)) + if clip_name2: + clip_paths.append(folder_paths.get_full_path_or_raise("text_encoders", clip_name2)) + + clip = None + if clip_paths: + print(f"[DiffusionModelLoader] Loading CLIP models") + clip = comfy.sd.load_clip( + ckpt_paths=clip_paths, + embedding_directory=folder_paths.get_folder_paths("embeddings"), + clip_type=comfy.sd.CLIPType.FLUX + ) + + # Generate simple model name string + model_string = os.path.splitext(model_name)[0] + + print(f"[DiffusionModelLoader] Successfully loaded diffusion model: {model_string}") + return (model, vae, clip, model_string) + + except Exception as e: + error_msg = f"Failed to load diffusion model '{model_name}': {str(e)}" + print(f"[DiffusionModelLoader] Error: {error_msg}") + raise RuntimeError(error_msg) \ No newline at end of file diff --git a/nodes/flux_image_saver.py b/nodes/flux_image_saver.py new file mode 100644 index 0000000..9f0c09c --- /dev/null +++ b/nodes/flux_image_saver.py @@ -0,0 +1,203 @@ +""" +FluxImageSaver Node +File: nodes/flux_image_saver.py +""" + +import os +import json +import re +import numpy as np +import folder_paths +from datetime import datetime +from typing import Dict, Any + + +class FluxImageSaver: + """ + FLUX-optimized image saver with organized folder structure and metadata support + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "model_string": ("STRING", {"forceInput": True}), + "processed_prompt": ("STRING", {"forceInput": True}), # Connect from wildcard encoder + "prefix": ("STRING", {"default": "flux_image"}), + "format": (["png", "jpg", "webp"], {"default": "png"}), + "quality": ("INT", {"default": 95, "min": 1, "max": 100}), + "save_metadata": ("BOOLEAN", {"default": True}), + "save_prompt": ("BOOLEAN", {"default": True, "tooltip": "Save the processed prompt as a separate text file"}), + }, + "hidden": { + "prompt_hidden": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO" + } + } + + RETURN_TYPES = () + RETURN_NAMES = () + FUNCTION = "save_images" + CATEGORY = "BawkNodes/image" + OUTPUT_NODE = True + DESCRIPTION = "Save FLUX-generated images with organized folder structure" + + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + + def save_images(self, images, model_string, processed_prompt, prefix="flux_image", + format="png", quality=95, save_metadata=True, save_prompt=True, + prompt_hidden=None, extra_pnginfo=None): + """ + Save images with FLUX-optimized organization and metadata + """ + try: + from PIL import Image + + print(f"[FluxImageSaver] Processing {len(images)} images") + + # Get current date for folder and filename + now = datetime.now() + date_str = now.strftime("%d-%m-%Y") + datetime_str = now.strftime("%d-%m-%Y_%H-%M-%S") + + # Clean model string for folder name + clean_model = self._clean_filename(model_string) + + # Create folder structure: [MODEL]-DD-MM-YYYY + folder_name = f"[{clean_model}]-{date_str}" + output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + os.makedirs(output_dir, exist_ok=True) + + saved_paths = [] + results = [] # For UI image preview + + # Save the processed prompt as a text file if enabled + if save_prompt and processed_prompt: + self._save_prompt_file(output_dir, processed_prompt, datetime_str, prefix) + + for i, image_tensor in enumerate(images): + print(f"[FluxImageSaver] Processing image {i+1}/{len(images)}, shape: {image_tensor.shape}") + + # Convert tensor to PIL Image + image_np = (255.0 * image_tensor.cpu().numpy()).astype(np.uint8) + pil_image = Image.fromarray(image_np) + + # Generate filename: prefix_DD-MM-YYYY_HH-MM-SS_001.ext + batch_suffix = f"_{i+1:03d}" if len(images) > 1 else "" + filename = f"{prefix}_{datetime_str}{batch_suffix}.{format}" + filepath = os.path.join(output_dir, filename) + + # Check if file exists and add counter if necessary + counter = 1 + original_filepath = filepath + while os.path.exists(filepath): + name_part = f"{prefix}_{datetime_str}{batch_suffix}_{counter:03d}" + filename = f"{name_part}.{format}" + filepath = os.path.join(output_dir, filename) + counter += 1 + + # Save image with format-specific options + save_kwargs = {} + if format.lower() == "jpg" or format.lower() == "jpeg": + save_kwargs = {"quality": quality, "optimize": True} + elif format.lower() == "webp": + save_kwargs = {"quality": quality, "method": 6} + elif format.lower() == "png": + save_kwargs = {"optimize": True} + # Add metadata to PNG + if save_metadata and extra_pnginfo: + from PIL.PngImagePlugin import PngInfo + metadata = PngInfo() + for key, value in extra_pnginfo.items(): + metadata.add_text(key, str(value)) + save_kwargs["pnginfo"] = metadata + + pil_image.save(filepath, format.upper(), **save_kwargs) + saved_paths.append(filepath) + + # Add to results for UI preview (relative to output directory) + results.append({ + "filename": filename, + "subfolder": folder_name, + "type": self.type + }) + + print(f"[FluxImageSaver] Saved: {filepath}") + + # Save metadata/prompt file if requested + if save_metadata and prompt_hidden: + self._save_metadata(filepath, prompt_hidden, extra_pnginfo, + model_string, format, quality, processed_prompt) + + print(f"[FluxImageSaver] Returning {len(results)} images for preview") + + return {"ui": {"images": results}} + + except Exception as e: + error_msg = f"Failed to save images: {str(e)}" + print(f"[FluxImageSaver] Error: {error_msg}") + import traceback + traceback.print_exc() + return {"ui": {"images": []}} + + def _clean_filename(self, filename): + """Clean filename for use in folder/file names""" + # Remove file extensions and clean up + clean = re.sub(r'\.[^.]*$', '', filename) # Remove extension + clean = re.sub(r'[<>:"/\\|?*]', '_', clean) # Replace invalid chars + clean = re.sub(r'_{2,}', '_', clean) # Collapse multiple underscores + return clean.strip('_') + + def _save_prompt_file(self, output_dir, processed_prompt, datetime_str, prefix): + """Save the processed prompt as a text file""" + try: + # Create prompt filename + prompt_filename = f"{prefix}_{datetime_str}_prompt.txt" + prompt_filepath = os.path.join(output_dir, prompt_filename) + + # Check if file exists and add counter if necessary + counter = 1 + while os.path.exists(prompt_filepath): + prompt_filename = f"{prefix}_{datetime_str}_prompt_{counter:03d}.txt" + prompt_filepath = os.path.join(output_dir, prompt_filename) + counter += 1 + + # Save the prompt + with open(prompt_filepath, 'w', encoding='utf-8') as f: + f.write(processed_prompt) + + print(f"[FluxImageSaver] Saved prompt: {prompt_filepath}") + + except Exception as e: + print(f"[FluxImageSaver] Failed to save prompt file: {str(e)}") + + def _save_metadata(self, image_path, prompt_hidden, extra_pnginfo, + model_string, format, quality, processed_prompt=None): + """Save metadata as JSON file""" + try: + # Create metadata filename + base_path = os.path.splitext(image_path)[0] + metadata_path = f"{base_path}_metadata.json" + + # Collect metadata + metadata = { + "timestamp": datetime.now().isoformat(), + "model": model_string, + "processed_prompt": processed_prompt, # Add the processed prompt + "format": format, + "quality": quality if format != "png" else None, + "workflow_prompt": prompt_hidden, + "extra_info": extra_pnginfo + } + + # Save as JSON + with open(metadata_path, 'w', encoding='utf-8') as f: + json.dump(metadata, f, indent=2, ensure_ascii=False) + + print(f"[FluxImageSaver] Metadata saved: {metadata_path}") + + except Exception as e: + print(f"[FluxImageSaver] Failed to save metadata: {str(e)}") \ No newline at end of file diff --git a/nodes/flux_prompt_saver.py b/nodes/flux_prompt_saver.py new file mode 100644 index 0000000..dbe8fd3 --- /dev/null +++ b/nodes/flux_prompt_saver.py @@ -0,0 +1,91 @@ +""" +FluxPromptSaver Node +File: nodes/flux_prompt_saver.py +""" + +import os +import json +import re +import folder_paths +from datetime import datetime +from typing import Tuple + + +class FluxPromptSaver: + """ + Standalone prompt saver for FLUX workflows + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_string": ("STRING", {"forceInput": True}), + "prompt": ("STRING", {"multiline": True}), + "filename_prefix": ("STRING", {"default": "flux_prompt"}), + }, + "optional": { + "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "generation_params": ("STRING", {"multiline": True, "default": ""}), + }, + "hidden": { + "extra_pnginfo": "EXTRA_PNGINFO" + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("saved_path",) + FUNCTION = "save_prompt" + CATEGORY = "BawkNodes/text" + OUTPUT_NODE = True + DESCRIPTION = "Save FLUX prompts and generation parameters" + + def save_prompt(self, model_string, prompt, filename_prefix="flux_prompt", + negative_prompt="", generation_params="", extra_pnginfo=None): + """Save prompt and parameters to organized text files""" + try: + # Get current date for folder structure + now = datetime.now() + date_str = now.strftime("%d-%m-%Y") + datetime_str = now.strftime("%d-%m-%Y_%H-%M-%S") + + # Clean model string for folder name + clean_model = self._clean_filename(model_string) + + # Create folder structure matching image saver + folder_name = f"[{clean_model}]-{date_str}" + output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + os.makedirs(output_dir, exist_ok=True) + + # Create filename + filename = f"{filename_prefix}_{datetime_str}.json" + filepath = os.path.join(output_dir, filename) + + # Prepare data + prompt_data = { + "timestamp": now.isoformat(), + "model": model_string, + "positive_prompt": prompt, + "negative_prompt": negative_prompt, + "generation_parameters": generation_params, + "workflow_info": extra_pnginfo + } + + # Save as JSON + with open(filepath, 'w', encoding='utf-8') as f: + json.dump(prompt_data, f, indent=2, ensure_ascii=False) + + print(f"[FluxPromptSaver] Saved: {filepath}") + return (f"Prompt saved to {filename}",) + + except Exception as e: + error_msg = f"Failed to save prompt: {str(e)}" + print(f"[FluxPromptSaver] Error: {error_msg}") + return (f"Error: {error_msg}",) + + def _clean_filename(self, filename): + """Clean filename for use in folder/file names""" + clean = re.sub(r'\.[^.]*$', '', filename) + clean = re.sub(r'[<>:"/\\|?*]', '_', clean) + clean = re.sub(r'_{2,}', '_', clean) + return clean.strip('_') \ No newline at end of file diff --git a/nodes/flux_wildcard_encode.py b/nodes/flux_wildcard_encode.py new file mode 100644 index 0000000..efbf0e2 --- /dev/null +++ b/nodes/flux_wildcard_encode.py @@ -0,0 +1,240 @@ +""" +FluxWildcardEncode Node with LoRA Support (using rgthree-style dynamic UI) +File: nodes/flux_wildcard_encode.py +""" + +import re +import random +import torch +import folder_paths +from typing import Tuple, Any + + +# Fallback implementation for flexible inputs +class FlexibleOptionalInputType(dict): + def __init__(self, type_func, data=None): + super().__init__() + self.data = data or {} + self.type = type_func + for key, value in self.data.items(): + self[key] = value + + def __contains__(self, key): + return True + + def __getitem__(self, key): + if key in self.data: + return self.data[key] + return self.type + + +def any_type(s=None): + return ("*",) + + +class FluxWildcardEncode: + """ + Enhanced FLUX wildcard text encoder with integrated LoRA loading. + Uses rgthree-style dynamic UI for clean LoRA management. + """ + + @classmethod + def INPUT_TYPES(cls): + # Get available LoRA files + lora_list = folder_paths.get_filename_list("loras") + lora_options = ["None"] + lora_list + + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "prompt": ("STRING", {"multiline": True, "default": ""}), + "wildcard_seed": ("INT", { + "default": -1, "min": -1, "max": 0xffffffffffffffff, + "tooltip": "Seed for wildcard selection. Use -1 to disable wildcards" + }), + }, + "optional": { + # LoRA 1 + "lora_1_on": ("BOOLEAN", {"default": False}), + "lora_1_name": (lora_options, {"default": "None"}), + "lora_1_strength": ("FLOAT", {"default": 1.00, "min": -10.0, "max": 10.0, "step": 0.01}), + + # LoRA 2 + "lora_2_on": ("BOOLEAN", {"default": False}), + "lora_2_name": (lora_options, {"default": "None"}), + "lora_2_strength": ("FLOAT", {"default": 1.00, "min": -10.0, "max": 10.0, "step": 0.01}), + + # LoRA 3 + "lora_3_on": ("BOOLEAN", {"default": False}), + "lora_3_name": (lora_options, {"default": "None"}), + "lora_3_strength": ("FLOAT", {"default": 1.00, "min": -10.0, "max": 10.0, "step": 0.01}), + + # LoRA 4 + "lora_4_on": ("BOOLEAN", {"default": False}), + "lora_4_name": (lora_options, {"default": "None"}), + "lora_4_strength": ("FLOAT", {"default": 1.00, "min": -10.0, "max": 10.0, "step": 0.01}), + + # LoRA 5 + "lora_5_on": ("BOOLEAN", {"default": False}), + "lora_5_name": (lora_options, {"default": "None"}), + "lora_5_strength": ("FLOAT", {"default": 1.00, "min": -10.0, "max": 10.0, "step": 0.01}), + + # LoRA 6 + "lora_6_on": ("BOOLEAN", {"default": False}), + "lora_6_name": (lora_options, {"default": "None"}), + "lora_6_strength": ("FLOAT", {"default": 1.00, "min": -10.0, "max": 10.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "CONDITIONING", "STRING") + RETURN_NAMES = ("MODEL", "CLIP", "CONDITIONING", "PROMPT_OUT") + FUNCTION = "encode_with_loras" + CATEGORY = "BawkNodes/conditioning" + DESCRIPTION = "🎲 FLUX Wildcard Encoder with Dynamic LoRA Support" + + def encode_with_loras(self, model, clip, prompt, wildcard_seed=-1, **kwargs): + """ + Encode text prompt with wildcard processing and dynamic LoRA loading + """ + try: + print(f"[FluxWildcardEncode] Starting encode with dynamic LoRAs") + print(f"[FluxWildcardEncode] Received kwargs: {list(kwargs.keys())}") + + # Step 1: Process wildcards if seed is provided + processed_prompt = prompt + if wildcard_seed != -1: + processed_prompt = self._process_wildcards(prompt, wildcard_seed) + if processed_prompt != prompt: + print(f"[FluxWildcardEncode] Processed wildcards in prompt") + + # Step 2: Apply LoRAs from fixed slots + working_model = model + working_clip = clip + ui_lora_count = 0 + + # Process 6 fixed LoRA slots + for i in range(1, 7): # LoRA 1-6 + on_key = f"lora_{i}_on" + name_key = f"lora_{i}_name" + strength_key = f"lora_{i}_strength" + + # Check if this LoRA slot is enabled and has a valid selection + lora_enabled = kwargs.get(on_key, False) + lora_name = kwargs.get(name_key, "None") + + if lora_enabled and lora_name and lora_name != "None": + strength = kwargs.get(strength_key, 1.00) + + # Skip if strength is zero + if strength == 0: + continue + + try: + from nodes import LoraLoader + working_model, working_clip = LoraLoader().load_lora( + working_model, working_clip, lora_name, strength, strength + ) + ui_lora_count += 1 + print(f"[FluxWildcardEncode] Applied LoRA {i}: {lora_name} (strength: {strength})") + except Exception as e: + print(f"[FluxWildcardEncode] Failed to load LoRA {i} ({lora_name}): {e}") + + if ui_lora_count > 0: + print(f"[FluxWildcardEncode] Applied {ui_lora_count} LoRAs") + else: + print(f"[FluxWildcardEncode] No LoRAs applied") + + # Step 3: Encode the processed prompt + if not processed_prompt.strip(): + print("[FluxWildcardEncode] Warning: Empty prompt") + empty_conditioning = [[torch.zeros((1, 77, 768)), {"pooled_output": torch.zeros((1, 768))}]] + return (working_model, working_clip, empty_conditioning, processed_prompt) + + # Encode using the LoRA-modified CLIP + tokens = working_clip.tokenize(processed_prompt) + conditioning, pooled = working_clip.encode_from_tokens(tokens, return_pooled=True) + conditioning = [[conditioning, {"pooled_output": pooled}]] + + print(f"[FluxWildcardEncode] Success: {ui_lora_count} LoRAs, conditioning shape: {conditioning[0][0].shape}") + return (working_model, working_clip, conditioning, processed_prompt) + + except Exception as e: + error_msg = f"Failed to encode with LoRAs: {str(e)}" + print(f"[FluxWildcardEncode] Error: {error_msg}") + import traceback + traceback.print_exc() + empty_conditioning = [[torch.zeros((1, 77, 768)), {"pooled_output": torch.zeros((1, 768))}]] + return (model, clip, empty_conditioning, prompt) + + def _process_wildcards(self, prompt, seed): + """ + Process wildcard syntax in prompts + Wildcard format: {option1|option2|option3} + """ + # Set seed for consistent wildcard selection + random.seed(seed) + + processed = prompt + + # Find and process wildcards in {option1|option2|option3} format + wildcard_pattern = r'\{([^}]+)\}' + + def replace_wildcard(match): + wildcard_content = match.group(1) + if '|' in wildcard_content: + options = [option.strip() for option in wildcard_content.split('|')] + if options: # Make sure we have options + return random.choice(options) + return match.group(0) # Return original if no valid options + + # Replace all wildcards + processed = re.sub(wildcard_pattern, replace_wildcard, processed) + + return processed + + def _get_lora_by_filename(self, file_path): + """Find LoRA file with fuzzy matching (rgthree style)""" + if not file_path: + return None + + lora_paths = folder_paths.get_filename_list('loras') + + # Direct match + if file_path in lora_paths: + return file_path + + # Match without extension + lora_paths_no_ext = [os.path.splitext(x)[0] for x in lora_paths] + if file_path in lora_paths_no_ext: + return lora_paths[lora_paths_no_ext.index(file_path)] + + # Force input without extension + file_path_no_ext = os.path.splitext(file_path)[0] + if file_path_no_ext in lora_paths_no_ext: + return lora_paths[lora_paths_no_ext.index(file_path_no_ext)] + + # Basename matching + lora_basenames = [os.path.basename(x) for x in lora_paths] + if file_path in lora_basenames: + return lora_paths[lora_basenames.index(file_path)] + + # Basename without extension + file_basename = os.path.basename(file_path) + if file_basename in lora_basenames: + return lora_paths[lora_basenames.index(file_basename)] + + # Basename no extension matching + lora_basenames_no_ext = [os.path.splitext(os.path.basename(x))[0] for x in lora_paths] + file_basename_no_ext = os.path.splitext(os.path.basename(file_path))[0] + if file_basename_no_ext in lora_basenames_no_ext: + return lora_paths[lora_basenames_no_ext.index(file_basename_no_ext)] + + # Fuzzy partial match + for lora_path in lora_paths: + if file_path in lora_path: + print(f"[FluxWildcardEncode] Fuzzy-matched '{file_path}' to '{lora_path}'") + return lora_path + + print(f"[FluxWildcardEncode] LoRA '{file_path}' not found") + return None \ No newline at end of file diff --git a/previews/efficient-workflow-preview.png b/previews/efficient-workflow-preview.png new file mode 100644 index 0000000..addf56b Binary files /dev/null and b/previews/efficient-workflow-preview.png differ diff --git a/pyproject.toml b/pyproject.toml index 18a9c27..f0d653c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,12 @@ [project] name = "comfyui-bawknodes" -version = "1.0.0" -description = "A collection of useful ComfyUI nodes including advanced diffusion model loader and more" +version = "2.0.0" +description = "A complete collection of FLUX-optimized ComfyUI nodes for enhanced AI image generation workflows." readme = "README.md" license = { file = "LICENSE" } requires-python = ">=3.8" authors = [ - { name = "Your Name", email = "your.email@example.com" } + { name = "judd", email = "" } ] keywords = ["comfyui", "diffusion", "model-loading", "ai", "flux", "nodes", "bawk"] classifiers = [ diff --git a/utils/flexible_inputs.py b/utils/flexible_inputs.py new file mode 100644 index 0000000..33f0de1 --- /dev/null +++ b/utils/flexible_inputs.py @@ -0,0 +1,60 @@ +""" +Flexible Input Utilities for Dynamic UI (based on rgthree) +File: utils/flexible_inputs.py +""" + +from typing import Union + + +class AnyType(str): + """A special class that is always equal in not equal comparisons. Credit to pythongosssss""" + + def __ne__(self, __value: object) -> bool: + return False + + +class FlexibleOptionalInputType(dict): + """A special class to make flexible nodes that pass data to our python handlers. + + Enables both flexible/dynamic input types (like for Any Switch) or a dynamic number of inputs + (like for Any Switch, Context Switch, Context Merge, Power Lora Loader, etc). + + Initially, ComfyUI only needed to return True for `__contains__` below, which told ComfyUI that + our node will handle the input, regardless of what it is. + + However, after https://github.com/comfyanonymous/ComfyUI/pull/2666 ComfyUI's execution changed + also checking the data for the key; specifically, the type which is the first tuple entry. This + type is supplied to our FlexibleOptionalInputType and returned for any non-data key. This can be a + real type, or use the AnyType for additional flexibility. + """ + + def __init__(self, type_func, data: Union[dict, None] = None): + """Initializes the FlexibleOptionalInputType. + + Args: + type_func: The flexible type to use when ComfyUI retrieves an unknown key (via `__getitem__`). + data: An optional dict to use as the basis. This is stored both in a `data` attribute, so we + can look it up without hitting our overrides, as well as iterated over and adding its key + and values to our `self` keys. This way, when looked at, we will appear to represent this + data. When used in an "optional" INPUT_TYPES, these are the starting optional node types. + """ + super().__init__() + self.data = data or {} + self.type = type_func + for key, value in self.data.items(): + self[key] = value + + def __contains__(self, key): + """Return True for any key, allowing ComfyUI to pass any input to our node.""" + return True + + def __getitem__(self, key): + """Return the type for any unknown key, or the stored value for known keys.""" + if key in self.data: + return self.data[key] + return self.type + + +def any_type(s=None): + """Helper function for any type inputs""" + return AnyType("*") \ No newline at end of file diff --git a/web/extensions/__init__.py b/web/extensions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/web/extensions/flux_wildcard_encoder.js b/web/extensions/flux_wildcard_encoder.js new file mode 100644 index 0000000..6cb96aa --- /dev/null +++ b/web/extensions/flux_wildcard_encoder.js @@ -0,0 +1,377 @@ +/** + * FLUX Wildcard Encoder Extension - rgthree-style Power LoRA Loader UI + * File: web/extensions/flux_wildcard_encoder.js + */ + +import { app } from "../../scripts/app.js"; +import { ComfyWidgets } from "../../scripts/widgets.js"; + +const NODE_TYPE = "FluxWildcardEncode"; +const MAX_LORAS = 10; // Maximum number of LoRAs + +class FluxWildcardEncoderNode { + constructor() { + this.loras = []; + this.loraCounter = 0; + } + + setup(node) { + this.node = node; + this.setupLoraUI(); + } + + setupLoraUI() { + // Remove default widgets that we'll replace + this.node.widgets = this.node.widgets?.filter(w => + !w.name.startsWith("lora_") && + w.name !== "toggle_all_loras" + ) || []; + + // Add "Toggle All" button at the top + const toggleAllWidget = ComfyWidgets.BOOLEAN(this.node, "toggle_all", ["BOOLEAN", { default: false }], app); + toggleAllWidget.name = "toggle_all"; + toggleAllWidget.callback = () => { + const enabled = toggleAllWidget.value; + this.loras.forEach(lora => { + if (lora.toggle) { + lora.toggle.value = enabled; + } + }); + this.node.setDirtyCanvas(true, true); + }; + + // Add initial LoRA row + this.addLoraRow(); + + // Add "Add LoRA" button + this.addLoraButton(); + + // Update node size + this.updateNodeSize(); + } + + addLoraRow() { + if (this.loras.length >= MAX_LORAS) { + console.warn(`Maximum ${MAX_LORAS} LoRAs reached`); + return; + } + + const loraIndex = this.loraCounter++; + const lora = { + index: loraIndex, + enabled: false, + name: "None", + strength: 1.0, + strengthTwo: 1.0 + }; + + // Create toggle widget + const toggleWidget = ComfyWidgets.BOOLEAN( + this.node, + `lora_${loraIndex}_toggle`, + ["BOOLEAN", { default: false }], + app + ); + toggleWidget.name = `lora_${loraIndex}_on`; + toggleWidget.callback = () => { + lora.enabled = toggleWidget.value; + this.updateNodeInputs(); + }; + + // Create LoRA selection widget + const loraOptions = this.getLoraOptions(); + const loraWidget = ComfyWidgets.COMBO( + this.node, + `lora_${loraIndex}_name`, + [loraOptions, { default: "None" }], + app + ); + loraWidget.name = `lora_${loraIndex}_lora`; + loraWidget.callback = () => { + lora.name = loraWidget.value; + this.updateNodeInputs(); + }; + + // Create strength widget + const strengthWidget = ComfyWidgets.NUMBER( + this.node, + `lora_${loraIndex}_strength`, + ["FLOAT", { default: 1.0, min: -10.0, max: 10.0, step: 0.01 }], + app + ); + strengthWidget.name = `lora_${loraIndex}_strength`; + strengthWidget.callback = () => { + lora.strength = strengthWidget.value; + lora.strengthTwo = strengthWidget.value; // Use same for both by default + this.updateNodeInputs(); + }; + + // Create remove button for this LoRA + const removeWidget = ComfyWidgets.BUTTON( + this.node, + `remove_lora_${loraIndex}`, + "🗑️", + () => this.removeLoraRow(loraIndex) + ); + removeWidget.name = `remove_lora_${loraIndex}`; + + // Store references + lora.toggle = toggleWidget; + lora.loraWidget = loraWidget; + lora.strengthWidget = strengthWidget; + lora.removeWidget = removeWidget; + + this.loras.push(lora); + this.updateNodeSize(); + this.updateNodeInputs(); + } + + removeLoraRow(index) { + // Find and remove the LoRA + const loraIndex = this.loras.findIndex(l => l.index === index); + if (loraIndex === -1) return; + + const lora = this.loras[loraIndex]; + + // Remove widgets + this.node.widgets = this.node.widgets.filter(w => + w !== lora.toggle && + w !== lora.loraWidget && + w !== lora.strengthWidget && + w !== lora.removeWidget + ); + + // Remove from array + this.loras.splice(loraIndex, 1); + + this.updateNodeSize(); + this.updateNodeInputs(); + this.node.setDirtyCanvas(true, true); + } + + addLoraButton() { + const addButton = ComfyWidgets.BUTTON( + this.node, + "add_lora", + "+ Add LoRA", + () => this.addLoraRow() + ); + addButton.name = "add_lora_button"; + } + + getLoraOptions() { + // Get LoRA files from ComfyUI + try { + const loraList = app.ui.settings.getSettingValue("Comfy.LoraList") || []; + return ["None", ...loraList]; + } catch (e) { + // Fallback - we'll populate this from the backend + return ["None"]; + } + } + + updateNodeInputs() { + // Clear existing LoRA inputs + if (this.node.inputs) { + this.node.inputs = this.node.inputs.filter(input => + !input.name.startsWith("lora_") + ); + } + + // Add current LoRA inputs to the node's input data + this.loras.forEach(lora => { + if (lora.enabled && lora.name !== "None") { + const loraData = { + on: lora.enabled, + lora: lora.name, + strength: lora.strength, + strengthTwo: lora.strengthTwo + }; + + // Store in node's widget values for serialization + if (!this.node.widgets_values) { + this.node.widgets_values = {}; + } + this.node.widgets_values[`lora_${lora.index}`] = loraData; + } + }); + } + + updateNodeSize() { + // Calculate required height based on number of widgets + const baseHeight = 100; // Base height for prompt and wildcard_seed + const loraRowHeight = 25; // Height per LoRA row + const buttonHeight = 25; // Height for buttons + + const totalHeight = baseHeight + + (this.loras.length * loraRowHeight) + + buttonHeight + + 25; // Toggle all button + + this.node.size = [ + Math.max(300, this.node.size[0]), // Minimum width + Math.max(totalHeight, this.node.size[1]) + ]; + + this.node.setDirtyCanvas(true, true); + } + + onDrawForeground(ctx) { + // Custom drawing for LoRA UI (similar to rgthree style) + if (!this.loras.length) return; + + const margin = 10; + let y = 80; // Start below the main inputs + + // Draw LoRA section header + ctx.fillStyle = "#555"; + ctx.fillRect(margin, y, this.node.size[0] - margin * 2, 1); + y += 10; + + // Draw each LoRA row + this.loras.forEach((lora, index) => { + this.drawLoraRow(ctx, lora, margin, y + index * 25); + }); + } + + drawLoraRow(ctx, lora, x, y) { + const rowHeight = 20; + const toggleWidth = 20; + const nameWidth = 150; + const strengthWidth = 60; + + // Draw toggle + this.drawToggle(ctx, x, y, toggleWidth, rowHeight, lora.enabled); + + // Draw LoRA name + ctx.fillStyle = lora.name !== "None" ? "#fff" : "#888"; + ctx.font = "12px Arial"; + ctx.fillText( + lora.name.length > 20 ? lora.name.substring(0, 17) + "..." : lora.name, + x + toggleWidth + 5, + y + 14 + ); + + // Draw strength value + ctx.fillStyle = "#fff"; + ctx.fillText( + lora.strength.toFixed(2), + x + toggleWidth + nameWidth + 10, + y + 14 + ); + } + + drawToggle(ctx, x, y, width, height, enabled) { + // Draw toggle button background + ctx.fillStyle = enabled ? "#4a9eff" : "#333"; + ctx.fillRect(x + 2, y + 2, width - 4, height - 4); + + // Draw toggle border + ctx.strokeStyle = "#666"; + ctx.lineWidth = 1; + ctx.strokeRect(x + 2, y + 2, width - 4, height - 4); + + // Draw checkmark if enabled + if (enabled) { + ctx.strokeStyle = "#fff"; + ctx.lineWidth = 2; + ctx.beginPath(); + ctx.moveTo(x + 6, y + height / 2); + ctx.lineTo(x + width / 2, y + height - 6); + ctx.lineTo(x + width - 4, y + 6); + ctx.stroke(); + } + } + + serialize() { + const data = {}; + this.loras.forEach(lora => { + if (lora.enabled && lora.name !== "None") { + data[`lora_${lora.index}`] = { + on: lora.enabled, + lora: lora.name, + strength: lora.strength, + strengthTwo: lora.strengthTwo + }; + } + }); + return data; + } + + deserialize(data) { + // Clear existing LoRAs + this.loras = []; + this.loraCounter = 0; + + // Recreate LoRAs from data + Object.keys(data).forEach(key => { + if (key.startsWith("lora_")) { + const loraData = data[key]; + this.addLoraRow(); + const lora = this.loras[this.loras.length - 1]; + + lora.enabled = loraData.on || false; + lora.name = loraData.lora || "None"; + lora.strength = loraData.strength || 1.0; + lora.strengthTwo = loraData.strengthTwo || lora.strength; + + // Update widgets + if (lora.toggle) lora.toggle.value = lora.enabled; + if (lora.loraWidget) lora.loraWidget.value = lora.name; + if (lora.strengthWidget) lora.strengthWidget.value = lora.strength; + } + }); + + this.updateNodeSize(); + this.updateNodeInputs(); + } +} + +// Extension registration +app.registerExtension({ + name: "BawkNodes.FluxWildcardEncoder", + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === NODE_TYPE) { + const onNodeCreated = nodeType.prototype.onNodeCreated; + + nodeType.prototype.onNodeCreated = function() { + const result = onNodeCreated?.apply(this, arguments); + + // Initialize our custom UI + this.fluxWildcardEncoder = new FluxWildcardEncoderNode(); + this.fluxWildcardEncoder.setup(this); + + // Override drawing + const originalOnDrawForeground = this.onDrawForeground; + this.onDrawForeground = function(ctx) { + originalOnDrawForeground?.apply(this, arguments); + this.fluxWildcardEncoder?.onDrawForeground(ctx); + }; + + return result; + }; + + // Override serialization + const originalSerialize = nodeType.prototype.serialize; + nodeType.prototype.serialize = function() { + const data = originalSerialize?.apply(this, arguments) || {}; + if (this.fluxWildcardEncoder) { + data.fluxWildcardEncoder = this.fluxWildcardEncoder.serialize(); + } + return data; + }; + + // Override deserialization + const originalConfigure = nodeType.prototype.configure; + nodeType.prototype.configure = function(data) { + originalConfigure?.apply(this, arguments); + if (this.fluxWildcardEncoder && data.fluxWildcardEncoder) { + this.fluxWildcardEncoder.deserialize(data.fluxWildcardEncoder); + } + }; + } + } +}); + +console.log("🎲 FLUX Wildcard Encoder extension loaded"); \ No newline at end of file diff --git a/workflows/example_bawknodes_workflow.json b/workflows/example_bawknodes_workflow.json new file mode 100644 index 0000000..bb66641 --- /dev/null +++ b/workflows/example_bawknodes_workflow.json @@ -0,0 +1,343 @@ +{ + "id": "5acb6057-af87-40fa-90de-4b9e3cdcd5e3", + "revision": 0, + "last_node_id": 35, + "last_link_id": 78, + "nodes": [ + { + "id": 30, + "type": "DiffusionModelLoader", + "pos": [ + -259.1785888671875, + 87.82862091064453 + ], + "size": [ + 342.2134704589844, + 214 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 74 + ] + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 77 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 75 + ] + }, + { + "name": "MODEL_STRING", + "type": "STRING", + "links": [ + 66 + ] + } + ], + "properties": { + "cnr_id": "comfyui-bawknodes", + "ver": "97bbbef153df94caebb97f054b8fd53de1861cf1", + "Node name for S&R": "DiffusionModelLoader" + }, + "widgets_values": [ + "flux1-dev.safetensors", + "ae.safetensors", + "t5xxl_fp16.safetensors", + "clip_l.safetensors", + "fp8_e4m3fn" + ], + "shape": 1 + }, + { + "id": 34, + "type": "BawkSampler", + "pos": [ + 788.8836059570312, + 87.903076171875 + ], + "size": [ + 288.7222595214844, + 622 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 72 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 73 + }, + { + "name": "vae", + "type": "VAE", + "link": 77 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 76 + ] + }, + { + "name": "latent", + "type": "LATENT", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-bawknodes", + "ver": "97bbbef153df94caebb97f054b8fd53de1861cf1", + "Node name for S&R": "BawkSampler" + }, + "widgets_values": [ + "FHD 16:9 - 1920x1080", + 4, + 417938486318098, + "randomize", + "euler", + "beta", + 30, + 3.5, + 0.5, + 0.3, + 1, + false, + 1920, + 1080 + ], + "shape": 1 + }, + { + "id": 31, + "type": "FluxImageSaver", + "pos": [ + 1109.04248046875, + 87.9061508178711 + ], + "size": [ + 536.477294921875, + 470.7662353515625 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 76 + }, + { + "name": "model_string", + "type": "STRING", + "link": 66 + }, + { + "name": "processed_prompt", + "type": "STRING", + "link": 78 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfyui-bawknodes", + "ver": "97bbbef153df94caebb97f054b8fd53de1861cf1", + "Node name for S&R": "FluxImageSaver" + }, + "widgets_values": [ + "NEW-CKH", + "png", + 100, + false, + true + ], + "shape": 1 + }, + { + "id": 35, + "type": "FluxWildcardEncode", + "pos": [ + 118.85038757324219, + 87.09760284423828 + ], + "size": [ + 633.321533203125, + 695.784912109375 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 74 + }, + { + "name": "clip", + "type": "CLIP", + "link": 75 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 72 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "links": null + }, + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 73 + ] + }, + { + "name": "PROMPT_OUT", + "type": "STRING", + "links": [ + 78 + ] + } + ], + "properties": { + "cnr_id": "comfyui-bawknodes", + "ver": "97bbbef153df94caebb97f054b8fd53de1861cf1", + "Node name for S&R": "FluxWildcardEncode" + }, + "widgets_values": [ + "", + -1, + false, + "None", + 1, + false, + "None", + 1, + false, + "None", + 1, + false, + "None", + 1, + false, + "None", + 1, + false, + "None", + 1 + ], + "shape": 1 + } + ], + "links": [ + [ + 66, + 30, + 3, + 31, + 1, + "STRING" + ], + [ + 72, + 35, + 0, + 34, + 0, + "MODEL" + ], + [ + 73, + 35, + 2, + 34, + 1, + "CONDITIONING" + ], + [ + 74, + 30, + 0, + 35, + 0, + "MODEL" + ], + [ + 75, + 30, + 2, + 35, + 1, + "CLIP" + ], + [ + 76, + 34, + 0, + 31, + 0, + "IMAGE" + ], + [ + 77, + 30, + 1, + 34, + 2, + "VAE" + ], + [ + 78, + 35, + 3, + 31, + 2, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8358208385450411, + "offset": [ + 299.5806636329205, + 171.9927929925298 + ] + }, + "frontendVersion": "1.23.4" + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/example_flux_workflow.json.json b/workflows/example_flux_workflow.json.json deleted file mode 100644 index b25eb71..0000000 --- a/workflows/example_flux_workflow.json.json +++ /dev/null @@ -1,320 +0,0 @@ -{ - "last_node_id": 10, - "last_link_id": 15, - "nodes": [ - { - "id": 1, - "type": "DiffusionModelLoader", - "pos": [50, 50], - "size": [400, 200], - "flags": {}, - "order": 0, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "MODEL", - "type": "MODEL", - "links": [1], - "slot_index": 0 - }, - { - "name": "VAE", - "type": "VAE", - "links": [2], - "slot_index": 1 - }, - { - "name": "CLIP", - "type": "CLIP", - "links": [3], - "slot_index": 2 - }, - { - "name": "MODEL_STRING", - "type": "STRING", - "links": [4], - "slot_index": 3 - } - ], - "properties": { - "Node name for S&R": "DiffusionModelLoader" - }, - "widgets_values": [ - "flux1-dev.safetensors", - "fp8_e4m3fn", - "ae.safetensors", - "t5xxl_fp16.safetensors", - "clip_l.safetensors" - ], - "title": "🚀 FLUX Model Loader" - }, - { - "id": 2, - "type": "CLIPTextEncode", - "pos": [500, 50], - "size": [400, 200], - "flags": {}, - "order": 1, - "mode": 0, - "inputs": [ - { - "name": "clip", - "type": "CLIP", - "link": 3 - } - ], - "outputs": [ - { - "name": "CONDITIONING", - "type": "CONDITIONING", - "links": [5], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "CLIPTextEncode" - }, - "widgets_values": [ - "A beautiful landscape with mountains and a lake, highly detailed, photorealistic" - ], - "title": "Positive Prompt" - }, - { - "id": 3, - "type": "CLIPTextEncode", - "pos": [500, 300], - "size": [400, 200], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [ - { - "name": "clip", - "type": "CLIP", - "link": 3 - } - ], - "outputs": [ - { - "name": "CONDITIONING", - "type": "CONDITIONING", - "links": [6], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "CLIPTextEncode" - }, - "widgets_values": [ - "blurry, low quality, distorted" - ], - "title": "Negative Prompt" - }, - { - "id": 4, - "type": "EmptyLatentImage", - "pos": [50, 300], - "size": [300, 100], - "flags": {}, - "order": 3, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "LATENT", - "type": "LATENT", - "links": [7], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "EmptyLatentImage" - }, - "widgets_values": [ - 1024, - 1024, - 1 - ], - "title": "Empty Latent Image" - }, - { - "id": 5, - "type": "KSampler", - "pos": [950, 50], - "size": [400, 300], - "flags": {}, - "order": 4, - "mode": 0, - "inputs": [ - { - "name": "model", - "type": "MODEL", - "link": 1 - }, - { - "name": "positive", - "type": "CONDITIONING", - "link": 5 - }, - { - "name": "negative", - "type": "CONDITIONING", - "link": 6 - }, - { - "name": "latent_image", - "type": "LATENT", - "link": 7 - } - ], - "outputs": [ - { - "name": "LATENT", - "type": "LATENT", - "links": [8], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "KSampler" - }, - "widgets_values": [ - 12345, - "randomize", - 20, - 1.0, - "euler", - "normal", - 1.0 - ], - "title": "KSampler" - }, - { - "id": 6, - "type": "VAEDecode", - "pos": [1400, 50], - "size": [200, 100], - "flags": {}, - "order": 5, - "mode": 0, - "inputs": [ - { - "name": "samples", - "type": "LATENT", - "link": 8 - }, - { - "name": "vae", - "type": "VAE", - "link": 2 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [9], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "VAEDecode" - }, - "title": "VAE Decode" - }, - { - "id": 7, - "type": "SaveImage", - "pos": [1650, 50], - "size": [300, 300], - "flags": {}, - "order": 6, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 9 - } - ], - "outputs": [], - "properties": { - "Node name for S&R": "SaveImage" - }, - "widgets_values": [ - "ComfyUI_JLoader" - ], - "title": "Save Image" - }, - { - "id": 8, - "type": "ShowText", - "pos": [50, 450], - "size": [400, 100], - "flags": {}, - "order": 7, - "mode": 0, - "inputs": [ - { - "name": "text", - "type": "STRING", - "link": 4 - } - ], - "outputs": [], - "properties": { - "Node name for S&R": "ShowText" - }, - "title": "Model Info Display" - } - ], - "links": [ - [1, 1, 0, 5, 0, "MODEL"], - [2, 1, 1, 6, 1, "VAE"], - [3, 1, 2, 2, 0, "CLIP"], - [4, 1, 3, 8, 0, "STRING"], - [5, 2, 0, 5, 1, "CONDITIONING"], - [6, 3, 0, 5, 2, "CONDITIONING"], - [7, 4, 0, 5, 3, "LATENT"], - [8, 5, 0, 6, 0, "LATENT"], - [9, 6, 0, 7, 0, "IMAGE"] - ], - "groups": [ - { - "title": "Model Loading", - "bounding": [25, 25, 450, 400], - "color": "#3f789e" - }, - { - "title": "Text Encoding", - "bounding": [475, 25, 450, 500], - "color": "#a1309b" - }, - { - "title": "Generation", - "bounding": [925, 25, 450, 350], - "color": "#b58b2a" - }, - { - "title": "Output", - "bounding": [1375, 25, 600, 350], - "color": "#8a8a8a" - } - ], - "config": {}, - "extra": { - "ds": { - "scale": 0.8, - "offset": [0, 0] - } - }, - "version": 0.4, - "workflow_info": { - "name": "FLUX Model Loading Example - Bawk Nodes", - "description": "Example workflow showing how to use the Bawk Nodes Diffusion Model Loader with FLUX models, including separate VAE and dual CLIP text encoders.", - "author": "ComfyUI Bawk Nodes", - "version": "1.0.0", - "tags": ["flux", "model-loading", "advanced", "fp8", "bawk-nodes"] - } -} \ No newline at end of file