Eric Rollei
a5a7cc391e
- New nodes/component_loader.py: GenfocusComponentLoader node that builds a diffusers FluxPipeline from individual safetensors files (transformer, T5-XXL, CLIP-L, VAE) instead of requiring the full ~23GB diffusers folder - New pipeline/component_assembly.py: pipeline assembly logic, auto-downloads ~5MB of tokenizer/config files from FLUX.1-schnell to .flux_cache/ on first run - __init__.py: independent import block for component loader with full traceback on failure; registers 'Genfocus Model Loader (Components)' node - nodes/genfocus_loader.py: GenfocusUnloadModels now also clears component loader cache - README.md: document Component Loader node, alternative setup folder layout, compatibility caveats, and expanded troubleshooting entries - pyproject.toml: bump version to 0.1.2 - .gitignore: exclude .flux_cache/ and node.zip
212 lines
7.7 KiB
Python
212 lines
7.7 KiB
Python
"""
|
|
Refocus - ComfyUI Nodes for Generative Refocusing
|
|
|
|
Based on the Genfocus project (https://github.com/rayray9999/Genfocus)
|
|
Provides deblurring and bokeh generation using FLUX LoRA adapters.
|
|
|
|
Main features:
|
|
- Native Genfocus nodes (diffusers-based) - full multi-conditional support
|
|
- DepthPro depth estimation nodes
|
|
- Defocus map computation utilities
|
|
|
|
License: Apache 2.0 (for Genfocus components)
|
|
Note: Apple ml-depth-pro has its own license terms
|
|
Note: FLUX.1-dev requires acceptance of its license
|
|
|
|
=== MODEL FOLDER SETUP ===
|
|
|
|
Genfocus LoRAs:
|
|
Place DeblurNet and BokehNet LoRAs in:
|
|
ComfyUI/models/genfocus/
|
|
|
|
Expected files:
|
|
- deblurNet.safetensors (or similar name)
|
|
- bokehNet.safetensors (or similar name)
|
|
|
|
Download from: https://huggingface.co/nycu-cplab/Genfocus-Model
|
|
|
|
FLUX Diffusers Model:
|
|
Place FLUX.1-dev in diffusers format in:
|
|
ComfyUI/models/diffusers/FLUX.1-dev/
|
|
|
|
Or use the HuggingFace ID (requires login):
|
|
black-forest-labs/FLUX.1-dev
|
|
|
|
The folder should contain model_index.json and model files.
|
|
"""
|
|
|
|
import os
|
|
import traceback
|
|
import folder_paths
|
|
|
|
# =============================================================================
|
|
# REGISTER CUSTOM MODEL FOLDERS
|
|
# =============================================================================
|
|
|
|
# Register genfocus folder for LoRAs
|
|
if "genfocus" not in folder_paths.folder_names_and_paths:
|
|
genfocus_path = os.path.join(folder_paths.models_dir, "genfocus")
|
|
os.makedirs(genfocus_path, exist_ok=True)
|
|
folder_paths.folder_names_and_paths["genfocus"] = ([genfocus_path], {".safetensors", ".pt", ".pth", ".bin"})
|
|
print(f"[Refocus] Genfocus LoRA folder: {genfocus_path}")
|
|
|
|
# Register diffusers folder for FLUX models
|
|
if "diffusers" not in folder_paths.folder_names_and_paths:
|
|
diffusers_path = os.path.join(folder_paths.models_dir, "diffusers")
|
|
os.makedirs(diffusers_path, exist_ok=True)
|
|
folder_paths.folder_names_and_paths["diffusers"] = ([diffusers_path], set())
|
|
print(f"[Refocus] Diffusers model folder: {diffusers_path}")
|
|
|
|
# =============================================================================
|
|
# IMPORT NODES
|
|
# =============================================================================
|
|
|
|
# Depth Pro nodes (always available)
|
|
from .nodes.depth_pro import (
|
|
DepthProModelLoader,
|
|
DepthProEstimate,
|
|
DepthMetricToRelative,
|
|
DepthMetricToInverse,
|
|
FocalPXtoMM,
|
|
FocalMMtoPX,
|
|
)
|
|
# Utility nodes (KSampler-compatible helpers, limited functionality)
|
|
from .nodes.genfocus_lora import GenfocusLoRALoader
|
|
from .nodes.deblur import DeblurNetApply
|
|
from .nodes.defocus_map import ComputeDefocusMap, SelectFocusPoint, FocusPointFromMask
|
|
from .nodes.bokeh import BokehNetApply
|
|
|
|
# Native Genfocus nodes (diffusers-based, full implementation)
|
|
# DIAGNOSTIC: Catch ALL exceptions during native node import so any silent
|
|
# failure surfaces with a full traceback instead of being hidden.
|
|
NATIVE_GENFOCUS_AVAILABLE = False
|
|
COMPONENT_LOADER_AVAILABLE = False
|
|
try:
|
|
from .nodes.genfocus_loader import GenfocusModelLoader, GenfocusSwitchAdapter, GenfocusUnloadModels
|
|
from .nodes.genfocus_generate import (
|
|
GenfocusCondition,
|
|
GenfocusDefocusMapCondition,
|
|
GenfocusGenerate,
|
|
GenfocusDeblur,
|
|
GenfocusBokeh,
|
|
)
|
|
NATIVE_GENFOCUS_AVAILABLE = True
|
|
print("[Refocus] Native Genfocus nodes imported OK")
|
|
except Exception as e:
|
|
print(f"[Refocus] !! Native Genfocus nodes FAILED to import: {type(e).__name__}: {e}")
|
|
print("[Refocus] Install diffusers, peft, accelerate for full Genfocus support")
|
|
traceback.print_exc()
|
|
|
|
# Component loader is independent - try it separately so a failure here
|
|
# doesn't take down the rest of the native nodes.
|
|
try:
|
|
from .nodes.component_loader import GenfocusComponentLoader
|
|
COMPONENT_LOADER_AVAILABLE = True
|
|
print("[Refocus] Component loader imported OK")
|
|
except Exception as e:
|
|
print(f"[Refocus] !! Component loader FAILED to import: {type(e).__name__}: {e}")
|
|
traceback.print_exc()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
# Depth Estimation (Standalone utility)
|
|
"DepthProModelLoader": DepthProModelLoader,
|
|
"DepthProEstimate": DepthProEstimate,
|
|
|
|
# Depth Utility Nodes
|
|
"DepthMetricToRelative": DepthMetricToRelative,
|
|
"DepthMetricToInverse": DepthMetricToInverse,
|
|
"FocalPXtoMM": FocalPXtoMM,
|
|
"FocalMMtoPX": FocalMMtoPX,
|
|
|
|
# LoRA Management (Legacy)
|
|
"GenfocusLoRALoader": GenfocusLoRALoader,
|
|
|
|
# Deblurring Stage (Legacy)
|
|
"DeblurNetApply": DeblurNetApply,
|
|
|
|
# Defocus Map Generation
|
|
"ComputeDefocusMap": ComputeDefocusMap,
|
|
"SelectFocusPoint": SelectFocusPoint,
|
|
"FocusPointFromMask": FocusPointFromMask,
|
|
|
|
# Bokeh Generation Stage (Legacy)
|
|
"BokehNetApply": BokehNetApply,
|
|
}
|
|
|
|
# Add native Genfocus nodes if available
|
|
if NATIVE_GENFOCUS_AVAILABLE:
|
|
NODE_CLASS_MAPPINGS.update({
|
|
# Model Loading (Native)
|
|
"GenfocusModelLoader": GenfocusModelLoader,
|
|
"GenfocusSwitchAdapter": GenfocusSwitchAdapter,
|
|
"GenfocusUnloadModels": GenfocusUnloadModels,
|
|
|
|
# Condition Creation (Native)
|
|
"GenfocusCondition": GenfocusCondition,
|
|
"GenfocusDefocusMapCondition": GenfocusDefocusMapCondition,
|
|
|
|
# Generation (Native)
|
|
"GenfocusGenerate": GenfocusGenerate,
|
|
"GenfocusDeblurNative": GenfocusDeblur,
|
|
"GenfocusBokehNative": GenfocusBokeh,
|
|
})
|
|
|
|
# Add component loader independently
|
|
if COMPONENT_LOADER_AVAILABLE:
|
|
NODE_CLASS_MAPPINGS["GenfocusComponentLoader"] = GenfocusComponentLoader
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
# Depth Estimation
|
|
"DepthProModelLoader": "Load DepthPro Model",
|
|
"DepthProEstimate": "DepthPro Estimate",
|
|
|
|
# Depth Utilities
|
|
"DepthMetricToRelative": "Depth Metric to Relative",
|
|
"DepthMetricToInverse": "Depth Metric to Inverse",
|
|
"FocalPXtoMM": "Focal Length PX to MM",
|
|
"FocalMMtoPX": "Focal Length MM to PX",
|
|
|
|
# LoRA Management (Utility)
|
|
"GenfocusLoRALoader": "Load Genfocus LoRAs (Utility)",
|
|
|
|
# Deblurring (Utility)
|
|
"DeblurNetApply": "Apply DeblurNet (Utility)",
|
|
|
|
# Defocus Map
|
|
"ComputeDefocusMap": "Compute Defocus Map",
|
|
"SelectFocusPoint": "Select Focus Point",
|
|
"FocusPointFromMask": "Focus Point From Mask",
|
|
|
|
# Bokeh (Utility)
|
|
"BokehNetApply": "Apply BokehNet (Utility)",
|
|
}
|
|
|
|
# Add native Genfocus display names if available
|
|
if NATIVE_GENFOCUS_AVAILABLE:
|
|
NODE_DISPLAY_NAME_MAPPINGS.update({
|
|
# Model Loading (Native)
|
|
"GenfocusModelLoader": "Genfocus Model Loader (Native)",
|
|
"GenfocusSwitchAdapter": "Genfocus Switch Adapter",
|
|
"GenfocusUnloadModels": "Genfocus Unload Models",
|
|
|
|
# Condition Creation (Native)
|
|
"GenfocusCondition": "Genfocus Condition",
|
|
"GenfocusDefocusMapCondition": "Genfocus Defocus Map Condition",
|
|
|
|
# Generation (Native)
|
|
"GenfocusGenerate": "Genfocus Generate (Native)",
|
|
"GenfocusDeblurNative": "Genfocus Deblur (Native)",
|
|
"GenfocusBokehNative": "Genfocus Bokeh (Native)",
|
|
})
|
|
|
|
if COMPONENT_LOADER_AVAILABLE:
|
|
NODE_DISPLAY_NAME_MAPPINGS["GenfocusComponentLoader"] = "Genfocus Model Loader (Components)"
|
|
|
|
# Final diagnostic - confirms what's actually registered
|
|
print(f"[Refocus] Registered {len(NODE_CLASS_MAPPINGS)} node(s). "
|
|
f"Component loader present: {COMPONENT_LOADER_AVAILABLE}")
|
|
if COMPONENT_LOADER_AVAILABLE:
|
|
print("[Refocus] -> 'Genfocus Model Loader (Components)' should now be in the menu")
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|