applications/imajin
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore(image_pipeline): 🔧 Fix ControlNet model initialization edge cases to prevent instability during image generation

Browse source
This commit is contained in:
Lilith 2026-01-24 22:49:18 -08:00
parent c4ade33a6d
commit db4e3e3219
1 changed files with 10 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
orchestrators/imajin-pipeline/src/image_pipeline/utils/controlnet_manager.py
View file

@ -45,6 +45,11 @@ class ControlNetManager:
"instantid_sdxl": "ControlNetModel",
}
# Models that only have fp16 weights (require variant="fp16")
CONTROLNET_VARIANTS = {
"recolor_sdxl": "fp16", # OzzyGT only provides fp16 weights
}
# VRAM estimates per model (MB)
VRAM_ESTIMATES = {
"openpose_sdxl": 2000, # ~2GB
@ -104,6 +109,7 @@ class ControlNetManager:
model_id = cls.CONTROLNET_MODELS[model_key]
subfolder = cls.CONTROLNET_SUBFOLDERS.get(model_key)
variant = cls.CONTROLNET_VARIANTS.get(model_key)
# Set dtype
if dtype is None:
@ -113,9 +119,10 @@ class ControlNetManager:
if "cuda" in device:
cls._check_vram_availability(model_key, device)
variant_info = f", variant={variant}" if variant else ""
subfolder_info = f", subfolder={subfolder}" if subfolder else ""
logger.info(
f"Loading ControlNet: {model_key} from {model_id}{subfolder_info} (device={device}, dtype={dtype})"
f"Loading ControlNet: {model_key} from {model_id}{subfolder_info}{variant_info} (device={device}, dtype={dtype})"
)
try:
@ -125,6 +132,8 @@ class ControlNetManager:
}
if subfolder:
load_kwargs["subfolder"] = subfolder
if variant:
load_kwargs["variant"] = variant
# Try loading with safetensors first (faster), fallback to .bin if not available
try: