Pixel-space diffusion text-to-image generation with REPA LoRA training. 6 nodes: Model Loader, Text Encoder Loader, Generate, Training Config, Train, LoRA Loader. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
"""ComfyUI-FL-PixelGen: Pixel-space diffusion text-to-image generation."""
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
|
|
logger = logging.getLogger("FL_PixelGen")
|
|
|
|
# Add src/ to sys.path so bundled PixelGen code can be imported as `src.models.*`
|
|
_this_dir = os.path.dirname(os.path.abspath(__file__))
|
|
_src_parent = _this_dir # src/ lives directly under this package
|
|
if _src_parent not in sys.path:
|
|
sys.path.insert(0, _src_parent)
|
|
|
|
# Register model folder with ComfyUI
|
|
try:
|
|
import folder_paths
|
|
|
|
models_dir = os.path.join(
|
|
folder_paths.models_dir, "pixelgen"
|
|
)
|
|
os.makedirs(models_dir, exist_ok=True)
|
|
folder_paths.add_model_folder_path("pixelgen", models_dir)
|
|
|
|
# Scan for user-placed checkpoints
|
|
from modules.model_info import AVAILABLE_MODELS
|
|
if os.path.isdir(models_dir):
|
|
for f in os.listdir(models_dir):
|
|
if f.endswith((".ckpt", ".pt", ".safetensors", ".pth")):
|
|
name = os.path.splitext(f)[0]
|
|
if name not in AVAILABLE_MODELS:
|
|
AVAILABLE_MODELS[name] = os.path.join(models_dir, f)
|
|
logger.info(f"Found local PixelGen model: {name}")
|
|
except Exception as e:
|
|
logger.warning(f"Could not register PixelGen model folder: {e}")
|
|
|
|
# Import and register nodes
|
|
try:
|
|
from .nodes import (
|
|
FL_PixelGen_ModelLoader,
|
|
FL_PixelGen_TextEncoderLoader,
|
|
FL_PixelGen_Generate,
|
|
FL_PixelGen_TrainingConfig,
|
|
FL_PixelGen_Train,
|
|
FL_PixelGen_LoraLoader,
|
|
)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FL_PixelGen_ModelLoader": FL_PixelGen_ModelLoader,
|
|
"FL_PixelGen_TextEncoderLoader": FL_PixelGen_TextEncoderLoader,
|
|
"FL_PixelGen_Generate": FL_PixelGen_Generate,
|
|
"FL_PixelGen_TrainingConfig": FL_PixelGen_TrainingConfig,
|
|
"FL_PixelGen_Train": FL_PixelGen_Train,
|
|
"FL_PixelGen_LoraLoader": FL_PixelGen_LoraLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FL_PixelGen_ModelLoader": "FL PixelGen Model Loader",
|
|
"FL_PixelGen_TextEncoderLoader": "FL PixelGen Text Encoder Loader",
|
|
"FL_PixelGen_Generate": "FL PixelGen Generate",
|
|
"FL_PixelGen_TrainingConfig": "FL PixelGen Training Config",
|
|
"FL_PixelGen_Train": "FL PixelGen Train",
|
|
"FL_PixelGen_LoraLoader": "FL PixelGen LoRA Loader",
|
|
}
|
|
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to load FL_PixelGen nodes: {e}")
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|